lou
318 lines
12 KiB
Python
318 lines
12 KiB
Python
"""
|
|
元数据生成模块 (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 _build_vol_init(src_path: Path, args: argparse.Namespace) -> int:
|
|
"""卷压缩模式: 流式切卷 -> 逐卷 7z 压缩 -> 逐卷加密
|
|
|
|
每卷独立 context (7z-encrypt:v1:vol:N), nonce 由 context 派生,
|
|
下载端按卷号重建 context 解密, 服务端零改动。
|
|
"""
|
|
import subprocess
|
|
import tempfile
|
|
from crypto import CONTEXT
|
|
|
|
vol_size = args.vol_compress * 1024 * 1024
|
|
engine = CryptoEngine()
|
|
try:
|
|
engine.get_key(args.key_id)
|
|
except KeyNotFoundError:
|
|
engine.generate_key(args.key_id)
|
|
print(f"[密钥库] 已生成新密钥 -> {engine.keyring_path}")
|
|
|
|
# 卷工作目录: 磁盘优先 (/root), Android /tmp 是 tmpfs 但卷小无碍, 统一临时目录
|
|
workdir = Path(tempfile.gettempdir())
|
|
init_path = src_path.with_name(src_path.name + ".init.json")
|
|
manifest_path = src_path.with_name(src_path.name + "_manifest.json")
|
|
src_st = src_path.stat().st_size
|
|
print(f"=== 卷压缩模式: {src_path.name} ({src_st / 1048576:.1f} MB), 每卷 {args.vol_compress} MB ===")
|
|
chunks: list[dict[str, Any]] = []
|
|
total_sha = hashlib.sha256()
|
|
idx = 0
|
|
pbar = tqdm(total=src_st, desc="[压缩+加密]", unit="B", unit_scale=True, leave=False)
|
|
with open(src_path, "rb") as src:
|
|
while True:
|
|
raw = src.read(vol_size)
|
|
if not raw:
|
|
break
|
|
idx += 1
|
|
raw_path = workdir / f"szvol_{src_path.name}_{idx:04d}.bin"
|
|
z7_path = raw_path.with_suffix(".7z")
|
|
chunk_path = src_path.with_name(f"{src_path.name}.chunk{idx:04d}")
|
|
try:
|
|
raw_path.write_bytes(raw)
|
|
if args.vol_no_7z:
|
|
# 已压缩格式: 切卷后直接加密 (不 7z)
|
|
with open(raw_path, "rb") as s, open(chunk_path, "wb") as d:
|
|
engine.encrypt_stream(
|
|
s, d, args.key_id, context=f"{CONTEXT.decode()}:vol:{idx}".encode())
|
|
else:
|
|
r = subprocess.run(
|
|
["7z", "a", "-y", "-bd", f"-mx={args.vol_mx}", str(z7_path), str(raw_path)],
|
|
capture_output=True, text=True,
|
|
)
|
|
if r.returncode != 0:
|
|
raise MetadataError(f"卷 {idx} 7z 压缩失败: {r.stderr[-200:]}")
|
|
raw_path.unlink()
|
|
with open(z7_path, "rb") as s, open(chunk_path, "wb") as d:
|
|
engine.encrypt_stream(
|
|
s, d, args.key_id, context=f"{CONTEXT.decode()}:vol:{idx}".encode())
|
|
z7_path.unlink(missing_ok=True)
|
|
with open(chunk_path, "rb") as f:
|
|
while True:
|
|
b = f.read(1 << 20)
|
|
if not b:
|
|
break
|
|
total_sha.update(b)
|
|
chunks.append({
|
|
"index": idx,
|
|
"size": chunk_path.stat().st_size,
|
|
"sha256": sha256_of(chunk_path),
|
|
"filename": str(chunk_path),
|
|
})
|
|
finally:
|
|
raw_path.unlink(missing_ok=True)
|
|
z7_path.unlink(missing_ok=True)
|
|
pbar.update(len(raw))
|
|
pbar.close()
|
|
|
|
with open(manifest_path, "w", encoding="utf-8") as f:
|
|
json.dump(chunks, f, indent=2, ensure_ascii=False)
|
|
enc_params: dict[str, Any] = {
|
|
"alg": "cobblestone-aes256gcm", "key_id": args.key_id,
|
|
"context": CONTEXT.decode(), "compressed": True, "vol_mode": True,
|
|
}
|
|
if args.vol_no_7z:
|
|
enc_params["vol_7z"] = False # 已压缩格式: 卷未 7z, 下载端直接拼接明文
|
|
file_name = args.file_name or src_path.name
|
|
if not args.no_encrypt_name:
|
|
file_name = engine.encrypt_name(file_name)
|
|
init_json = build_init_json(
|
|
file_name=file_name,
|
|
file_size=src_st,
|
|
total_sha256=total_sha.hexdigest(),
|
|
enc=enc_params,
|
|
chunks=chunks,
|
|
)
|
|
with open(init_path, "w", encoding="utf-8") as f:
|
|
json.dump(init_json, f, indent=2, ensure_ascii=False)
|
|
print(f"[完成] {idx} 卷 -> {init_path.name}, 清单 -> {manifest_path.name}")
|
|
return 0
|
|
|
|
|
|
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_params 加 compressed, 下载端自动解压)")
|
|
ap.add_argument("--vol-compress", type=int, default=0,
|
|
help="卷压缩模式 MB: 流式切卷(每卷MB) -> 逐卷7z压缩 -> 逐卷加密 (规避大文件OOM)")
|
|
ap.add_argument("--vol-no-7z", action="store_true",
|
|
help="卷模式跳过 7z (已压缩格式: 切卷+加密不压缩)")
|
|
ap.add_argument("--vol-mx", type=int, default=9,
|
|
help="卷内 7z 压缩级别 (默认 9; 备份照片视频 mx=1 快速)")
|
|
ap.add_argument("--no-encrypt-name", action="store_true",
|
|
help="不加密文件名 (默认加密成 enc: 前缀)")
|
|
args = ap.parse_args()
|
|
|
|
src_path = Path(args.file).resolve()
|
|
if not src_path.exists():
|
|
print(f"[错误] 文件不存在: {src_path}")
|
|
return 1
|
|
|
|
# 卷压缩模式: 流式切卷 -> 逐卷 7z 压缩 -> 逐卷加密 (每卷内存 16MB 级, 规避大文件 OOM)
|
|
if args.vol_compress > 0:
|
|
return _build_vol_init(src_path, args)
|
|
|
|
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 解压
|
|
file_name = args.file_name or src_path.name
|
|
if not args.no_encrypt_name:
|
|
# 元数据零知识: 文件名加密成 'enc:' 前缀 (服务端只存密文, 客户端解密还原)
|
|
file_name = engine.encrypt_name(file_name)
|
|
init_json = build_init_json(
|
|
file_name=file_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())
|