xsk-education-management/scripts/sync_to_vps.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
from dataclasses import dataclass
from datetime import datetime
import os
from pathlib import Path
import shlex
import subprocess
import sys
import time


DEFAULT_LOCAL_DIR = Path("/Users/yangdawei/Desktop/新时空业务源数据/新时空课程记录与课时账户")
DEFAULT_REMOTE_HOST = "121.199.172.246"
DEFAULT_REMOTE_USER = "root"
DEFAULT_REMOTE_PORT = 22222
DEFAULT_REMOTE_DIR = "/root/新时空教务管理系统"
DEFAULT_IDENTITY_FILE = Path.home() / ".ssh" / "xsk_records_vps_ed25519"
SYNC_FILES = ("classnotes.txt", "学生课时账户.md")


@dataclass(frozen=True)
class Config:
    local_dir: Path
    remote_host: str
    remote_user: str
    remote_port: int
    remote_dir: str
    identity_file: Path | None
    interval: float
    use_sshpass: bool


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="同步新时空课程记录数据到 VPS")
    parser.add_argument("--once", action="store_true", help="只同步一次后退出")
    parser.add_argument("--interval", type=float, default=float(os.getenv("XSK_SYNC_INTERVAL", "3")))
    parser.add_argument("--local-dir", default=os.getenv("XSK_LOCAL_DIR", str(DEFAULT_LOCAL_DIR)))
    parser.add_argument("--remote-host", default=os.getenv("XSK_REMOTE_HOST", DEFAULT_REMOTE_HOST))
    parser.add_argument("--remote-user", default=os.getenv("XSK_REMOTE_USER", DEFAULT_REMOTE_USER))
    parser.add_argument("--remote-port", type=int, default=int(os.getenv("XSK_REMOTE_PORT", str(DEFAULT_REMOTE_PORT))))
    parser.add_argument("--remote-dir", default=os.getenv("XSK_REMOTE_DIR", DEFAULT_REMOTE_DIR))
    parser.add_argument(
        "--identity-file",
        default=os.getenv("XSK_IDENTITY_FILE", str(DEFAULT_IDENTITY_FILE) if DEFAULT_IDENTITY_FILE.exists() else ""),
        help="SSH 私钥路径；默认兼容使用历史文件 ~/.ssh/xsk_records_vps_ed25519（如果存在）",
    )
    parser.add_argument(
        "--use-sshpass",
        action="store_true",
        default=os.getenv("XSK_USE_SSHPASS", "") == "1",
        help="从 XSK_SSH_PASSWORD 读取密码并通过 sshpass 连接；推荐改用 SSH key",
    )
    return parser.parse_args()


def log(message: str) -> None:
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {message}", flush=True)


def run_command(command: list[str]) -> subprocess.CompletedProcess[str]:
    env = os.environ.copy()
    if env.get("XSK_SSH_PASSWORD") and not env.get("SSHPASS"):
        env["SSHPASS"] = env["XSK_SSH_PASSWORD"]
    result = subprocess.run(command, text=True, capture_output=True, check=False, env=env)
    if result.returncode != 0:
        safe_command = " ".join(shlex.quote(part) for part in command if part != os.getenv("XSK_SSH_PASSWORD", ""))
        raise RuntimeError(
            f"命令失败({result.returncode}): {safe_command}\n"
            f"STDOUT: {result.stdout.strip()}\nSTDERR: {result.stderr.strip()}"
        )
    return result


def base_remote(config: Config) -> str:
    return f"{config.remote_user}@{config.remote_host}"


def ssh_options(config: Config) -> list[str]:
    options = [
        "-p",
        str(config.remote_port),
        "-o",
        "StrictHostKeyChecking=accept-new",
    ]
    if config.identity_file:
        options.extend(["-i", str(config.identity_file), "-o", "IdentitiesOnly=yes"])
    return options


def ssh_prefix(config: Config) -> list[str]:
    command: list[str] = []
    if config.use_sshpass:
        password = os.getenv("XSK_SSH_PASSWORD")
        if not password:
            raise RuntimeError("启用 --use-sshpass 时必须设置 XSK_SSH_PASSWORD")
        command.extend(["sshpass", "-e"])
    command.extend(["ssh", *ssh_options(config), base_remote(config)])
    return command


def rsync_prefix(config: Config) -> list[str]:
    command: list[str] = []
    if config.use_sshpass:
        password = os.getenv("XSK_SSH_PASSWORD")
        if not password:
            raise RuntimeError("启用 --use-sshpass 时必须设置 XSK_SSH_PASSWORD")
        command.extend(["sshpass", "-e"])
    ssh_command = " ".join(shlex.quote(part) for part in ["ssh", *ssh_options(config)])
    command.extend(
        [
            "rsync",
            "-az",
            "-e",
            ssh_command,
        ]
    )
    return command


def remote_shell_quote(value: str) -> str:
    return shlex.quote(value)


def ensure_remote_dirs(config: Config) -> None:
    data_dir = f"{config.remote_dir.rstrip('/')}/data"
    run_command(ssh_prefix(config) + [f"mkdir -p {remote_shell_quote(data_dir)}"])


def sync_once(config: Config) -> None:
    ensure_remote_dirs(config)
    data_dir = f"{config.remote_dir.rstrip('/')}/data"
    for filename in SYNC_FILES:
        source = config.local_dir / filename
        if not source.exists():
            raise FileNotFoundError(f"本地文件不存在: {source}")
        temp_name = f".{filename}.tmp"
        remote_temp = f"{base_remote(config)}:{data_dir}/{temp_name}"
        run_command(rsync_prefix(config) + [str(source), remote_temp])
        run_command(
            ssh_prefix(config)
            + [
                "mv "
                f"{remote_shell_quote(data_dir + '/' + temp_name)} "
                f"{remote_shell_quote(data_dir + '/' + filename)}"
            ]
        )
        log(f"已同步 {source.name}")


def file_signature(path: Path) -> tuple[int, int]:
    stat = path.stat()
    return stat.st_mtime_ns, stat.st_size


def current_signatures(local_dir: Path) -> dict[str, tuple[int, int]]:
    return {filename: file_signature(local_dir / filename) for filename in SYNC_FILES}


def watch(config: Config) -> None:
    log("启动新时空教务管理系统数据同步监听")
    signatures: dict[str, tuple[int, int]] = {}
    while True:
        try:
            next_signatures = current_signatures(config.local_dir)
            if next_signatures != signatures:
                time.sleep(0.4)
                sync_once(config)
                signatures = current_signatures(config.local_dir)
        except Exception as exc:
            log(f"同步失败: {exc}")
        time.sleep(config.interval)


def main() -> int:
    args = parse_args()
    config = Config(
        local_dir=Path(args.local_dir).expanduser(),
        remote_host=args.remote_host,
        remote_user=args.remote_user,
        remote_port=args.remote_port,
        remote_dir=args.remote_dir,
        identity_file=Path(args.identity_file).expanduser() if args.identity_file else None,
        interval=args.interval,
        use_sshpass=args.use_sshpass,
    )
    try:
        if args.once:
            sync_once(config)
        else:
            watch(config)
    except KeyboardInterrupt:
        return 130
    except Exception as exc:
        log(f"退出: {exc}")
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())