lou
143 lines
4.6 KiB
Python
143 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
分卷器 (splitter)
|
|
模块:客户端 / 分卷器
|
|
输入:密文流(文件或 stdin)
|
|
输出:N 个卷文件 + 卷清单 (index, size, sha256)
|
|
|
|
用法:
|
|
cd /home/lou/文档/7z-encrypt
|
|
.venv/bin/python splitter.py -i index.mp4.enc -o index.mp4.enc -s 10485760
|
|
或: cat index.mp4.enc | .venv/bin/python splitter.py -o index.mp4.enc -s 10485760
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
from typing import IO, Any, Callable
|
|
|
|
|
|
def _read_block(stream: IO[bytes], size: int) -> bytes:
|
|
"""循环读取直到攒满 size 字节或 EOF。
|
|
管道/stdin 不保证一次 read 返回满 size, 必须循环攒块, 否则卷大小不固定。
|
|
"""
|
|
buf = b""
|
|
while len(buf) < size:
|
|
chunk = stream.read(size - len(buf))
|
|
if not chunk:
|
|
break
|
|
buf += chunk
|
|
return buf
|
|
|
|
|
|
def split_stream(
|
|
stream: IO[bytes],
|
|
chunk_size: int,
|
|
output_prefix: str,
|
|
total_size: int | None = None,
|
|
progress: Callable[[int], Any] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""从输入流切分卷文件, 返回清单 [{index, size, sha256, filename}]
|
|
|
|
Args:
|
|
stream: 二进制输入流 (open(...,'rb') 或 sys.stdin.buffer)
|
|
chunk_size: 每卷最大字节数
|
|
output_prefix: 输出前缀, 卷文件命名为 <prefix>.partXXXX
|
|
total_size: 输入总大小 (进度条用, 可 None)
|
|
progress: 进度回调, 每写一卷后调 progress(已写字节)
|
|
|
|
Returns:
|
|
list: 清单
|
|
"""
|
|
manifest: list[dict[str, Any]] = []
|
|
part_num = 1
|
|
done = 0
|
|
last_pct = -1
|
|
label = "分卷"
|
|
|
|
while True:
|
|
data = _read_block(stream, chunk_size)
|
|
if not data:
|
|
break
|
|
|
|
sha256_hex = hashlib.sha256(data).hexdigest()
|
|
size = len(data)
|
|
filename = f"{output_prefix}.part{part_num:04d}"
|
|
|
|
with open(filename, "wb") as f:
|
|
f.write(data)
|
|
|
|
manifest.append({
|
|
"index": part_num,
|
|
"size": size,
|
|
"sha256": sha256_hex,
|
|
"filename": filename,
|
|
})
|
|
part_num += 1
|
|
done += size
|
|
|
|
# 实时进度: 有回调交给调用方 (tqdm), 否则内部打印
|
|
if progress is not None:
|
|
progress(done)
|
|
elif total_size:
|
|
pct = done * 100 // total_size
|
|
if pct != last_pct:
|
|
last_pct = pct
|
|
sys.stdout.write(
|
|
f"\r {label}: {pct}% ({done / 1048576:.1f}/{total_size / 1048576:.1f} MB)"
|
|
)
|
|
sys.stdout.flush()
|
|
elif part_num % 16 == 0:
|
|
sys.stdout.write(f"\r {label}: 已切 {part_num} 卷 ({done / 1048576:.1f} MB)")
|
|
sys.stdout.flush()
|
|
|
|
if total_size:
|
|
sys.stdout.write(f"\r {label}: 100% ({total_size / 1048576:.1f} MB)\n")
|
|
sys.stdout.flush()
|
|
else:
|
|
sys.stdout.write(f"\r {label}: 完成, 共 {part_num - 1} 卷 ({done / 1048576:.1f} MB)\n")
|
|
sys.stdout.flush()
|
|
|
|
return manifest
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="分卷器: 把数据流切成 N 个卷, 每卷带 SHA-256")
|
|
parser.add_argument("-i", "--input", help="输入文件路径 (省略则从 stdin 读取)")
|
|
parser.add_argument("-o", "--output-prefix", required=True,
|
|
help="输出卷文件前缀 (如 'archive' -> archive.part0001)")
|
|
parser.add_argument("-s", "--chunk-size", type=int, default=10 * 1024 * 1024,
|
|
help="每卷最大字节数 (默认 10MB)")
|
|
parser.add_argument("-m", "--manifest",
|
|
help="清单 JSON 保存路径 (默认 <输出前缀>_manifest.json)")
|
|
args = parser.parse_args()
|
|
|
|
if args.chunk_size <= 0:
|
|
print(f"[错误] chunk-size 必须大于 0, 收到: {args.chunk_size}")
|
|
return 1
|
|
|
|
if args.input:
|
|
with open(args.input, "rb") as f:
|
|
total = f.seek(0, 2)
|
|
f.seek(0)
|
|
manifest = split_stream(f, args.chunk_size, args.output_prefix, total_size=total)
|
|
else:
|
|
manifest = split_stream(sys.stdin.buffer, args.chunk_size, args.output_prefix)
|
|
|
|
if not manifest:
|
|
print("[错误] 输入为空, 没有产出任何卷")
|
|
return 1
|
|
|
|
manifest_path = args.manifest if args.manifest else f"{args.output_prefix}_manifest.json"
|
|
with open(manifest_path, "w", encoding="utf-8") as f:
|
|
json.dump(manifest, f, indent=2, ensure_ascii=False)
|
|
|
|
total_mb = sum(p["size"] for p in manifest) / 1048576
|
|
print(f"分卷完成: {len(manifest)} 卷, 共 {total_mb:.1f} MB, 清单 -> {manifest_path}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|