7z-encrypt/crypto.py

"""
加密引擎模块 (Crypto Engine)

流式 AES-256-GCM 分块加密,基于 cryptography.cobblestone (C2SP chunked-encryption)。
每块 16KiB 独立 GCM 认证,流式处理不整块进内存,适合大文件。

流程对应模块图 (02-加密引擎):
    B(取密钥) -> C(初始化流式上下文) -> D(逐块加密) -> F(输出参数 + 密文流)

格式说明:
    cobblestone 输出流自带头部 + 分块结构 (magic/版本/随机 IV/密文+tag),
    IV 由库内部管理,不需要也不应该手动维护 iv 列表。
    解密时只需密钥 + context,finalize() 验证整体完整性,篡改即抛 InvalidTag。
"""

import json
import os
from pathlib import Path
from typing import IO, Any, Callable

from cryptography.cobblestone import Cobblestone256Decryptor, Cobblestone256Encryptor
from cryptography.exceptions import InvalidTag

# 应用绑定上下文: 防止密文被挪用到其他系统/版本
CONTEXT = b"7z-encrypt:v1"

# 默认密钥库位置 (配置外置原则, 权限 600)
DEFAULT_KEYRING = Path(__file__).resolve().parent / "config" / "keyring.json"


class CryptoError(RuntimeError):
    """加密/解密过程错误"""


class KeyNotFoundError(CryptoError):
    """密钥 ID 不存在"""


class IntegrityError(CryptoError):
    """解密验证失败 (密钥错误或数据被篡改)"""


class CryptoEngine:
    """流式加密引擎, 密钥持久化到 keyring 文件"""

    def __init__(self, keyring_path: str | Path | None = None) -> None:
        self.keyring_path = Path(keyring_path) if keyring_path else DEFAULT_KEYRING
        self._keyring = self._load_keyring()

    # ---------- 密钥库 ----------

    def _load_keyring(self) -> dict[str, bytes]:
        """从 keyring 文件加载密钥 (hex -> bytes)"""
        if not self.keyring_path.exists():
            return {}
        data = json.loads(self.keyring_path.read_text(encoding="utf-8"))
        return {kid: bytes.fromhex(hex_str) for kid, hex_str in data.items()}

    def _save_keyring(self) -> None:
        """持久化密钥库, 文件权限 600 (仅本人可读写)"""
        self.keyring_path.parent.mkdir(parents=True, exist_ok=True)
        data = {kid: key.hex() for kid, key in self._keyring.items()}
        self.keyring_path.write_text(json.dumps(data, indent=2), encoding="utf-8")
        os.chmod(self.keyring_path, 0o600)

    def get_key(self, key_id: str) -> bytes:
        """步骤 B: 取密钥。密钥不存在抛 KeyNotFoundError"""
        try:
            return self._keyring[key_id]
        except KeyError:
            raise KeyNotFoundError(f"密钥ID '{key_id}' 未找到 (keyring: {self.keyring_path})") from None

    def generate_key(self, key_id: str) -> bytes:
        """生成新密钥 (os.urandom 32B) 并持久化。已存在则不覆盖, 抛错防误用"""
        if key_id in self._keyring:
            raise CryptoError(f"密钥ID '{key_id}' 已存在, 不覆盖 (需轮换请换新 key_id)")
        key = os.urandom(32)
        self._keyring[key_id] = key
        self._save_keyring()
        return key

    # ---------- 加解密 ----------

    def encrypt_stream(
        self,
        src: IO[bytes],
        dst: IO[bytes],
        key_id: str,
        context: bytes = CONTEXT,
        progress: Callable[[int], Any] | None = None,
    ) -> dict[str, Any]:
        """步骤 A & F: 流式加密。src 读明文 -> dst 写密文, 返回加密参数

        Args:
            src: 可读二进制流 (文件/BytesIO)
            dst: 可写二进制流
            key_id: 密钥标识
            context: 应用绑定上下文 (默认 7z-encrypt:v1)
            progress: 进度回调, 每读块后调 progress(已读字节)
        """
        key = self.get_key(key_id)
        enc = Cobblestone256Encryptor(key, context)
        done = 0
        while True:
            chunk = src.read(1 << 16)  # 64KiB 读块
            if not chunk:
                break
            dst.write(enc.update(chunk))
            done += len(chunk)
            if progress is not None:
                progress(done)
        dst.write(enc.finalize())
        return {"alg": "cobblestone-aes256gcm", "key_id": key_id, "context": context.decode()}

    def decrypt_stream(
        self,
        src: IO[bytes],
        dst: IO[bytes],
        key_id: str,
        context: bytes = CONTEXT,
        progress: Callable[[int], Any] | None = None,
    ) -> None:
        """解密流 (进度回调: 每读块后调 progress(已读字节))"""
        key = self.get_key(key_id)
        try:
            dec = Cobblestone256Decryptor(key, context)
            done = 0
            while True:
                chunk = src.read(1 << 16)
                if not chunk:
                    break
                dst.write(dec.update(chunk))  # 块级认证: 篡改在这里就抛 InvalidTag
                done += len(chunk)
                if progress is not None:
                    progress(done)
            dst.write(dec.finalize())  # finalize 返回剩余明文
        except InvalidTag as e:
            raise IntegrityError("解密验证失败: 密钥不匹配或密文被篡改") from e


# ================= 使用示例 (Usage Example) =================
if __name__ == "__main__":
    import io

    engine = CryptoEngine()

    # 1. 确保密钥存在 (首次运行生成并持久化)
    try:
        engine.get_key("default_key")
    except KeyNotFoundError:
        engine.generate_key("default_key")
        print("[密钥库] 已生成 default_key ->", engine.keyring_path)

    # 2. 模拟流式加密 (大文件场景: 分块读入)
    plain = (b"This is a streaming encryption test. " * 10000)  # ~330KB
    src = io.BytesIO(plain)
    dst = io.BytesIO()
    params = engine.encrypt_stream(src, dst, key_id="default_key")

    print("=== 加密输出 (JSON) ===")
    print(json.dumps(params, indent=2))
    print(f"明文 {len(plain)}B -> 密文 {len(dst.getvalue())}B (含头部+分块tag开销)")

    # 3. 流式解密还原
    out = io.BytesIO()
    dst.seek(0)
    engine.decrypt_stream(dst, out, key_id="default_key")
    restored = out.getvalue()
    assert restored == plain, "解密结果与明文不一致!"
    print("解密还原: 一致 ✓")

    # 4. 篡改检测: 改密文一个字节, finalize 必须抛 IntegrityError
    tampered = bytearray(dst.getvalue())
    tampered[len(tampered) // 2] ^= 0xFF
    try:
        out2 = io.BytesIO()
        engine.decrypt_stream(io.BytesIO(bytes(tampered)), out2, key_id="default_key")
        print("篡改检测: 未检出! (异常)")
    except IntegrityError:
        print("篡改检测: 捕获 IntegrityError ✓")