7z-encrypt/backup.py

#!/usr/bin/env python3
"""一键备份: 相册/聊天记录/文档 -> 私有加密服务器

增量备份 (mtime+size 对比本地清单) + 批量 tar + 复用加密上传管线。
手机 ZeroTermux/PRoot 可用: /sdcard/DCIM 等外部存储直接可见。

聊天记录注意: PRoot 无 root 读不到微信数据库 (/data/data), 需先把
微信"备份与迁移/导出"的文件放到手机存储, 再配置 backup.chat 目录。

用法:
  sz-backup                    # 备份全部 (photos+documents+download+chat)
  sz-backup photos             # 只备份相册
  sz-backup list               # 已备份清单
  sz-backup restore photos     # 恢复相册
"""
import argparse
import json
import os
import subprocess
import sys
import tarfile
from datetime import datetime
from pathlib import Path
from typing import Any

from tqdm import tqdm

STATE_DIR = Path.home() / ".local/share/7z-encrypt" / "backup"
BATCH_MB = 500  # 每批 tar 上限
BATCH_FILES = 500  # 每批文件数上限
# 高压缩率类型 (聊天记录/文档) 用 -mx=9, 照片/视频批次 -mx=1 快速
COMPRESS_SUFFIX = {".db", ".sqlite", ".sqlite3", ".sql", ".txt", ".json", ".xml", ".csv", ".log", ".md", ".html"}


def _7z(args: list[str]) -> None:
    """调系统 7z (p7zip-full), 实时透传进度输出 (7z 自带百分比进度条)"""
    r = subprocess.run(["7z"] + args)
    if r.returncode != 0:
        raise RuntimeError(f"7z 失败 (exit {r.returncode})")


def _cfg() -> dict[str, Any]:
    """读配置文件 backup 段 (XDG, 同客户端)"""
    env = os.environ.get("SZ_CONFIG")
    if env:
        p = Path(env).expanduser()
    else:
        xdg = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config")
        p = Path(xdg) / "7z-encrypt" / "config.json"
    try:
        return dict(json.loads(p.read_text(encoding="utf-8")).get("backup", {}))
    except (OSError, json.JSONDecodeError):
        return {}


def _manifest_path(name: str) -> Path:
    return STATE_DIR / f"{name}.json"


