feat: 图片迁移脚本

2025-12-20 16:41:17 +08:00
parent 64e5c4fe26
commit c7b22884ed
6 changed files with 1396 additions and 0 deletions
--- a/scripts/image_migrate.py
+++ b/scripts/image_migrate.py
@@ -0,0 +1,751 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+图片迁移脚本 - 通过SSH或本地目录迁移图片到远程服务器
+保持相对路径结构一致
+"""
+
+import argparse
+import json
+import os
+import shutil
+import sys
+import tempfile
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+# paramiko 导入被延迟到实际需要时
+paramiko: Any = None
+
+
+def ensure_paramiko() -> None:
+    """确保paramiko已导入，如果未安装则报错"""
+    global paramiko
+    if paramiko is None:
+        try:
+            import paramiko as paramiko_module
+
+            paramiko = paramiko_module
+        except ImportError:
+            print("错误: 需要安装paramiko库")
+            print("请运行: pip install paramiko")
+            sys.exit(1)
+
+
+class ServerConfig:
+    """服务器配置（支持SSH或本地目录）"""
+
+    def __init__(
+        self,
+        config_type: str = "ssh",  # "ssh" 或 "local"
+        host: str = "",
+        port: int = 22,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        key_file: Optional[str] = None,
+        base_dir: Optional[str] = None,
+    ):
+        """
+        初始化服务器配置
+
+        Args:
+            config_type: 配置类型，'ssh' 表示SSH服务器，'local' 表示本地目录
+            host: 服务器地址（对于local类型可为None）
+            port: 服务器端口（仅SSH类型）
+            username: 用户名（仅SSH类型）
+            password: 密码（仅SSH类型）
+            key_file: 私钥文件路径（仅SSH类型）
+            base_dir: 基础目录路径
+        """
+        self.config_type = config_type
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.key_file = key_file
+        self.base_dir = base_dir.rstrip("/") if base_dir else ""
+
+    def to_dict(self) -> Dict:
+        """转换为字典"""
+        return {
+            "type": self.config_type,
+            "host": self.host,
+            "port": self.port,
+            "username": self.username,
+            "password": self.password,
+            "key_file": self.key_file,
+            "base_dir": self.base_dir,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "ServerConfig":
+        """从字典创建"""
+        ssh_data = data.get("ssh", {})
+        return cls(
+            config_type=data.get("type", "ssh"),
+            host=ssh_data.get("host", "localhost"),
+            port=ssh_data.get("port", 22),
+            username=ssh_data.get("username", ""),
+            password=ssh_data.get("password", ""),
+            key_file=ssh_data.get("key_file", ""),
+            base_dir=data.get("base_dir", ""),
+        )
+
+    def is_local(self) -> bool:
+        """是否是本地目录配置"""
+        return self.config_type == "local"
+
+    def __str__(self) -> str:
+        if self.is_local():
+            return f"LocalDirectory(base_dir={self.base_dir})"
+        else:
+            return f"SSHServer(host={self.host}:{self.port}, user={self.username}, base_dir={self.base_dir})"
+
+
+class ImageMigrator:
+    """图片迁移器"""
+
+    def __init__(
+        self,
+        source_config: ServerConfig,
+        target_config: ServerConfig,
+        verbose: bool = False,
+        overwrite: bool = False,
+    ) -> None:
+        """
+        初始化迁移器
+
+        Args:
+            source_config: 源服务器配置（可以是本地目录或SSH服务器）
+            target_config: 目标服务器配置（必须为SSH服务器）
+            verbose: 是否显示详细输出
+            overwrite: 是否覆盖已存在的文件
+        """
+        self.source_config = source_config
+        self.target_config = target_config
+        self.verbose = verbose
+        self.overwrite = overwrite
+
+        # SSH客户端连接（仅用于SSH类型的服务器）
+        self.source_client: Any = None
+        self.target_client: Any = None
+        self.source_sftp: Any = None
+        self.target_sftp: Any = None
+
+        # 统计信息
+        self.stats = {
+            "total": 0,
+            "success": 0,
+            "failed": 0,
+            "skipped": 0,
+            "bytes_transferred": 0,
+        }
+
+    def connect(self) -> bool:
+        """连接到服务器"""
+        try:
+            # 确保paramiko已导入
+            ensure_paramiko()
+
+            # 连接目标服务器（必须为SSH）
+            if self.target_config.is_local():
+                print("错误: 目标服务器必须为SSH服务器，不能是本地目录")
+                return False
+
+            self.target_client = paramiko.SSHClient()
+            self.target_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+            if self.verbose:
+                print(
+                    f"连接到目标服务器: {self.target_config.host}:{self.target_config.port}"
+                )
+
+            if self.target_config.key_file:
+                key = paramiko.RSAKey.from_private_key_file(self.target_config.key_file)
+
+                self.target_client.connect(
+                    hostname=self.target_config.host,
+                    port=self.target_config.port,
+                    username=self.target_config.username,
+                    pkey=key,
+                )
+            else:
+                self.target_client.connect(
+                    hostname=self.target_config.host,
+                    port=self.target_config.port,
+                    username=self.target_config.username,
+                    password=self.target_config.password,
+                )
+
+            self.target_sftp = self.target_client.open_sftp()
+
+            # 连接源服务器（如果是SSH类型）
+            if not self.source_config.is_local():
+                self.source_client = paramiko.SSHClient()
+                self.source_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+                if self.verbose:
+                    print(
+                        f"连接到源服务器: {self.source_config.host}:{self.source_config.port}"
+                    )
+
+                if self.source_config.key_file:
+                    key = paramiko.RSAKey.from_private_key_file(
+                        self.source_config.key_file
+                    )
+
+                    self.source_client.connect(
+                        hostname=self.source_config.host,
+                        port=self.source_config.port,
+                        username=self.source_config.username,
+                        pkey=key,
+                    )
+                else:
+                    self.source_client.connect(
+                        hostname=self.source_config.host,
+                        port=self.source_config.port,
+                        username=self.source_config.username,
+                        password=self.source_config.password,
+                    )
+
+                self.source_sftp = self.source_client.open_sftp()
+
+            if self.verbose:
+                source_type = (
+                    "本地目录" if self.source_config.is_local() else "SSH服务器"
+                )
+                print(f"连接成功! 源: {source_type}, 目标: SSH服务器")
+            return True
+
+        except Exception as e:
+            print(f"连接失败: {e}")
+            self.close()
+            return False
+
+    def close(self) -> None:
+        """关闭所有连接"""
+        if self.source_sftp:
+            self.source_sftp.close()
+        if self.source_client:
+            self.source_client.close()
+        if self.target_sftp:
+            self.target_sftp.close()
+        if self.target_client:
+            self.target_client.close()
+
+        if self.verbose:
+            print("连接已关闭")
+
+    def ensure_target_directory(self, remote_dir: str) -> bool:
+        """确保目标目录存在（递归创建）"""
+        try:
+            # 检查目录是否存在
+            try:
+                self.target_sftp.stat(remote_dir)
+
+                return True
+            except FileNotFoundError:
+                # 目录不存在，递归创建
+                parts = remote_dir.strip("/").split("/")
+                current_path = ""
+
+                for part in parts:
+                    if not part:
+                        continue
+                    current_path = f"{current_path}/{part}" if current_path else part
+                    try:
+                        self.target_sftp.stat(current_path)
+
+                    except FileNotFoundError:
+                        if self.verbose:
+                            print(f"创建目录: {current_path}")
+                        self.target_sftp.mkdir(current_path)
+
+                return True
+        except Exception as e:
+            print(f"创建目录失败 {remote_dir}: {e}")
+            return False
+
+    def get_source_file_info(self, source_path: str) -> Tuple[bool, Optional[int]]:
+        """
+        获取源文件信息
+
+        Returns:
+            Tuple[exists: bool, size: Optional[int]]
+        """
+        try:
+            if self.source_config.is_local():
+                # 本地文件
+                if os.path.exists(source_path):
+                    return True, os.path.getsize(source_path)
+                else:
+                    return False, 0
+            else:
+                # SSH远程文件
+                source_attr = self.source_sftp.stat(source_path)
+                return True, source_attr.st_size
+
+        except Exception:
+            return False, 0
+
+    def read_source_file(self, source_path: str, temp_path: str) -> bool:
+        """
+        读取源文件到临时文件
+
+        Returns:
+            bool: 是否成功
+        """
+        try:
+            if self.source_config.is_local():
+                # 从本地文件复制
+                shutil.copy2(source_path, temp_path)
+                return True
+            else:
+                # 从SSH服务器下载
+                self.source_sftp.get(source_path, temp_path)
+
+                return True
+        except Exception as e:
+            print(f"读取源文件失败 {source_path}: {e}")
+            return False
+
+    def transfer_file(self, relative_path: str) -> bool:
+        """
+        传输单个文件
+
+        Args:
+            relative_path: 相对路径（相对于基础目录）
+
+        Returns:
+            bool: 是否成功
+        """
+        try:
+            # 构建完整路径
+            source_path = f"{self.source_config.base_dir}/{relative_path}"
+            target_path = f"{self.target_config.base_dir}/{relative_path}"
+
+            if self.verbose:
+                source_type = "本地" if self.source_config.is_local() else "远程"
+                print(f"传输: {relative_path}")
+                print(f"  源({source_type}): {source_path}")
+                print(f"  目标(远程): {target_path}")
+
+            # 检查源文件是否存在并获取大小
+            source_exists, file_size = self.get_source_file_info(source_path)
+            if not source_exists:
+                print(f"错误: 源文件不存在: {source_path}")
+                self.stats["failed"] += 1
+                return False
+
+            # 检查目标文件是否已存在（根据overwrite选项处理）
+            try:
+                self.target_sftp.stat(target_path)
+
+                if not self.overwrite:
+                    if self.verbose:
+                        print(f"跳过: 目标文件已存在: {target_path}")
+                    self.stats["skipped"] += 1
+                    return True
+                else:
+                    if self.verbose:
+                        print(f"覆盖: 目标文件已存在: {target_path}")
+            except FileNotFoundError:
+                pass
+
+            # 确保目标目录存在
+            target_dir = os.path.dirname(target_path)
+            if target_dir and not self.ensure_target_directory(target_dir):
+                print(f"错误: 无法创建目标目录: {target_dir}")
+                self.stats["failed"] += 1
+                return False
+
+            # 传输文件
+            start_time = time.time()
+            temp_path = None
+
+            try:
+                # 创建临时文件
+                with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+                    temp_path = temp_file.name
+
+                # 读取源文件到临时文件
+                if not self.read_source_file(source_path, temp_path):
+                    self.stats["failed"] += 1
+                    return False
+
+                # 上传到目标服务器
+                self.target_sftp.put(temp_path, target_path)
+
+            finally:
+                # 清理临时文件
+                if temp_path and os.path.exists(temp_path):
+                    try:
+                        os.remove(temp_path)
+                    except Exception:
+                        pass
+
+            # 记录统计
+            transfer_time = time.time() - start_time
+            speed = (file_size or 0) / transfer_time / 1024 if transfer_time > 0 else 0
+
+            if self.verbose:
+                print(
+                    f"  完成: {file_size} bytes, 耗时: {transfer_time:.2f}s, 速度: {speed:.2f} KB/s"
+                )
+
+            self.stats["success"] += 1
+            self.stats["bytes_transferred"] += file_size or 0
+            return True
+
+        except Exception as e:
+            print(f"传输失败 {relative_path}: {e}")
+            self.stats["failed"] += 1
+            return False
+
+    def migrate_images(self, image_paths: List[str]) -> Dict[str, int]:
+        """
+        迁移图片列表
+
+        Args:
+            image_paths: 相对路径列表
+
+        Returns:
+            Dict: 统计信息
+        """
+        if not self.connect():
+            return self.stats
+
+        try:
+            self.stats["total"] = len(image_paths)
+
+            source_type = "本地目录" if self.source_config.is_local() else "SSH服务器"
+            print(f"开始从{source_type}迁移 {len(image_paths)} 个图片到SSH服务器...")
+
+            for i, relative_path in enumerate(image_paths, 1):
+                # 清理路径
+                relative_path = relative_path.strip()
+                if not relative_path:
+                    continue
+
+                # 显示进度
+                print(f"[{i}/{len(image_paths)}] {relative_path}", end="")
+
+                # 传输文件
+                success = self.transfer_file(relative_path)
+
+                if success:
+                    print(" ✓")
+                else:
+                    print(" ✗")
+
+            # 打印统计信息
+            print("\n" + "=" * 50)
+            print("迁移完成!")
+            print(f"总计: {self.stats['total']}")
+            print(f"成功: {self.stats['success']}")
+            print(f"失败: {self.stats['failed']}")
+            print(f"跳过: {self.stats['skipped']}")
+            print(f"传输字节: {self.stats['bytes_transferred']:,}")
+
+            return self.stats
+
+        finally:
+            self.close()
+
+
+def load_config(config_file: str) -> Dict:
+    """加载配置文件"""
+    try:
+        with open(config_file, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        return config
+    except Exception as e:
+        print(f"加载配置文件失败 {config_file}: {e}")
+        return {}
+
+
+def save_config(config_file: str, config: Dict):
+    """保存配置文件"""
+    try:
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(config, f, indent=2, ensure_ascii=False)
+        print(f"配置文件已保存: {config_file}")
+    except Exception as e:
+        print(f"保存配置文件失败 {config_file}: {e}")
+
+
+def create_example_config(config_file: str):
+    """创建示例配置文件"""
+    example_config = {
+        "source": {
+            "type": "local",  # 可以是 "local" 或 "ssh"
+            "base_dir": "/path/to/local/images",
+            # 如果是SSH类型，需要以下配置
+            "ssh": {
+                "host": "source-server.example.com",
+                "port": 22,
+                "username": "username",
+                "password": "your_password_here",
+                "key_file": "/path/to/private/key",
+            },
+        },
+        "target": {
+            "type": "ssh",  # 目标必须是SSH类型
+            "base_dir": "/var/www/html/images",
+            "ssh": {
+                "host": "target-server.example.com",
+                "port": 22,
+                "username": "username",
+                "password": "your_password_here",
+                "key_file": "/path/to/private/key",
+            },
+        },
+    }
+
+    save_config(config_file, example_config)
+    print("已创建示例配置文件")
+    print("注意: 源服务器可以是本地目录(local)或SSH服务器(ssh)")
+    print("      目标服务器必须是SSH服务器(ssh)")
+
+
+def read_image_paths(image_list_file: str) -> List[str]:
+    """从文件读取图片路径列表"""
+    paths = []
+    try:
+        with open(image_list_file, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#"):
+                    paths.append(line)
+        print(f"从文件读取 {len(paths)} 个图片路径")
+    except Exception as e:
+        print(f"读取图片路径文件失败 {image_list_file}: {e}")
+    return paths
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="从本地目录或SSH服务器迁移图片到远程SSH服务器，保持相对路径结构",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+示例:
+  # 使用配置文件迁移（源为本地目录）
+  python image_migrate.py --config config.json --input images.txt
+
+  # 使用配置文件迁移（源为SSH服务器）
+  # 配置文件中 source.type="ssh"
+
+  # 直接指定参数迁移（源为本地目录）
+  python image_migrate.py \\
+    --source-type local --source-dir /path/to/local/images \\
+    --target-host dst.example.com --target-user user2 --target-dir /home/user/images \\
+    --input images.txt
+
+  # 直接指定参数迁移（源为SSH服务器）
+  python image_migrate.py \\
+    --source-type ssh --source-host src.example.com --source-user user1 --source-dir /var/www/images \\
+    --target-host dst.example.com --target-user user2 --target-dir /home/user/images \\
+    --input images.txt
+
+  # 创建示例配置文件
+  python image_migrate.py --create-config config.json
+
+  # 从命令行直接指定图片路径
+  python image_migrate.py --config config.json images/product1.jpg images/product2.jpg
+        """,
+    )
+
+    # 配置文件选项
+    parser.add_argument("--config", help="配置文件路径")
+    parser.add_argument("--create-config", metavar="FILE", help="创建示例配置文件")
+
+    # 源服务器选项
+    parser.add_argument(
+        "--source-type",
+        choices=["local", "ssh"],
+        default="local",
+        help="源服务器类型: local(本地目录) 或 ssh(SSH服务器)，默认: local",
+    )
+    parser.add_argument("--source-host", help="源服务器地址（仅SSH类型需要）")
+    parser.add_argument(
+        "--source-port",
+        type=int,
+        default=22,
+        help="源服务器端口（仅SSH类型，默认: 22）",
+    )
+    parser.add_argument("--source-user", help="源服务器用户名（仅SSH类型需要）")
+    parser.add_argument("--source-password", help="源服务器密码（仅SSH类型需要）")
+    parser.add_argument("--source-key", help="源服务器私钥文件路径（仅SSH类型需要）")
+    parser.add_argument("--source-dir", help="源服务器图片基础目录")
+
+    # 目标服务器选项
+    parser.add_argument("--target-host", help="目标服务器地址")
+    parser.add_argument(
+        "--target-port", type=int, default=22, help="目标服务器端口（默认: 22）"
+    )
+    parser.add_argument("--target-user", help="目标服务器用户名")
+    parser.add_argument("--target-password", help="目标服务器密码")
+    parser.add_argument("--target-key", help="目标服务器私钥文件路径")
+    parser.add_argument("--target-dir", help="目标服务器图片基础目录")
+
+    # 输入选项
+    parser.add_argument(
+        "--input", "-i", help="包含图片路径列表的文件（每行一个相对路径）"
+    )
+    parser.add_argument("--verbose", "-v", action="store_true", help="显示详细输出")
+    parser.add_argument("--overwrite", action="store_true", help="覆盖已存在的文件")
+
+    # 直接传递图片路径
+    parser.add_argument("image_paths", nargs="*", help="直接指定图片相对路径")
+
+    args = parser.parse_args()
+
+    # 创建示例配置文件
+    if args.create_config:
+        create_example_config(args.create_config)
+        return
+
+    # 检查必要参数（如果不是创建配置文件）
+    if not args.config and not (
+        args.source_dir and args.target_host and args.target_user and args.target_dir
+    ):
+        print("错误: 必须提供配置文件或指定必要的服务器参数")
+        print("必要参数: --source-dir, --target-host, --target-user, --target-dir")
+        parser.print_help()
+        sys.exit(1)
+
+    # 加载配置
+    config = {}
+    if args.config:
+        config = load_config(args.config)
+        if not config:
+            print("错误: 无法加载配置文件")
+            sys.exit(1)
+
+    # 构建源服务器配置
+    source_config_data = config.get("source", {})
+
+    # 命令行参数覆盖配置文件
+    if args.source_type:
+        source_config_data["type"] = args.source_type
+
+    # 如果是SSH类型，需要SSH配置
+    if args.source_type == "ssh" or source_config_data.get("type") == "ssh":
+        ssh_config = source_config_data.get("ssh", {})
+
+        if args.source_host:
+            ssh_config["host"] = args.source_host
+        if args.source_port:
+            ssh_config["port"] = args.source_port
+        if args.source_user:
+            ssh_config["username"] = args.source_user
+        if args.source_password:
+            ssh_config["password"] = args.source_password
+        if args.source_key:
+            ssh_config["key_file"] = args.source_key
+
+        source_config_data["ssh"] = ssh_config
+
+        # 检查必要的SSH参数
+        if not ssh_config.get("host"):
+            print("错误: 源服务器为SSH类型，必须指定 --source-host")
+            sys.exit(1)
+        if not ssh_config.get("username"):
+            print("错误: 源服务器为SSH类型，必须指定 --source-user")
+            sys.exit(1)
+    else:
+        # 本地类型，不需要SSH配置
+        source_config_data.pop("ssh", None)
+
+    # 设置基础目录
+    if args.source_dir:
+        source_config_data["base_dir"] = args.source_dir
+    elif "base_dir" not in source_config_data:
+        print("错误: 必须指定源服务器图片基础目录 (--source-dir)")
+        sys.exit(1)
+
+    # 构建目标服务器配置
+    target_config_data = config.get("target", {})
+    target_config_data["type"] = "ssh"  # 目标必须是SSH
+
+    ssh_config = target_config_data.get("ssh", {})
+
+    # 命令行参数覆盖配置文件
+    if args.target_host:
+        ssh_config["host"] = args.target_host
+    if args.target_port:
+        ssh_config["port"] = args.target_port
+    if args.target_user:
+        ssh_config["username"] = args.target_user
+    if args.target_password:
+        ssh_config["password"] = args.target_password
+    if args.target_key:
+        ssh_config["key_file"] = args.target_key
+
+    target_config_data["ssh"] = ssh_config
+
+    # 设置基础目录
+    if args.target_dir:
+        target_config_data["base_dir"] = args.target_dir
+    elif "base_dir" not in target_config_data:
+        print("错误: 必须指定目标服务器图片基础目录 (--target-dir)")
+        sys.exit(1)
+
+    # 检查必要的目标SSH参数
+    if not ssh_config.get("host"):
+        print("错误: 必须指定目标服务器地址 (--target-host)")
+        sys.exit(1)
+    if not ssh_config.get("username"):
+        print("错误: 必须指定目标服务器用户名 (--target-user)")
+        sys.exit(1)
+
+    # 创建服务器配置对象
+    try:
+        source_config = ServerConfig.from_dict(source_config_data)
+        target_config = ServerConfig.from_dict(target_config_data)
+
+        if args.verbose:
+            print(f"源服务器配置: {source_config}")
+            print(f"目标服务器配置: {target_config}")
+    except Exception as e:
+        print(f"创建服务器配置失败: {e}")
+        sys.exit(1)
+
+    # 获取图片路径列表
+    image_paths = []
+
+    # 从文件读取
+    if args.input:
+        file_paths = read_image_paths(args.input)
+        image_paths.extend(file_paths)
+
+    # 从命令行参数添加
+    if args.image_paths:
+        image_paths.extend(args.image_paths)
+
+    if not image_paths:
+        print("错误: 没有指定要迁移的图片路径")
+        print("请使用 --input 指定文件或直接在命令行提供路径")
+        sys.exit(1)
+
+    # 去重
+    image_paths = list(set(image_paths))
+
+    # 创建并运行迁移器
+    migrator = ImageMigrator(
+        source_config=source_config,
+        target_config=target_config,
+        verbose=args.verbose,
+        overwrite=args.overwrite,
+    )
+
+    # 执行迁移
+    stats = migrator.migrate_images(image_paths)
+
+    # 如果有失败，返回错误代码
+    if stats["failed"] > 0:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()