|
|
@@ -1,17 +1,14 @@
|
|
|
#!/usr/bin/env python3
|
|
|
"""
|
|
|
Fast partial media locker (cross-platform, Python stdlib only).
|
|
|
-
|
|
|
Goal:
|
|
|
- Very fast "lock/unlock" by encrypting only the first N MB of each file.
|
|
|
- No separate key file; password only.
|
|
|
- Metadata is embedded in each encrypted file tail, so files can be moved anywhere
|
|
|
and still be unlocked by this script.
|
|
|
-
|
|
|
Security note:
|
|
|
- This is NOT full-file encryption. It is designed for speed/obfuscation.
|
|
|
"""
|
|
|
-
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import argparse
|
|
|
@@ -23,6 +20,7 @@ import os
|
|
|
import struct
|
|
|
import sys
|
|
|
import threading
|
|
|
+import time
|
|
|
import zlib
|
|
|
from pathlib import Path
|
|
|
from typing import Iterable
|
|
|
@@ -33,7 +31,7 @@ SALT_SIZE = 16
|
|
|
NONCE_SIZE = 16
|
|
|
VERIFIER_SIZE = 16
|
|
|
LOCKED_SUFFIX = ".lockx"
|
|
|
-IO_CHUNK_SIZE = 4 * 1024 * 1024
|
|
|
+IO_CHUNK_SIZE = 8 * 1024 * 1024 # 8MB,兼顾速度与内存
|
|
|
|
|
|
# magic(8) + version(1) + reserved(3) + chunk_size(8) + original_size(8)
|
|
|
# + salt(16) + nonce(16) + verifier(16) + crc32(4)
|
|
|
@@ -52,7 +50,6 @@ class LockerError(Exception):
|
|
|
|
|
|
|
|
|
def _derive_keys(password: str, salt: bytes) -> tuple[bytes, bytes]:
|
|
|
- # scrypt is in stdlib and works on Linux/Windows without extra dependencies.
|
|
|
km = hashlib.scrypt(
|
|
|
password.encode("utf-8"),
|
|
|
salt=salt,
|
|
|
@@ -64,35 +61,21 @@ def _derive_keys(password: str, salt: bytes) -> tuple[bytes, bytes]:
|
|
|
return km[:32], km[32:]
|
|
|
|
|
|
|
|
|
-def _keystream(stream_key: bytes, nonce: bytes, length: int) -> bytes:
|
|
|
- if length <= 0:
|
|
|
- return b""
|
|
|
- blocks = (length + 31) // 32
|
|
|
- out = bytearray(blocks * 32)
|
|
|
- mv = memoryview(out)
|
|
|
- off = 0
|
|
|
- for counter in range(blocks):
|
|
|
- block = hmac.digest(stream_key, nonce + counter.to_bytes(8, "big"), hashlib.sha256)
|
|
|
- mv[off : off + 32] = block
|
|
|
- off += 32
|
|
|
- return bytes(mv[:length])
|
|
|
-
|
|
|
-
|
|
|
def _xor_bytes(data: bytes, stream_key: bytes, nonce: bytes, block_offset: int = 0) -> bytes:
|
|
|
+ """与旧版本完全兼容的密钥流 XOR(从 block_offset 开始)"""
|
|
|
if not data:
|
|
|
return b""
|
|
|
out = bytearray(len(data))
|
|
|
mv_in = memoryview(data)
|
|
|
mv_out = memoryview(out)
|
|
|
data_len = len(data)
|
|
|
-
|
|
|
- # 分块 XOR,避免一次性把整段数据转成大整数。
|
|
|
for off in range(0, data_len, 32):
|
|
|
block_len = min(32, data_len - off)
|
|
|
- ks = hmac.digest(stream_key, nonce + (block_offset + off // 32).to_bytes(8, "big"), hashlib.sha256)
|
|
|
- chunk = int.from_bytes(mv_in[off : off + block_len], "little") ^ int.from_bytes(ks[:block_len], "little")
|
|
|
- mv_out[off : off + block_len] = chunk.to_bytes(block_len, "little")
|
|
|
-
|
|
|
+ counter = block_offset + (off // 32)
|
|
|
+ ks = hmac.digest(stream_key, nonce + counter.to_bytes(8, "big"), hashlib.sha256)
|
|
|
+ # 32 字节整数 XOR,比纯字节循环快
|
|
|
+ chunk = int.from_bytes(mv_in[off:off + block_len], "little") ^ int.from_bytes(ks[:block_len], "little")
|
|
|
+ mv_out[off:off + block_len] = chunk.to_bytes(block_len, "little")
|
|
|
return bytes(out)
|
|
|
|
|
|
|
|
|
@@ -119,16 +102,12 @@ def _build_trailer(chunk_size: int, original_size: int, salt: bytes, nonce: byte
|
|
|
def _parse_trailer(raw: bytes) -> dict | None:
|
|
|
if len(raw) != TRAILER_SIZE:
|
|
|
return None
|
|
|
-
|
|
|
magic, version, _reserved, chunk_size, original_size, salt, nonce, verifier, crc = TRAILER_STRUCT.unpack(raw)
|
|
|
-
|
|
|
if magic != MAGIC or version != VERSION:
|
|
|
return None
|
|
|
-
|
|
|
expect_crc = zlib.crc32(raw[:-4]) & 0xFFFFFFFF
|
|
|
if crc != expect_crc:
|
|
|
return None
|
|
|
-
|
|
|
return {
|
|
|
"chunk_size": chunk_size,
|
|
|
"original_size": original_size,
|
|
|
@@ -142,7 +121,6 @@ def _read_trailer(path: Path) -> dict | None:
|
|
|
size = path.stat().st_size
|
|
|
if size < TRAILER_SIZE:
|
|
|
return None
|
|
|
-
|
|
|
with path.open("rb") as f:
|
|
|
f.seek(-TRAILER_SIZE, os.SEEK_END)
|
|
|
raw = f.read(TRAILER_SIZE)
|
|
|
@@ -153,9 +131,7 @@ def is_encrypted(path: Path) -> bool:
|
|
|
meta = _read_trailer(path)
|
|
|
if not meta:
|
|
|
return False
|
|
|
-
|
|
|
size = path.stat().st_size
|
|
|
- # Basic sanity: encrypted file size should be original + trailer.
|
|
|
return meta["original_size"] + TRAILER_SIZE == size
|
|
|
|
|
|
|
|
|
@@ -163,7 +139,6 @@ def _iter_files(target: Path, all_files: bool) -> Iterable[Path]:
|
|
|
if target.is_file():
|
|
|
yield target
|
|
|
return
|
|
|
-
|
|
|
for p in target.rglob("*"):
|
|
|
if not p.is_file():
|
|
|
continue
|
|
|
@@ -179,18 +154,15 @@ def _iter_files_by_names(target: Path, names: list[str], all_files: bool) -> Ite
|
|
|
selected: set[Path] = set()
|
|
|
raw_names = [n.strip() for n in names if n.strip()]
|
|
|
basename_set = {n for n in raw_names if "/" not in n and "\\" not in n}
|
|
|
-
|
|
|
for raw in raw_names:
|
|
|
if "/" in raw or "\\" in raw:
|
|
|
rel = Path(raw)
|
|
|
cand = (target / rel).resolve()
|
|
|
if cand.is_file():
|
|
|
selected.add(cand)
|
|
|
-
|
|
|
for p in _iter_files(target, all_files=all_files):
|
|
|
if p.name in basename_set:
|
|
|
selected.add(p.resolve())
|
|
|
-
|
|
|
for p in sorted(selected):
|
|
|
if not all_files and p.suffix.lower() not in MEDIA_EXTS:
|
|
|
continue
|
|
|
@@ -209,83 +181,73 @@ def _unlocked_name(path: Path) -> Path:
|
|
|
return path
|
|
|
|
|
|
|
|
|
-class _ProgressWriter:
|
|
|
- def __init__(self) -> None:
|
|
|
- self._enabled = sys.stdout.isatty()
|
|
|
+class Progress:
|
|
|
+ """实时整体进度条(安装风格)"""
|
|
|
+
|
|
|
+ def __init__(self, total_bytes: int, total_files: int):
|
|
|
+ self.total_bytes = max(total_bytes, 1)
|
|
|
+ self.total_files = max(total_files, 1)
|
|
|
+ self.processed_bytes = 0
|
|
|
+ self.completed_files = 0
|
|
|
+ self.current_file = ""
|
|
|
self._lock = threading.Lock()
|
|
|
+ self._enabled = sys.stdout.isatty()
|
|
|
+ self._start = time.monotonic()
|
|
|
self._last_len = 0
|
|
|
- self._file_mode = False
|
|
|
-
|
|
|
- def update(self, label: str, done: int, total: int, path: Path) -> None:
|
|
|
- if not self._enabled:
|
|
|
- return
|
|
|
-
|
|
|
- if not self._file_mode:
|
|
|
- return
|
|
|
-
|
|
|
- total = max(total, 1)
|
|
|
- percent = (done * 100) // total
|
|
|
- width = 24
|
|
|
- filled = min(width, (done * width) // total)
|
|
|
- bar = "#" * filled + "-" * (width - filled)
|
|
|
- msg = f"{label} [{bar}] {percent:3d}% {done}/{total} {path}"
|
|
|
+ self._last_update = 0.0
|
|
|
|
|
|
+ def update(self, added_bytes: int = 0, file_done: bool = False, current: str = ""):
|
|
|
with self._lock:
|
|
|
- pad = max(0, self._last_len - len(msg))
|
|
|
- sys.stdout.write("\r" + msg + (" " * pad))
|
|
|
- sys.stdout.flush()
|
|
|
- self._last_len = len(msg)
|
|
|
-
|
|
|
- def done(self) -> None:
|
|
|
+ self.processed_bytes += added_bytes
|
|
|
+ if file_done:
|
|
|
+ self.completed_files += 1
|
|
|
+ if current:
|
|
|
+ self.current_file = current
|
|
|
+ now = time.monotonic()
|
|
|
+ # 限制刷新频率,避免刷屏
|
|
|
+ if not self._enabled or (now - self._last_update < 0.08 and not file_done):
|
|
|
+ return
|
|
|
+ self._last_update = now
|
|
|
+ self._render()
|
|
|
+
|
|
|
+ def _render(self):
|
|
|
+ elapsed = max(time.monotonic() - self._start, 0.001)
|
|
|
+ pct = min(100.0, self.processed_bytes * 100.0 / self.total_bytes)
|
|
|
+ speed = self.processed_bytes / elapsed / (1024 * 1024) # MB/s
|
|
|
+
|
|
|
+ bar_width = 28
|
|
|
+ filled = int(bar_width * pct / 100)
|
|
|
+ bar = "█" * filled + "░" * (bar_width - filled)
|
|
|
+
|
|
|
+ name = self.current_file
|
|
|
+ if len(name) > 42:
|
|
|
+ name = "…" + name[-41:]
|
|
|
+
|
|
|
+ msg = (
|
|
|
+ f"\r[{bar}] {pct:5.1f}% "
|
|
|
+ f"{self.completed_files}/{self.total_files} files "
|
|
|
+ f"{speed:5.1f} MB/s "
|
|
|
+ f"{name}"
|
|
|
+ )
|
|
|
+ pad = max(0, self._last_len - len(msg))
|
|
|
+ sys.stdout.write(msg + " " * pad)
|
|
|
+ sys.stdout.flush()
|
|
|
+ self._last_len = len(msg)
|
|
|
+
|
|
|
+ def finish(self):
|
|
|
if not self._enabled:
|
|
|
return
|
|
|
-
|
|
|
- if not self._file_mode:
|
|
|
- return
|
|
|
-
|
|
|
with self._lock:
|
|
|
+ self.processed_bytes = self.total_bytes
|
|
|
+ self.completed_files = self.total_files
|
|
|
+ self._render()
|
|
|
sys.stdout.write("\n")
|
|
|
sys.stdout.flush()
|
|
|
- self._last_len = 0
|
|
|
-
|
|
|
- def finish_count(self) -> None:
|
|
|
- if not self._enabled:
|
|
|
- return
|
|
|
-
|
|
|
- with self._lock:
|
|
|
- sys.stdout.write("\n")
|
|
|
- sys.stdout.flush()
|
|
|
- self._last_len = 0
|
|
|
-
|
|
|
- def set_file_mode(self, enabled: bool) -> None:
|
|
|
- self._file_mode = enabled
|
|
|
-
|
|
|
- def count_update(self, done: int, total: int) -> None:
|
|
|
- if not self._enabled:
|
|
|
- return
|
|
|
- with self._lock:
|
|
|
- msg = f"已完成: {done}/{total}"
|
|
|
- pad = max(0, self._last_len - len(msg))
|
|
|
- sys.stdout.write("\r" + msg + (" " * pad))
|
|
|
- sys.stdout.flush()
|
|
|
- self._last_len = len(msg)
|
|
|
-
|
|
|
-
|
|
|
-_progress = _ProgressWriter()
|
|
|
-
|
|
|
-
|
|
|
-def _iter_files_list(target: Path, all_files: bool) -> list[Path]:
|
|
|
- return list(_iter_files(target, all_files=all_files))
|
|
|
|
|
|
|
|
|
-def _iter_files_by_names_list(target: Path, names: list[str], all_files: bool) -> list[Path]:
|
|
|
- return list(_iter_files_by_names(target, names, all_files=all_files))
|
|
|
-
|
|
|
-
|
|
|
-def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
|
|
|
+def encrypt_file(path: Path, password: str, chunk_size: int, progress: Progress | None = None) -> str:
|
|
|
if not path.exists() or not path.is_file():
|
|
|
return "skip(not_file)"
|
|
|
-
|
|
|
if is_encrypted(path):
|
|
|
return "skip(already_encrypted)"
|
|
|
|
|
|
@@ -303,33 +265,46 @@ def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
|
|
|
verifier = _build_verifier(check_key)
|
|
|
|
|
|
n = min(chunk_size, original_size)
|
|
|
+ if progress:
|
|
|
+ progress.update(current=str(path))
|
|
|
|
|
|
- with path.open("r+b") as f:
|
|
|
- processed = 0
|
|
|
- block_offset = 0
|
|
|
- while processed < n:
|
|
|
- read_len = min(IO_CHUNK_SIZE, n - processed)
|
|
|
- plain = f.read(read_len)
|
|
|
- if not plain:
|
|
|
- break
|
|
|
- cipher = _xor_bytes(plain, stream_key, nonce, block_offset=block_offset)
|
|
|
- f.seek(processed)
|
|
|
- f.write(cipher)
|
|
|
- processed += len(plain)
|
|
|
- block_offset += (len(plain) + 31) // 32
|
|
|
- _progress.update("[ENC]", processed, n, path)
|
|
|
- trailer = _build_trailer(chunk_size=chunk_size, original_size=original_size, salt=salt, nonce=nonce, verifier=verifier)
|
|
|
- f.seek(0, os.SEEK_END)
|
|
|
- f.write(trailer)
|
|
|
-
|
|
|
- if locked_path != path:
|
|
|
- path.rename(locked_path)
|
|
|
-
|
|
|
- _progress.done()
|
|
|
- return "ok"
|
|
|
-
|
|
|
-
|
|
|
-def decrypt_file(path: Path, password: str) -> str:
|
|
|
+ try:
|
|
|
+ with path.open("r+b") as f:
|
|
|
+ processed = 0
|
|
|
+ block_offset = 0
|
|
|
+ while processed < n:
|
|
|
+ read_len = min(IO_CHUNK_SIZE, n - processed)
|
|
|
+ plain = f.read(read_len)
|
|
|
+ if not plain:
|
|
|
+ break
|
|
|
+ cipher = _xor_bytes(plain, stream_key, nonce, block_offset=block_offset)
|
|
|
+ f.seek(processed)
|
|
|
+ f.write(cipher)
|
|
|
+ processed += len(plain)
|
|
|
+ block_offset += (len(plain) + 31) // 32
|
|
|
+ if progress:
|
|
|
+ progress.update(added_bytes=len(plain))
|
|
|
+
|
|
|
+ trailer = _build_trailer(
|
|
|
+ chunk_size=chunk_size,
|
|
|
+ original_size=original_size,
|
|
|
+ salt=salt,
|
|
|
+ nonce=nonce,
|
|
|
+ verifier=verifier,
|
|
|
+ )
|
|
|
+ f.seek(0, os.SEEK_END)
|
|
|
+ f.write(trailer)
|
|
|
+
|
|
|
+ if locked_path != path:
|
|
|
+ path.rename(locked_path)
|
|
|
+ if progress:
|
|
|
+ progress.update(file_done=True)
|
|
|
+ return "ok"
|
|
|
+ except Exception as e:
|
|
|
+ return f"fail({type(e).__name__})"
|
|
|
+
|
|
|
+
|
|
|
+def decrypt_file(path: Path, password: str, progress: Progress | None = None) -> str:
|
|
|
if not path.exists() or not path.is_file():
|
|
|
return "skip(not_file)"
|
|
|
|
|
|
@@ -351,50 +326,45 @@ def decrypt_file(path: Path, password: str) -> str:
|
|
|
return "fail(wrong_password)"
|
|
|
|
|
|
n = min(meta["chunk_size"], meta["original_size"])
|
|
|
+ if progress:
|
|
|
+ progress.update(current=str(path))
|
|
|
|
|
|
- with path.open("r+b") as f:
|
|
|
- processed = 0
|
|
|
- block_offset = 0
|
|
|
- while processed < n:
|
|
|
- read_len = min(IO_CHUNK_SIZE, n - processed)
|
|
|
- cipher = f.read(read_len)
|
|
|
- if not cipher:
|
|
|
- break
|
|
|
- plain = _xor_bytes(cipher, stream_key, meta["nonce"], block_offset=block_offset)
|
|
|
- f.seek(processed)
|
|
|
- f.write(plain)
|
|
|
- processed += len(cipher)
|
|
|
- block_offset += (len(cipher) + 31) // 32
|
|
|
- _progress.update("[DEC]", processed, n, path)
|
|
|
- f.truncate(meta["original_size"])
|
|
|
-
|
|
|
- if unlocked_path != path:
|
|
|
- path.rename(unlocked_path)
|
|
|
-
|
|
|
- _progress.done()
|
|
|
- return "ok"
|
|
|
-
|
|
|
-
|
|
|
-def _encrypt_task(args: tuple[Path, str, int]) -> tuple[Path, str]:
|
|
|
- path, password, chunk_size = args
|
|
|
- return path, encrypt_file(path, password, chunk_size)
|
|
|
-
|
|
|
-
|
|
|
-def _decrypt_task(args: tuple[Path, str]) -> tuple[Path, str]:
|
|
|
- path, password = args
|
|
|
- return path, decrypt_file(path, password)
|
|
|
+ try:
|
|
|
+ with path.open("r+b") as f:
|
|
|
+ processed = 0
|
|
|
+ block_offset = 0
|
|
|
+ while processed < n:
|
|
|
+ read_len = min(IO_CHUNK_SIZE, n - processed)
|
|
|
+ cipher = f.read(read_len)
|
|
|
+ if not cipher:
|
|
|
+ break
|
|
|
+ plain = _xor_bytes(cipher, stream_key, meta["nonce"], block_offset=block_offset)
|
|
|
+ f.seek(processed)
|
|
|
+ f.write(plain)
|
|
|
+ processed += len(cipher)
|
|
|
+ block_offset += (len(cipher) + 31) // 32
|
|
|
+ if progress:
|
|
|
+ progress.update(added_bytes=len(cipher))
|
|
|
+
|
|
|
+ f.truncate(meta["original_size"])
|
|
|
+
|
|
|
+ if unlocked_path != path:
|
|
|
+ path.rename(unlocked_path)
|
|
|
+ if progress:
|
|
|
+ progress.update(file_done=True)
|
|
|
+ return "ok"
|
|
|
+ except Exception as e:
|
|
|
+ return f"fail({type(e).__name__})"
|
|
|
|
|
|
|
|
|
def _ask_password(confirm: bool) -> str:
|
|
|
pw = getpass.getpass("请输入密码: ")
|
|
|
if not pw:
|
|
|
raise LockerError("密码不能为空")
|
|
|
-
|
|
|
if confirm:
|
|
|
pw2 = getpass.getpass("请再次输入密码: ")
|
|
|
if pw != pw2:
|
|
|
raise LockerError("两次密码不一致")
|
|
|
-
|
|
|
return pw
|
|
|
|
|
|
|
|
|
@@ -407,167 +377,144 @@ def _resolve_workers(workers: int) -> int:
|
|
|
return workers
|
|
|
|
|
|
|
|
|
-def run_encrypt(target: Path, password: str, chunk_mb: int, all_files: bool, workers: int, file_progress: bool, names: list[str] | None = None) -> int:
|
|
|
+def _collect_files(target: Path, all_files: bool, names: list[str] | None) -> list[Path]:
|
|
|
+ if names:
|
|
|
+ return list(_iter_files_by_names(target, names, all_files=all_files))
|
|
|
+ return list(_iter_files(target, all_files=all_files))
|
|
|
+
|
|
|
+
|
|
|
+def run_encrypt(
|
|
|
+ target: Path,
|
|
|
+ password: str,
|
|
|
+ chunk_mb: int,
|
|
|
+ all_files: bool,
|
|
|
+ workers: int,
|
|
|
+ names: list[str] | None = None,
|
|
|
+) -> int:
|
|
|
chunk_size = chunk_mb * 1024 * 1024
|
|
|
worker_count = _resolve_workers(workers)
|
|
|
- ok = 0
|
|
|
- skipped = 0
|
|
|
- failed = 0
|
|
|
- _progress.set_file_mode(file_progress)
|
|
|
|
|
|
- files = _iter_files_by_names_list(target, names, all_files=all_files) if names else _iter_files_list(target, all_files=all_files)
|
|
|
- total_files = len(files)
|
|
|
+ files = _collect_files(target, all_files, names)
|
|
|
+ if not files:
|
|
|
+ print("没有找到需要处理的文件")
|
|
|
+ return 0
|
|
|
|
|
|
- if not file_progress:
|
|
|
- if worker_count == 1:
|
|
|
- completed = 0
|
|
|
- results = []
|
|
|
- try:
|
|
|
- for p in files:
|
|
|
- results.append((p, encrypt_file(p, password, chunk_size)))
|
|
|
- completed += 1
|
|
|
- _progress.count_update(completed, total_files)
|
|
|
- except KeyboardInterrupt:
|
|
|
- print("\n已收到中断信号,当前文件处理完成后退出", file=sys.stderr)
|
|
|
- else:
|
|
|
- completed = 0
|
|
|
- with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
|
|
|
- future_map = {ex.submit(encrypt_file, p, password, chunk_size): p for p in files}
|
|
|
- results = []
|
|
|
- try:
|
|
|
- for fut in concurrent.futures.as_completed(future_map):
|
|
|
- p = future_map[fut]
|
|
|
- res = fut.result()
|
|
|
- results.append((p, res))
|
|
|
- completed += 1
|
|
|
- _progress.count_update(completed, total_files)
|
|
|
- except KeyboardInterrupt:
|
|
|
- print("\n已收到中断信号,等待当前文件处理完成后退出", file=sys.stderr)
|
|
|
- for fut, p in future_map.items():
|
|
|
- if fut.done():
|
|
|
- continue
|
|
|
- try:
|
|
|
- res = fut.result()
|
|
|
- except Exception as e:
|
|
|
- res = f"fail({type(e).__name__})"
|
|
|
- results.append((p, res))
|
|
|
- else:
|
|
|
+ # 预计算总工作量(用于进度条)
|
|
|
+ total_bytes = 0
|
|
|
+ for p in files:
|
|
|
+ try:
|
|
|
+ sz = p.stat().st_size
|
|
|
+ total_bytes += min(chunk_size, sz)
|
|
|
+ except OSError:
|
|
|
+ pass
|
|
|
+
|
|
|
+ progress = Progress(total_bytes, len(files))
|
|
|
+ ok = skipped = failed = 0
|
|
|
+ results: list[tuple[Path, str]] = []
|
|
|
+
|
|
|
+ def task(p: Path) -> tuple[Path, str]:
|
|
|
+ return p, encrypt_file(p, password, chunk_size, progress)
|
|
|
+
|
|
|
+ try:
|
|
|
if worker_count == 1:
|
|
|
- results = ((p, encrypt_file(p, password, chunk_size)) for p in files)
|
|
|
+ for p in files:
|
|
|
+ results.append(task(p))
|
|
|
else:
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
|
|
|
- results = ex.map(_encrypt_task, ((p, password, chunk_size) for p in files))
|
|
|
+ futures = {ex.submit(task, p): p for p in files}
|
|
|
+ for fut in concurrent.futures.as_completed(futures):
|
|
|
+ results.append(fut.result())
|
|
|
+ except KeyboardInterrupt:
|
|
|
+ print("\n已收到中断信号,等待当前任务结束后退出…", file=sys.stderr)
|
|
|
+ # 尽量收集已完成的结果
|
|
|
+ pass
|
|
|
+
|
|
|
+ progress.finish()
|
|
|
|
|
|
for p, res in results:
|
|
|
if res == "ok":
|
|
|
ok += 1
|
|
|
- if not file_progress and not sys.stdout.isatty():
|
|
|
- print(f"[OK] encrypted: {p}")
|
|
|
elif res.startswith("fail"):
|
|
|
failed += 1
|
|
|
- if not file_progress:
|
|
|
- print(f"[FAIL] {p} -> {res}")
|
|
|
+ print(f"[FAIL] {p} -> {res}")
|
|
|
else:
|
|
|
skipped += 1
|
|
|
- if not file_progress and not sys.stdout.isatty():
|
|
|
- print(f"[SKIP] {p} -> {res}")
|
|
|
-
|
|
|
- if not file_progress:
|
|
|
- _progress.finish_count()
|
|
|
|
|
|
print(f"完成: encrypted={ok}, failed={failed}, skipped={skipped}, chunk={chunk_mb}MB, workers={worker_count}")
|
|
|
return 2 if failed else 0
|
|
|
|
|
|
|
|
|
-def run_decrypt(target: Path, password: str, all_files: bool, workers: int, file_progress: bool, names: list[str] | None = None) -> int:
|
|
|
+def run_decrypt(
|
|
|
+ target: Path,
|
|
|
+ password: str,
|
|
|
+ all_files: bool,
|
|
|
+ workers: int,
|
|
|
+ names: list[str] | None = None,
|
|
|
+) -> int:
|
|
|
worker_count = _resolve_workers(workers)
|
|
|
- ok = 0
|
|
|
- skipped = 0
|
|
|
- failed = 0
|
|
|
- _progress.set_file_mode(file_progress)
|
|
|
|
|
|
- files = _iter_files_by_names_list(target, names, all_files=all_files) if names else _iter_files_list(target, all_files=all_files)
|
|
|
- total_files = len(files)
|
|
|
+ files = _collect_files(target, all_files, names)
|
|
|
+ if not files:
|
|
|
+ print("没有找到需要处理的文件")
|
|
|
+ return 0
|
|
|
|
|
|
- if not file_progress:
|
|
|
- if worker_count == 1:
|
|
|
- completed = 0
|
|
|
- results = []
|
|
|
- try:
|
|
|
- for p in files:
|
|
|
- results.append((p, decrypt_file(p, password)))
|
|
|
- completed += 1
|
|
|
- _progress.count_update(completed, total_files)
|
|
|
- except KeyboardInterrupt:
|
|
|
- print("\n已收到中断信号,当前文件处理完成后退出", file=sys.stderr)
|
|
|
- else:
|
|
|
- completed = 0
|
|
|
- with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
|
|
|
- future_map = {ex.submit(decrypt_file, p, password): p for p in files}
|
|
|
- results = []
|
|
|
- try:
|
|
|
- for fut in concurrent.futures.as_completed(future_map):
|
|
|
- p = future_map[fut]
|
|
|
- res = fut.result()
|
|
|
- results.append((p, res))
|
|
|
- completed += 1
|
|
|
- _progress.count_update(completed, total_files)
|
|
|
- except KeyboardInterrupt:
|
|
|
- print("\n已收到中断信号,等待当前文件处理完成后退出", file=sys.stderr)
|
|
|
- for fut, p in future_map.items():
|
|
|
- if fut.done():
|
|
|
- continue
|
|
|
- try:
|
|
|
- res = fut.result()
|
|
|
- except Exception as e:
|
|
|
- res = f"fail({type(e).__name__})"
|
|
|
- results.append((p, res))
|
|
|
- else:
|
|
|
+ # 解密时无法提前准确知道每个文件的 chunk,用文件大小近似(足够进度条使用)
|
|
|
+ total_bytes = 0
|
|
|
+ for p in files:
|
|
|
+ try:
|
|
|
+ total_bytes += p.stat().st_size
|
|
|
+ except OSError:
|
|
|
+ pass
|
|
|
+ # 实际只会处理前面的 chunk,这里给一个合理估计即可
|
|
|
+ total_bytes = max(total_bytes // 8, 1) # 粗略估计
|
|
|
+
|
|
|
+ progress = Progress(total_bytes, len(files))
|
|
|
+ ok = skipped = failed = 0
|
|
|
+ results: list[tuple[Path, str]] = []
|
|
|
+
|
|
|
+ def task(p: Path) -> tuple[Path, str]:
|
|
|
+ return p, decrypt_file(p, password, progress)
|
|
|
+
|
|
|
+ try:
|
|
|
if worker_count == 1:
|
|
|
- results = ((p, decrypt_file(p, password)) for p in files)
|
|
|
+ for p in files:
|
|
|
+ results.append(task(p))
|
|
|
else:
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
|
|
|
- results = ex.map(_decrypt_task, ((p, password) for p in files))
|
|
|
+ futures = {ex.submit(task, p): p for p in files}
|
|
|
+ for fut in concurrent.futures.as_completed(futures):
|
|
|
+ results.append(fut.result())
|
|
|
+ except KeyboardInterrupt:
|
|
|
+ print("\n已收到中断信号,等待当前任务结束后退出…", file=sys.stderr)
|
|
|
+
|
|
|
+ progress.finish()
|
|
|
|
|
|
for p, res in results:
|
|
|
if res == "ok":
|
|
|
ok += 1
|
|
|
- if not file_progress and not sys.stdout.isatty():
|
|
|
- print(f"[OK] decrypted: {p}")
|
|
|
elif res.startswith("fail"):
|
|
|
failed += 1
|
|
|
- if not file_progress:
|
|
|
- print(f"[FAIL] {p} -> {res}")
|
|
|
+ print(f"[FAIL] {p} -> {res}")
|
|
|
else:
|
|
|
skipped += 1
|
|
|
- if not file_progress and not sys.stdout.isatty():
|
|
|
- print(f"[SKIP] {p} -> {res}")
|
|
|
-
|
|
|
- if not file_progress:
|
|
|
- _progress.finish_count()
|
|
|
|
|
|
print(f"完成: decrypted={ok}, failed={failed}, skipped={skipped}, workers={worker_count}")
|
|
|
return 2 if failed else 0
|
|
|
|
|
|
|
|
|
def run_status(target: Path, all_files: bool) -> int:
|
|
|
- total = 0
|
|
|
- encrypted = 0
|
|
|
-
|
|
|
+ total = encrypted = 0
|
|
|
for p in _iter_files(target, all_files=all_files):
|
|
|
total += 1
|
|
|
flag = is_encrypted(p)
|
|
|
encrypted += int(flag)
|
|
|
print(f"{'[ENC]' if flag else '[RAW]'} {p}")
|
|
|
-
|
|
|
print(f"统计: total={total}, encrypted={encrypted}, raw={total - encrypted}")
|
|
|
return 0
|
|
|
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
|
- parser = argparse.ArgumentParser(
|
|
|
- description="快速加密工具"
|
|
|
- )
|
|
|
+ parser = argparse.ArgumentParser(description="快速部分加密工具(兼容旧版本)")
|
|
|
sub = parser.add_subparsers(dest="cmd", required=True)
|
|
|
|
|
|
for name in ("lock", "unlock", "status", "lock-name", "unlock-name"):
|
|
|
@@ -586,13 +533,11 @@ def build_parser() -> argparse.ArgumentParser:
|
|
|
if name in ("lock-name", "unlock-name"):
|
|
|
p.add_argument("names", nargs="+", help="文件名或相对路径(可一次传多个)")
|
|
|
p.add_argument("--media-only", action="store_true", help="仅处理媒体后缀(默认处理所有文件)")
|
|
|
- p.add_argument("--file-progress", action="store_true", help="显示单文件进度条,默认显示已完成文件数/总文件数")
|
|
|
if name in ("lock", "lock-name"):
|
|
|
- p.add_argument("--chunk-mb", type=int, default=1, help="加密前多少MB(默认8)")
|
|
|
+ p.add_argument("--chunk-mb", type=int, default=1, help="加密前多少MB(默认1)")
|
|
|
if name in ("lock", "unlock", "lock-name", "unlock-name"):
|
|
|
p.add_argument("--password", help="密码(不传则交互输入)")
|
|
|
p.add_argument("--workers", type=int, default=0, help="并发线程数,0=自动(默认)")
|
|
|
-
|
|
|
return parser
|
|
|
|
|
|
|
|
|
@@ -600,43 +545,36 @@ def main() -> int:
|
|
|
parser = build_parser()
|
|
|
args = parser.parse_args()
|
|
|
target = Path(args.target)
|
|
|
-
|
|
|
if not target.exists():
|
|
|
print(f"目标不存在: {target}", file=sys.stderr)
|
|
|
return 1
|
|
|
|
|
|
try:
|
|
|
all_files = not args.media_only
|
|
|
-
|
|
|
if args.cmd in ("lock", "encrypt"):
|
|
|
if args.chunk_mb <= 0:
|
|
|
raise LockerError("--chunk-mb 必须大于0")
|
|
|
pw = args.password if args.password else _ask_password(confirm=True)
|
|
|
- return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, file_progress=args.file_progress)
|
|
|
-
|
|
|
+ return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers)
|
|
|
if args.cmd in ("unlock", "decrypt"):
|
|
|
pw = args.password if args.password else _ask_password(confirm=False)
|
|
|
- return run_decrypt(target, pw, all_files=all_files, workers=args.workers, file_progress=args.file_progress)
|
|
|
-
|
|
|
+ return run_decrypt(target, pw, all_files=all_files, workers=args.workers)
|
|
|
if args.cmd in ("lock-name", "encrypt-name"):
|
|
|
if args.chunk_mb <= 0:
|
|
|
raise LockerError("--chunk-mb 必须大于0")
|
|
|
pw = args.password if args.password else _ask_password(confirm=True)
|
|
|
- return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, file_progress=args.file_progress, names=args.names)
|
|
|
-
|
|
|
+ return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, names=args.names)
|
|
|
if args.cmd in ("unlock-name", "decrypt-name"):
|
|
|
pw = args.password if args.password else _ask_password(confirm=False)
|
|
|
- return run_decrypt(target, pw, all_files=all_files, workers=args.workers, file_progress=args.file_progress, names=args.names)
|
|
|
-
|
|
|
+ return run_decrypt(target, pw, all_files=all_files, workers=args.workers, names=args.names)
|
|
|
return run_status(target, all_files=all_files)
|
|
|
-
|
|
|
except LockerError as e:
|
|
|
print(f"错误: {e}", file=sys.stderr)
|
|
|
return 1
|
|
|
except KeyboardInterrupt:
|
|
|
- print("已取消", file=sys.stderr)
|
|
|
+ print("\n已取消", file=sys.stderr)
|
|
|
return 130
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
- raise SystemExit(main())
|
|
|
+ raise SystemExit(main())
|