单文件进度条

fast_media_lock.py

@@ -22,6 +22,7 @@ import hmac
 import os
 import struct
 import sys
+import threading
 import zlib
 from pathlib import Path
 from typing import Iterable
@@ -32,6 +33,7 @@ SALT_SIZE = 16
 NONCE_SIZE = 16
 VERIFIER_SIZE = 16
 LOCKED_SUFFIX = ".lockx"
+IO_CHUNK_SIZE = 4 * 1024 * 1024
 
 # magic(8) + version(1) + reserved(3) + chunk_size(8) + original_size(8)
 # + salt(16) + nonce(16) + verifier(16) + crc32(4)
@@ -76,7 +78,7 @@ def _keystream(stream_key: bytes, nonce: bytes, length: int) -> bytes:
     return bytes(mv[:length])
 
 
-def _xor_bytes(data: bytes, stream_key: bytes, nonce: bytes) -> bytes:
+def _xor_bytes(data: bytes, stream_key: bytes, nonce: bytes, block_offset: int = 0) -> bytes:
     if not data:
         return b""
     out = bytearray(len(data))
@@ -87,7 +89,7 @@ def _xor_bytes(data: bytes, stream_key: bytes, nonce: bytes) -> bytes:
     # 分块 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)
+        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")
 
@@ -207,6 +209,42 @@ def _unlocked_name(path: Path) -> Path:
     return path
 
 
+class _ProgressWriter:
+    def __init__(self) -> None:
+        self._enabled = sys.stdout.isatty()
+        self._lock = threading.Lock()
+        self._last_len = 0
+
+    def update(self, label: str, done: int, total: int, path: Path) -> None:
+        if not self._enabled:
+            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}"
+
+        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:
+        if not self._enabled:
+            return
+
+        with self._lock:
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+            self._last_len = 0
+
+
+_progress = _ProgressWriter()
+
+
 def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
     if not path.exists() or not path.is_file():
         return "skip(not_file)"
@@ -230,10 +268,19 @@ def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
     n = min(chunk_size, original_size)
 
     with path.open("r+b") as f:
-        plain = f.read(n)
-        cipher = _xor_bytes(plain, stream_key, nonce)
-        f.seek(0)
-        f.write(cipher)
+        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)
@@ -241,6 +288,7 @@ def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
     if locked_path != path:
         path.rename(locked_path)
 
+    _progress.done()
     return "ok"
 
 
@@ -268,15 +316,25 @@ def decrypt_file(path: Path, password: str) -> str:
     n = min(meta["chunk_size"], meta["original_size"])
 
     with path.open("r+b") as f:
-        cipher = f.read(n)
-        plain = _xor_bytes(cipher, stream_key, meta["nonce"])
-        f.seek(0)
-        f.write(plain)
+        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"
 
 
@@ -329,13 +387,15 @@ def run_encrypt(target: Path, password: str, chunk_mb: int, all_files: bool, wor
     for p, res in results:
         if res == "ok":
             ok += 1
-            print(f"[OK] encrypted: {p}")
+            if not sys.stdout.isatty():
+                print(f"[OK] encrypted: {p}")
         elif res.startswith("fail"):
             failed += 1
             print(f"[FAIL] {p} -> {res}")
         else:
             skipped += 1
-            print(f"[SKIP] {p} -> {res}")
+            if not sys.stdout.isatty():
+                print(f"[SKIP] {p} -> {res}")
 
     print(f"完成: encrypted={ok}, failed={failed}, skipped={skipped}, chunk={chunk_mb}MB, workers={worker_count}")
     return 2 if failed else 0
@@ -357,13 +417,15 @@ def run_decrypt(target: Path, password: str, all_files: bool, workers: int, name
     for p, res in results:
         if res == "ok":
             ok += 1
-            print(f"[OK] decrypted: {p}")
+            if not sys.stdout.isatty():
+                print(f"[OK] decrypted: {p}")
         elif res.startswith("fail"):
             failed += 1
             print(f"[FAIL] {p} -> {res}")
         else:
             skipped += 1
-            print(f"[SKIP] {p} -> {res}")
+            if not sys.stdout.isatty():
+                print(f"[SKIP] {p} -> {res}")
 
     print(f"完成: decrypted={ok}, failed={failed}, skipped={skipped}, workers={worker_count}")
     return 2 if failed else 0