7z-encrypt/metadata.py

"""
元数据生成模块 (metadata)
模块:客户端 / 元数据生成
输入:文件信息 + 加密参数 + 卷清单
输出:init json (file_name/file_size/chunk_count/total_sha256/enc/chunks)

init json 是服务端建任务的唯一依据 (POST /init 请求体),字段缺一不可,
组装后先自校验再输出。chunks 剥离本地 filename (服务端不关心客户端路径)。

用法:
    cd /home/lou/文档/7z-encrypt
    .venv/bin/python metadata.py [文件] [--chunk-size MB] [--key-id 密钥ID]
"""

import argparse
import hashlib
import json
import sys
from pathlib import Path
from typing import Any

from tqdm import tqdm

from crypto import CryptoEngine, KeyNotFoundError
from splitter import split_stream

SHA256_HEX_LEN = 64
DEFAULT_CHUNK_MB = 10


class MetadataError(ValueError):
    """init json 字段缺失或非法"""


def _require_str(obj: dict[str, Any], key: str) -> str:
    v = obj.get(key)
    if not isinstance(v, str) or not v:
        raise MetadataError(f"字段 '{key}' 缺失或非空字符串")
    return v


def _require_pos_int(obj: dict[str, Any], key: str) -> int:
    v = obj.get(key)
    if not isinstance(v, int) or isinstance(v, bool) or v <= 0:
        raise MetadataError(f"字段 '{key}' 缺失或非正整数")
    return v


def _require_sha256(obj: dict[str, Any], key: str) -> str:
    v = _require_str(obj, key)
    if len(v) != SHA256_HEX_LEN:
        raise MetadataError(f"字段 '{key}' 长度不是 {SHA256_HEX_LEN} (SHA-256 hex)")
    return v


def build_init_json(
    file_name: str,
    file_size: int,
    total_sha256: str,
    enc: dict[str, str],
    chunks: list[dict[str, Any]],
) -> dict[str, Any]:
    """组装并校验 init json

    Args:
        file_name: 原始文件名
        file_size: 原始文件字节数
        total_sha256: 密文整体 SHA-256 (服务端合并后验证)
        enc: 加密参数 {alg, key_id, context} (来自 crypto.encrypt_stream)
        chunks: 卷清单 (来自 splitter, 含 filename; 组装时剥离)

    Returns:
        dict: init json {file_name, file_size, chunk_count, total_sha256, enc, chunks}

    Raises:
        MetadataError: 任一必填字段缺失/非法/卷序号不连续
    """
    # enc 校验
    if not enc:
        raise MetadataError("字段 'enc' 缺失或为空")
    for k in ("alg", "key_id", "context"):
        _require_str(enc, k)

    # chunks 校验 (序号必须 1..N 连续)
    if not chunks:
        raise MetadataError("字段 'chunks' 缺失或为空")
    clean_chunks: list[dict[str, Any]] = []
    for i, ch in enumerate(chunks, start=1):
        if ch.get("index") != i:
            raise MetadataError(
                f"chunks[{i}] 序号不连续 (期望 {i}, 实际 {ch.get('index')})"
            )
        clean_chunks.append({
            "index": i,
            "size": _require_pos_int(ch, "size"),
            "sha256": _require_sha256(ch, "sha256"),
        })

    # 整体校验
    if not file_name:
        raise MetadataError("file_name 缺失或为空")
    if isinstance(file_size, bool) or file_size <= 0:
        raise MetadataError("file_size 缺失或非法")
    if len(total_sha256) != SHA256_HEX_LEN:
        raise MetadataError(f"total_sha256 长度不是 {SHA256_HEX_LEN} (SHA-256 hex)")

    return {
        "file_name": file_name,
        "file_size": file_size,
        "chunk_count": len(clean_chunks),
        "total_sha256": total_sha256,
        "enc": enc,
        "chunks": clean_chunks,
    }


def sha256_of(path: str | Path) -> str:
    """流式计算文件 SHA-256 (不整块进内存)"""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while True:
            chunk = f.read(1 << 16)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser(description="元数据生成: 加密 -> 分卷 -> 组装 init json")
    ap.add_argument("file", nargs="?", default="index.mp4", help="要处理的文件 (默认 index.mp4)")
    ap.add_argument("--chunk-size", type=int, default=DEFAULT_CHUNK_MB,
                    help=f"每卷大小 MB (默认 {DEFAULT_CHUNK_MB})")
    ap.add_argument("--key-id", default="default_key", help="密钥ID (默认 default_key)")
    ap.add_argument("--file-name", help="init json 的 file_name (默认=输入文件名, 压缩场景传原名)")
    ap.add_argument("--compressed", action="store_true",
                    help="标记已 7z 压缩 (enc 参数带 compressed=true, 下载端据此解压)")
    args = ap.parse_args()

    src_path = Path(args.file).resolve()
    if not src_path.exists():
        print(f"[错误] 文件不存在: {src_path}")
        return 1

    chunk_size = args.chunk_size * 1024 * 1024
    enc_path = src_path.with_name(src_path.name + ".enc")
    manifest_path = src_path.with_name(src_path.name + "_manifest.json")
    init_path = src_path.with_name(src_path.name + ".init.json")

    print("=== 元数据生成 (加密 -> 分卷 -> init json) ===")
    print(f"文件: {src_path.name} ({src_path.stat().st_size / 1048576:.1f} MB), 卷大小: {args.chunk_size} MB")

    # 1. 加密 -> 加密参数
    engine = CryptoEngine()
    try:
        engine.get_key(args.key_id)
    except KeyNotFoundError:
        engine.generate_key(args.key_id)
        print(f"[密钥库] 已生成新密钥 -> {engine.keyring_path}")
    with open(src_path, "rb") as src, open(enc_path, "wb") as dst:
        pbar = tqdm(total=src_path.stat().st_size, desc="[加密]", unit="B",
                    unit_scale=True, leave=False)
        enc_params = engine.encrypt_stream(src, dst, args.key_id,
                                           progress=lambda n: pbar.update(n))
        pbar.close()
    print(f"[1/4] 加密完成 -> {enc_path.name} ({enc_path.stat().st_size / 1048576:.1f} MB)")

    # 2. 分卷 -> 卷清单
    with open(enc_path, "rb") as f:
        pbar = tqdm(total=enc_path.stat().st_size, desc="[分卷]", unit="B",
                    unit_scale=True, leave=False)
        manifest = split_stream(f, chunk_size, str(enc_path),
                                total_size=enc_path.stat().st_size,
                                progress=lambda n: pbar.update(n))
        pbar.close()
    with open(manifest_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)
    print(f"[2/4] 分卷完成: {len(manifest)} 卷, 清单 -> {manifest_path.name}")

    # 3. 密文整体 SHA-256
    total_sha256 = sha256_of(enc_path)
    print(f"[3/4] 密文 SHA-256: {total_sha256}")

    # 4. 组装 init json
    if args.compressed:
        enc_params["compressed"] = True  # 下载端据此 7z 解压
    init_json = build_init_json(
        file_name=args.file_name or src_path.name,
        file_size=src_path.stat().st_size,
        total_sha256=total_sha256,
        enc=enc_params,
        chunks=manifest,
    )
    with open(init_path, "w", encoding="utf-8") as f:
        json.dump(init_json, f, indent=2, ensure_ascii=False)
    print(f"[4/4] init json -> {init_path.name}")
    print("\n=== init json ===")
    print(json.dumps(init_json, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())