bitool 0.1.2__py3-none-any.whl

bitool/scripts/sshcopyid.py

@@ -0,0 +1,491 @@
+"""类似 ssh-copy-id 的 SSH 密钥部署工具."""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import logging
+import re
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
+
+from bitool.utils import RunResult, execute
+
+# 主机名验证正则表达式 (支持域名、IPv4、localhost)
+HOSTNAME_PATTERN = re.compile(
+    r"^localhost$|"  # localhost
+    r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$|"  # 域名
+    r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$"  # IPv4
+)
+
+# 用户名验证正则表达式
+USERNAME_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_-]*$")
+
+
+def validate_hostname(hostname: str) -> bool:
+    """验证主机名格式是否安全.
+
+    Args:
+        hostname: 待验证的主机名
+
+    Returns:
+        True 如果格式有效,否则 False
+    """
+    return bool(HOSTNAME_PATTERN.match(hostname))
+
+
+def validate_username(username: str) -> bool:
+    """验证用户名格式是否安全.
+
+    Args:
+        username: 待验证的用户名
+
+    Returns:
+        True 如果格式有效,否则 False
+    """
+    return bool(USERNAME_PATTERN.match(username)) and len(username) <= 64
+
+
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+logger = logging.getLogger(__name__)
+
+
+MIN_KEY_PARTS = 2
+MIN_PORT = 1
+MAX_PORT = 65535
+
+
+@dataclass
+class SSHCopyIDConfig:
+    """SSH 密钥部署配置."""
+
+    hostname: str
+    username: str
+    password: str
+    port: int = 22
+    public_key_path: str = "~/.ssh/id_rsa.pub"
+    timeout: int = 30
+    connect_timeout: int = 10
+
+    @cached_property
+    def expanded_key_path(self) -> Path:
+        """获取展开后的公钥文件路径."""
+        return Path(self.public_key_path).expanduser()
+
+    @cached_property
+    def ssh_command_parts(self) -> list[str]:
+        """生成基础 SSH 命令部分."""
+        return [
+            "ssh",
+            "-p",
+            str(self.port),
+            "-o",
+            "StrictHostKeyChecking=no",
+            "-o",
+            "UserKnownHostsFile=/dev/null",
+            "-o",
+            f"ConnectTimeout={self.connect_timeout}",
+            f"{self.username}@{self.hostname}",
+        ]
+
+
+class SSHAuthenticationError(Exception):
+    """SSH 认证失败时抛出."""
+
+
+class SSHConnectionError(Exception):
+    """SSH 连接失败时抛出."""
+
+
+class SSHKeyNotFoundError(Exception):
+    """SSH 公钥文件未找到时抛出."""
+
+
+def deploy_with_native_ssh(config: SSHCopyIDConfig, pub_key: str) -> None:
+    """使用原生 SSH 部署 SSH 密钥,无需 sshpass."""
+    import shlex
+
+    # 为 SSH 密钥部署创建临时脚本
+    # 提取密钥类型和密钥数据以实现更可靠的匹配
+    key_parts = pub_key.split()
+    if len(key_parts) >= MIN_KEY_PARTS:
+        key_type = key_parts[0]
+        key_data = key_parts[1]
+        grep_pattern = f"{key_type}.*{key_data}"
+    else:
+        # 用于非标准密钥格式的后备方案
+        grep_pattern = pub_key
+
+    # 安全转义 shell 变量以防止命令注入
+    safe_grep_pattern = shlex.quote(grep_pattern)
+    safe_pub_key = shlex.quote(pub_key)
+
+    script_content = f"""#!/bin/bash
+mkdir -p ~/.ssh
+chmod 700 ~/.ssh
+cd ~/.ssh
+touch authorized_keys
+chmod 600 authorized_keys
+grep -qF {safe_grep_pattern} authorized_keys 2>/dev/null || \\
+    echo {safe_pub_key} >> authorized_keys
+"""
+
+    with tempfile.NamedTemporaryFile(
+        encoding="utf-8", mode="w", suffix=".sh", delete=False
+    ) as script_file:
+        script_file.write(script_content)
+        script_path = script_file.name
+
+    try:
+        # 使用 SSH 执行部署脚本,通过管道传递给 bash
+        command = [
+            "ssh",
+            *config.ssh_command_parts[:-1],  # 除主机名外的所有部分
+            f"{config.username}@{config.hostname}",
+            "bash -s",
+        ]
+
+        logger.info("尝试使用原生 SSH 进行部署...")
+        # 注意:此处已通过 validate_hostname() 和 validate_username() 验证输入
+        # 使用列表形式且 shell=False 是安全的
+        process = execute(
+            command, check=False, capture_output=True, text=True, timeout=config.timeout
+        )
+
+        if process.returncode != 0:
+            if "Permission denied" in process.stderr:
+                msg = "原生 SSH 认证失败"
+                raise SSHAuthenticationError(msg)
+            msg = "原生 SSH 部署失败:%s", process.stderr
+            raise SSHConnectionError(msg)
+
+        logger.info("使用原生 SSH 方法成功部署 SSH 密钥")
+
+    finally:
+        # 清理临时文件
+        with contextlib.suppress(BaseException):
+            Path(script_path).unlink()
+
+
+def _read_public_key(key_path: Path) -> str:
+    """读取并验证公钥文件.
+
+    Returns
+    -------
+        公钥内容
+
+    Raises
+    ------
+        SSHKeyNotFoundError: 密钥文件未找到或为空
+        SSHConnectionError: 读取密钥文件失败
+    """
+    # 验证密钥文件是否存在
+    if not key_path.exists():
+        msg = "公钥文件未找到:%s", key_path
+        raise SSHKeyNotFoundError(msg)
+
+    try:
+        with key_path.open(encoding="utf-8") as f:
+            pub_key = f.read().strip()
+
+    except (OSError, UnicodeDecodeError) as e:
+        msg = "读取公钥文件失败:%s"
+        raise SSHConnectionError(msg) from e
+
+    # 验证密钥内容
+    if not pub_key:
+        msg = "公钥文件为空:%s", key_path
+        raise SSHKeyNotFoundError(msg)
+
+    return pub_key
+
+
+def _extract_key_pattern(pub_key: str) -> str:
+    """从公钥中提取 grep 模式用于重复检测.
+
+    Args:
+        pub_key: SSH 公钥内容
+
+    Returns
+    -------
+        用于 grep 匹配的模式 (key_type.*key_data)
+    """
+    key_parts = pub_key.split()
+    if len(key_parts) >= MIN_KEY_PARTS:
+        key_type = key_parts[0]
+        key_data = key_parts[1]
+        return f"{key_type}.*{key_data}"
+    # 用于非标准密钥格式的后备方案
+    return pub_key
+
+
+def _build_ssh_command(
+    config: SSHCopyIDConfig, grep_pattern: str, pub_key: str
+) -> list[str]:
+    """使用 sshpass 构建 SSH 命令用于密钥部署.
+
+    Args:
+        config: SSH 配置
+        grep_pattern: 用于 grep 匹配以避免重复密钥的模式
+        pub_key: 公钥内容
+
+    Returns
+    -------
+        用于 subprocess 的命令列表 (sshpass + ssh + 远程命令)
+    """
+    import shlex
+
+    # 安全转义 shell 参数以防止命令注入
+    safe_grep_pattern = shlex.quote(grep_pattern)
+    safe_pub_key = shlex.quote(pub_key)
+
+    return [
+        "sshpass",
+        "-p",
+        config.password,
+        *config.ssh_command_parts,
+        (
+            f"mkdir -p ~/.ssh && chmod 700 ~/.ssh && "
+            f"cd ~/.ssh && touch authorized_keys && "
+            f"chmod 600 authorized_keys && "
+            f"grep -qF {safe_grep_pattern} "
+            f"authorized_keys 2>/dev/null || "
+            f"echo {safe_pub_key} >> authorized_keys"
+        ),
+    ]
+
+
+def _handle_ssh_error(process: RunResult, config: SSHCopyIDConfig) -> None:
+    """处理 SSH 错误并抛出适当的异常.
+
+    Args:
+        process: 已完成的具有返回码和标准错误的 RunResult
+        config: SSH 配置用于错误上下文
+
+    Raises
+    ------
+        SSHAuthenticationError: 认证失败 (Permission denied)
+        SSHConnectionError: 连接失败 (refused, network unreachable 等)
+    """
+    if "Permission denied" in process.stderr:
+        msg = "认证失败,请检查用户名或密码"
+        raise SSHAuthenticationError(msg)
+    if "Connection refused" in process.stderr:
+        msg = "连接被拒绝:%s:%d", config.hostname, config.port
+        raise SSHConnectionError(msg)
+    if "Network is unreachable" in process.stderr:
+        msg = "网络不可达:%s", config.hostname
+        raise SSHConnectionError(msg)
+    msg = "SSH 执行失败:%s", process.stderr.strip()
+    raise SSHConnectionError(msg)
+
+
+def _try_alternative_methods(config: SSHCopyIDConfig, pub_key: str) -> None:
+    """当 sshpass 不可用时尝试替代部署方法.
+
+    尝试使用原生 SSH（无需密码认证）作为后备。
+    如果所有方法都失败,提供 sshpass 的安装说明。
+
+    Args:
+        config: SSH 配置
+        pub_key: 公钥内容
+
+    Raises
+    ------
+        SystemExit: 如果所有方法都失败,以代码 1 退出
+    """
+    logger.warning("未找到 sshpass 工具。尝试替代部署方法...")
+
+    try:
+        deploy_with_native_ssh(config, pub_key)
+    except (OSError, subprocess.SubprocessError):
+        logger.exception(
+            "原生 SSH 部署也失败了。请安装 sshpass 或使用手动方法:\n"
+            "安装方法:\n"
+            "Windows: 下载 Windows 版 sshpass 或使用 WSL\n"
+            "Ubuntu/Debian: sudo apt-get install sshpass\n"
+            "CentOS/RHEL: sudo yum install sshpass\n"
+            "macOS: brew install hudochenkov/sshpass/sshpass\n"
+            f"或手动方式:ssh-copy-id -p {config.port} {config.username}@{config.hostname}"
+        )
+        sys.exit(1)
+    else:
+        logger.info("使用原生 SSH 方法成功部署 SSH 密钥")
+
+
+def ssh_copy_id(config: SSHCopyIDConfig) -> None:
+    """将 SSH 公钥部署到远程服务器.
+
+    Args:
+        config: SSH copy ID 配置
+
+    Raises
+    ------
+        SSHAuthenticationError: 认证失败
+        SSHConnectionError: 连接失败
+        SSHKeyNotFoundError: 公钥文件未找到
+        Exception: 其他异常
+        ValueError: 端口号无效或输入验证失败
+    """
+    # 读取本地公钥内容
+    pub_key = _read_public_key(config.expanded_key_path)
+
+    # 安全警告:StrictHostKeyChecking=no 禁用主机密钥验证
+    # 仅在受信任的网络环境中使用此选项
+    logger.warning("安全警告:禁用 SSH 主机密钥验证。确保您处于受信任的网络环境中.")
+
+    # 验证配置
+    if not config.hostname or not config.username or not config.password:
+        msg = "主机名、用户名和密码不能为空"
+        raise ValueError(msg)
+
+    # 验证主机名和用户名格式以防止命令注入
+    if not validate_hostname(config.hostname):
+        msg = "主机名格式无效"
+        raise ValueError(msg)
+
+    if not validate_username(config.username):
+        msg = "用户名格式无效"
+        raise ValueError(msg)
+
+    if config.port < MIN_PORT or config.port > MAX_PORT:
+        msg = "无效的端口号:%d, 必须在 1-65535 范围内", config.port
+        raise ValueError(msg)
+
+    # 提取密钥模式用于可靠匹配
+    grep_pattern = _extract_key_pattern(pub_key)
+
+    try:
+        # 使用 sshpass 执行远程命令
+        command = _build_ssh_command(config, grep_pattern, pub_key)
+
+        # 注意:此处已通过 validate_hostname() 和 validate_username() 验证输入
+        # 使用列表形式且 shell=False 是安全的
+        process = execute(
+            command, check=False, capture_output=True, text=True, timeout=config.timeout
+        )
+
+        if process.returncode != 0:
+            _handle_ssh_error(process, config)
+
+    except FileNotFoundError:
+        # 如果找不到 sshpass，尝试替代方法
+        _try_alternative_methods(config, pub_key)
+    except (subprocess.TimeoutExpired, TimeoutError) as e:
+        msg = f"SSH 连接超时, 已等待 {config.timeout} 秒"
+        raise SSHConnectionError(msg) from e
+    except (OSError, subprocess.SubprocessError) as e:
+        msg = f"SSH 操作失败:{e}"
+        raise SSHConnectionError(msg) from e
+
+
+def parse_args() -> argparse.Namespace:
+    """解析命令行参数.
+
+    Returns
+    -------
+        已解析的参数
+    """
+    parser = argparse.ArgumentParser(
+        description="将 SSH 公钥部署到远程服务器",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+示例:
+  sshcopyid 192.168.1.100 user password
+  sshcopyid 192.168.1.100 user password -p 2222
+  sshcopyid 192.168.1.100 user password -k ~/.ssh/id_ed25519.pub
+        """,
+    )
+    parser.add_argument("hostname", help="远程服务器主机名或 IP 地址")
+    parser.add_argument("username", help="远程服务器用户名")
+    parser.add_argument("password", help="远程服务器密码")
+    parser.add_argument("-p", "--port", type=int, default=22, help="SSH 端口 (默认:22)")
+    parser.add_argument(
+        "-k",
+        "--keypath",
+        type=str,
+        default="~/.ssh/id_rsa.pub",
+        help="公钥文件路径 (默认:~/.ssh/id_rsa.pub)",
+    )
+    parser.add_argument(
+        "-t", "--timeout", type=int, default=30, help="SSH 操作超时秒数 (默认:30)"
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", help="启用详细日志")
+    return parser.parse_args()
+
+
+def setup_logging(*, verbose: bool = False) -> None:
+    """设置日志配置.
+
+    Args:
+        verbose: 启用详细日志
+    """
+    level = logging.DEBUG if verbose else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+
+def main() -> None:
+    """运行命令行界面入口点."""
+    args = parse_args()
+    setup_logging(verbose=args.verbose)
+
+    # 创建配置对象
+    config = SSHCopyIDConfig(
+        hostname=args.hostname,
+        username=args.username,
+        password=args.password,
+        port=args.port,
+        public_key_path=args.keypath,
+        timeout=args.timeout,
+    )
+
+    # 在执行前验证配置
+    if not config.hostname or not config.username or not config.password:
+        logger.error("主机名、用户名和密码不能为空")
+        sys.exit(1)
+
+    if config.port < MIN_PORT or config.port > MAX_PORT:
+        logger.error("无效的端口号:%d, 必须在 1-65535 范围内", config.port)
+        sys.exit(1)
+
+    # 验证公钥文件
+    if not config.expanded_key_path.exists():
+        logger.error("公钥文件不存在:%s", config.expanded_key_path)
+        sys.exit(1)
+
+    if not config.expanded_key_path.is_file():
+        logger.error("指定路径不是文件:%s", config.expanded_key_path)
+        sys.exit(1)
+
+    try:
+        logger.info(
+            "正在部署 SSH 密钥到 %s@%s:%d",
+            config.username,
+            config.hostname,
+            config.port,
+        )
+        logger.info("使用公钥:%s", config.expanded_key_path)
+
+        ssh_copy_id(config)
+        logger.info("SSH 密钥部署成功完成")
+
+    except SSHAuthenticationError:
+        logger.exception("认证失败")
+        sys.exit(2)
+    except SSHConnectionError:
+        logger.exception("连接失败")
+        sys.exit(3)
+    except SSHKeyNotFoundError:
+        logger.exception("密钥文件错误")
+        sys.exit(4)
+    except (OSError, subprocess.SubprocessError):
+        logger.exception("意外错误")
+        sys.exit(1)