def _load_manifest(name: str) -> dict[str, Any]:
    try:
        return json.loads(_manifest_path(name).read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {"files": {}, "batches": {}}


def _save_manifest(name: str, m: dict[str, Any]) -> None:
    _manifest_path(name).parent.mkdir(parents=True, exist_ok=True)
    _manifest_path(name).write_text(json.dumps(m, ensure_ascii=False, indent=1), encoding="utf-8")


def _server() -> str:
    """服务器地址: 配置 server.url (备份复用客户端配置)"""
    env = os.environ.get("SZ_CONFIG")
    if env:
        p = Path(env).expanduser()
    else:
        xdg = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config")
        p = Path(xdg) / "7z-encrypt" / "config.json"
    try:
        return str(json.loads(p.read_text(encoding="utf-8")).get("server", {}).get("url", ""))
    except (OSError, json.JSONDecodeError):
        return ""


def _token() -> str:
    """登录 token: 配置文件 server.token"""
    env = os.environ.get("SZ_CONFIG")
    if env:
        p = Path(env).expanduser()
    else:
        xdg = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config")
        p = Path(xdg) / "7z-encrypt" / "config.json"
    try:
        return str(json.loads(p.read_text(encoding="utf-8")).get("server", {}).get("token", ""))
    except (OSError, json.JSONDecodeError):
        return ""


def _src_dirs(name: str) -> list[Path]:
    """备份源目录: 配置 override 默认"""
    defaults = {
        "photos": ["~/storage/dcim", "~/storage/pictures", "/sdcard/DCIM", "/sdcard/Pictures"],
        "documents": ["/sdcard/Documents"],
        "download": ["/sdcard/Download"],
        "chat": [],
    }
    dirs = _cfg().get(name, defaults.get(name, []))
    return [Path(d).expanduser() for d in dirs]


def _scan_new(name: str, manifest: dict[str, Any]) -> list[Path]:
    """扫描新增/修改文件 (mtime+size 对比)"""
    known = manifest.get("files", {})
    new_files: list[Path] = []
    for src in _src_dirs(name):
        if not src.exists():
            continue
        for p in src.rglob("*"):
            if not p.is_file():
                continue
            try:
                st = p.stat()
            except OSError:
                continue
            key = str(p)
            if key in known and known[key][0] == st.st_size and known[key][1] == int(st.st_mtime):
                continue
            new_files.append(p)
    return new_files


def _compress(path: Path) -> bool:
    return path.suffix.lower() in COMPRESS_SUFFIX


def _batch(files: list[Path]) -> list[list[Path]]:
    """按大小/数量分批"""
    batches: list[list[Path]] = []
    cur: list[Path] = []
    total = 0
    for f in files:
        size = f.stat().st_size
        if cur and (total + size > BATCH_MB * 1024 * 1024 or len(cur) >= BATCH_FILES):
            batches.append(cur)
            cur, total = [], 0
        cur.append(f)
        total += size
    if cur:
        batches.append(cur)
    return batches


def _upload_tar(tar_path: Path, chunk_mb: int) -> str:
    """加密 -> 分卷 -> 上传 (复用 metadata/splitter/crypto/transfer 库, 编译版自包含)"""
    from metadata import build_init_json, sha256_of
    from splitter import split_stream
    from crypto import CryptoEngine, KeyNotFoundError
    import transfer as transfer_mod

    engine = CryptoEngine()
    try:
        engine.get_key("default_key")
    except KeyNotFoundError:
        engine.generate_key("default_key")
    chunk_size = chunk_mb * 1024 * 1024
    enc_path = tar_path.with_suffix(tar_path.suffix + ".enc")
    print(f"[备份] 加密批次 ...")
    with open(tar_path, "rb") as src, open(enc_path, "wb") as dst:
        pbar = tqdm(total=tar_path.stat().st_size, desc="[加密]", unit="B",
                    unit_scale=True, leave=False)
        enc_params = engine.encrypt_stream(src, dst, "default_key",
                                           progress=lambda n: pbar.update(n))
        pbar.close()
    print(f"[备份] 分卷 (每卷 {chunk_mb}MB) ...")
    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()
    total_sha = sha256_of(enc_path)
    init_json = build_init_json(
        file_name=engine.encrypt_name(tar_path.name),
        file_size=tar_path.stat().st_size,
        total_sha256=total_sha,
        enc=enc_params,
        chunks=manifest,
    )
    chunk_files = {ch["index"]: Path(ch["filename"]).resolve() for ch in manifest}
    client = transfer_mod.TransferClient(_server(), token=_token())
    receipt = client.transfer(init_json, chunk_files)
    enc_path.unlink(missing_ok=True)
    for ch in manifest:
        Path(ch["filename"]).unlink(missing_ok=True)
    return str(receipt["file_id"])


def _prune_manifest(name: str, manifest: dict[str, Any]) -> None:
    """服务器同步: 服务器已不存在的批次从 manifest 移除, 其文件重新纳入增量

    服务器文件被删 (手动/清理) 后 manifest 仍认为已备份会导致"无增量更新"
    但实际恢复不到——以服务器 file_id 为准剪枝, 下次备份自动重传。
    """
    from transfer import TransferClient
    try:
        client = TransferClient(_server(), token=_token())
        server_ids = {f["file_id"] for f in client.list_files()}
    except Exception:
        return  # 服务器不可达: 保留 manifest, 不阻塞备份
    batches = manifest.get("batches", {})
    stale = [bn for bn, b in batches.items() if b.get("server_id") not in server_ids]
    if not stale:
        return
    for bn in stale:
        for fp in batches[bn].get("files", []):
            manifest.setdefault("files", {}).pop(fp, None)
        batches.pop(bn)
    print(f"[备份] 服务器同步: 移除 {len(stale)} 个失效批次 (服务器文件已删), 将重新备份")
    _save_manifest(name, manifest)


def backup(name: str, chunk_mb: int = 10) -> int:
    manifest = _load_manifest(name)
    _prune_manifest(name, manifest)
    files = _scan_new(name, manifest)
    if not files:
        print(f"[备份] {name}: 无新增文件, 已是最新")
        return 0
    batches = _batch(files)
    print(f"[备份] {name}: 新增 {len(files)} 个文件 ({sum(f.stat().st_size for f in files) / 1048576:.1f} MB), {len(batches)} 批")
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    for i, batch in enumerate(batches, 1):
        tar_path = STATE_DIR / f"{name}_{ts}_{i:02d}.tar"
        tar_path.parent.mkdir(parents=True, exist_ok=True)
        with tarfile.open(tar_path, "w") as tf:
            for f in batch:
                try:
                    tf.add(f, arcname=str(f))
                except (OSError, tarfile.TarError) as e:
                    print(f"  [跳过] {f}: {e}")
        # 7z 压缩 (与上传文件一致): 含文本类 -mx=9, 照片视频 -mx=1
        mx = 9 if any(_compress(f) for f in batch) else 1
        seven = tar_path.with_suffix(".tar.7z")
        print(f"[备份] 7z 压缩 (mx={mx}) ...")
        _7z(["a", "-y", "-bd", f"-mx={mx}", str(seven), str(tar_path)])
        tar_path.unlink(missing_ok=True)
        tar_path = seven
        size_mb = tar_path.stat().st_size / 1048576
        print(f"[备份] 批次 {i}/{len(batches)}: {tar_path.name} ({size_mb:.1f} MB), 上传中 ...")
        try:
            fid = _upload_tar(tar_path, chunk_mb)
        except RuntimeError as e:
            print(f"[错误] {e}")
            return 1
        files_in_batch = [str(f) for f in batch]
        for f in batch:
            try:
                st = f.stat()
                manifest.setdefault("files", {})[str(f)] = [st.st_size, int(st.st_mtime)]
            except OSError:
                pass
        manifest.setdefault("batches", {})[tar_path.name] = {
            "server_id": fid, "size": tar_path.stat().st_size, "files": files_in_batch,
        }
        _save_manifest(name, manifest)
        print(f"[备份] 批次 {i} 完成: file_id={fid}")
        tar_path.unlink(missing_ok=True)
    print(f"[备份] {name} 完成, 共 {len(files)} 个文件")
    return 0


def _download_batch(file_id: str, out_dir: Path) -> Path:
    """下载 + 解密备份批次 (复用 transfer/crypto 库), 返回还原的 tar 路径"""
    import transfer as transfer_mod
    from crypto import CryptoEngine

    client = transfer_mod.TransferClient(_server(), token=_token())
    enc_path, enc_params = client.download(file_id, out_dir / ".sz-backup.enc")
    engine = CryptoEngine()
    final = out_dir / str(enc_params.get("file_name", "backup.tar"))
    with open(enc_path, "rb") as src, open(final, "wb") as dst:
        pbar = tqdm(total=enc_path.stat().st_size, desc="[解密]", unit="B",
                    unit_scale=True, leave=False)
        engine.decrypt_stream(src, dst, enc_params["key_id"],
                              enc_params.get("context", "7z-encrypt:v1").encode(),
                              progress=lambda n: pbar.update(n))
        pbar.close()
    enc_path.unlink(missing_ok=True)
    return final


def restore(name: str, out_root: Path | None = None) -> int:
    manifest = _load_manifest(name)
    batches = manifest.get("batches", {})
    if not batches:
        print(f"[恢复] {name}: 无备份记录")
        return 1
    print(f"[恢复] {name}: {len(batches)} 个备份批次")
    for tar_name, info in batches.items():
        dl = STATE_DIR / f"restore_{tar_name}"
        dl.mkdir(parents=True, exist_ok=True)
        try:
            tar_src = _download_batch(info["server_id"], dl)
        except Exception as e:
            print(f"[错误] 下载 {tar_name} 失败: {e}")
            return 1
        # 7z 解压 -> tar -> 解包 (路径校验: 只解回 manifest 记录的文件)
        if tar_src.suffix == ".7z":
            _7z(["x", "-y", "-bd", f"-o{dl}", str(tar_src)])
            tar_src.unlink(missing_ok=True)
            tar_src = dl / tar_name
        if not tar_src.exists():
            print(f"[错误] 未找到 tar: {dl}")
            return 1
        with tarfile.open(tar_src) as tf:
            for member in tf.getmembers():
                if member.name not in info.get("files", []):
                    continue
                if out_root is not None:
                    rel = member.name.lstrip("/")
                    tf.extract(member, path=out_root / rel, filter="data")
                else:
                    tf.extract(member, path=Path(member.name).parent, filter="data")
        print(f"[恢复] 批次完成: {tar_name}")
    print(f"[恢复] {name} 完成")
    return 0


def list_backups() -> int:
    if not STATE_DIR.exists():
        print("[备份] 暂无备份记录")
        return 0
    for mp in sorted(STATE_DIR.glob("*.json")):
        name = mp.stem
        m = _load_manifest(name)
        batches = m.get("batches", {})
        total = sum(b["size"] for b in batches.values())
        print(f"  {name}: {len(m.get('files', {}))} 文件, {len(batches)} 批, {total / 1048576:.1f} MB 已备份")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="一键备份: 相册/聊天记录/文档 -> 加密服务器")
    ap.add_argument("action", nargs="?", default="all", help="备份项: photos/documents/download/chat/all/list/restore")
    ap.add_argument("target", nargs="?", help="restore 的备份项")
    ap.add_argument("--chunk-size", type=int, default=10, help="分卷大小 MB (默认 10)")
    args = ap.parse_args()

    if args.action == "list":
        return list_backups()
    if args.action == "restore":
        if not args.target:
            print("[错误] 用法: sz-backup restore <photos|documents|download|chat>")
            return 1
        return restore(args.target)
    names = ["photos", "documents", "download", "chat"] if args.action == "all" else [args.action]
    if args.action not in ("photos", "documents", "download", "chat", "all"):
        print(f"[错误] 未知备份项: {args.action} (photos/documents/download/chat/all)")
        return 1
    ok = True
    for n in names:
        if not _src_dirs(n):
            print(f"[备份] {n}: 无源目录 (config backup.{n} 为空, 跳过)")
            continue
        ok = backup(n, args.chunk_size) == 0 and ok
    return 0 if ok else 1


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