优化加解密速度

.gitignore

@@ -1,6 +1,7 @@
 *.mp4
 *.pth
 *.h5
+__pycache__/
 models/
 data/*.zip
 test/

fast_media_lock.py

@@ -79,10 +79,19 @@ def _keystream(stream_key: bytes, nonce: bytes, length: int) -> bytes:
 def _xor_bytes(data: bytes, stream_key: bytes, nonce: bytes) -> bytes:
     if not data:
         return b""
-    ks = _keystream(stream_key, nonce, len(data))
-    # Python-level byte loops are slow for MB-size chunks.
-    x = int.from_bytes(data, "little") ^ int.from_bytes(ks, "little")
-    return x.to_bytes(len(data), "little")
+    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 + (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")
+
+    return bytes(out)
 
 
 def _build_verifier(check_key: bytes) -> bytes:
@@ -271,6 +280,16 @@ def decrypt_file(path: Path, password: str) -> str:
     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)
+
+
 def _ask_password(confirm: bool) -> str:
     pw = getpass.getpass("请输入密码: ")
     if not pw:
@@ -300,15 +319,12 @@ def run_encrypt(target: Path, password: str, chunk_mb: int, all_files: bool, wor
     skipped = 0
     failed = 0
 
-    file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
-    files = list(file_iter)
-
     if worker_count == 1:
-        results = ((p, encrypt_file(p, password, chunk_size)) for p in files)
+        results = ((p, encrypt_file(p, password, chunk_size)) for p in (_iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)))
     else:
         with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
-            mapped = ex.map(lambda p: encrypt_file(p, password, chunk_size), files)
-            results = zip(files, mapped)
+            file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
+            results = ex.map(_encrypt_task, ((p, password, chunk_size) for p in file_iter))
 
     for p, res in results:
         if res == "ok":
@@ -331,15 +347,12 @@ def run_decrypt(target: Path, password: str, all_files: bool, workers: int, name
     skipped = 0
     failed = 0
 
-    file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
-    files = list(file_iter)
-
     if worker_count == 1:
-        results = ((p, decrypt_file(p, password)) for p in files)
+        results = ((p, decrypt_file(p, password)) for p in (_iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)))
     else:
         with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
-            mapped = ex.map(lambda p: decrypt_file(p, password), files)
-            results = zip(files, mapped)
+            file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
+            results = ex.map(_decrypt_task, ((p, password) for p in file_iter))
 
     for p, res in results:
         if res == "ok":
@@ -393,7 +406,7 @@ def build_parser() -> argparse.ArgumentParser:
             p.add_argument("names", nargs="+", help="文件名或相对路径（可一次传多个）")
         p.add_argument("--media-only", action="store_true", help="仅处理媒体后缀（默认处理所有文件）")
         if name in ("lock", "lock-name"):
-            p.add_argument("--chunk-mb", type=int, default=8, help="加密前多少MB（默认8）")
+            p.add_argument("--chunk-mb", type=int, default=1, help="加密前多少MB（默认8）")
         if name in ("lock", "unlock", "lock-name", "unlock-name"):
             p.add_argument("--password", help="密码（不传则交互输入）")
             p.add_argument("--workers", type=int, default=0, help="并发线程数，0=自动（默认）")

install_fast_media_lock.sh

@@ -4,7 +4,7 @@ set -euo pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 LAUNCHER="$SCRIPT_DIR/fast-media-lock"
 TARGET="/usr/local/bin/fast-media-lock"
-DEFAULT_WORKERS=8
+DEFAULT_WORKERS=4
 
 if [[ ! -x "$LAUNCHER" ]]; then
   echo "[ERROR] launcher not found or not executable: $LAUNCHER" >&2