拿来即用的 Python SSH+SFTP 实现类

Posted 2023-03-08 何小有

一个拿来即用的 Python SSH 和 SFTP 实现类，可用于：

• 与 Linux 服务器建立持续交互的 SSH 会话
• 从 Linux 服务器下载远程文件
• 上传本地文件到 Linux 服务器

新创建一个 linux_client.py 文件，编写下面的 LinuxClient 类代码：

import re
import time
# 通过 pip install paramiko 命令安装 paramiko 库
from paramiko import SSHClient, RSAKey, AutoAddPolicy, ssh_exception

class LinuxClient():
    def __init__(
        self,
        hostname: str,
        port: int = 22,
        private_key_file: str = None,
        username: str = None,
        password: str = None,
        tcp_timeout: int = 15,
    ) -> None:
        """初始化 LinuxClient 实例对象

        按以下优先级选择身份认证的方式:
            - 使用 `username` 和 `private_key_file` 认证
            - 使用 `username` 和 `password` 认证

        Args:
            hostname (str): 连接的服务器
            port (int, optional): 连接的服务器端口. 默认值为 22.
            private_key_file (str, optional): 认证私钥文件路径. 默认值为 None.
            username (str, optional): 认证用户名. 默认值为 None.
            password (str, optional): 认证密码. 默认值为 None.
            tcp_timeout (int, optional): TCP连接的超时秒数. 默认值为 15.
        """
        self.hostname = hostname
        self.port = port
        self.__ssh_transport = SSHClient()  # paramiko.client.SSHClient
        # 首次连接时自动信任远程机器, 建立 SSH 互信通道
        self.__ssh_transport.set_missing_host_key_policy(AutoAddPolicy())
        self.__ssh_session = None  # paramiko.channel.Channel
        self.__sftp = None  # paramiko.sftp_client.SFTPClient
        self._init_results = (True, '初始化成功')
        if not username:
            self._init_results = (False, '认证用户名不能为空')
            return
        try:
            self.__create_ssh_connect(username, private_key_file, password, tcp_timeout)
        except Exception as e:
            self._init_results = (False, str(e))
        if self._init_results[0]:
            # 建立持续交互的 SSH 会话
            self.__ssh_session = self.__ssh_transport.get_transport().open_session()
            self.__ssh_session.get_pty()
            self.__ssh_session.invoke_shell()
            # 使用 SSH 连接创建 SFTP 连接
            self.__sftp = self.__ssh_transport.open_sftp()

    def __create_ssh_connect(self, username, private_key_file, password, tcp_timeout):
        """建立 SSH 连接并建立会话"""
        try:
            if private_key_file:  # 私钥认证
                self.__ssh_transport.connect(
                    self.hostname,
                    port=self.port,
                    username=username,
                    pkey=RSAKey.from_private_key_file(private_key_file),  # 导入私钥
                    banner_timeout=60,  # 等待 SSHBanner 出现的超时秒数
                    timeout=tcp_timeout,
                )
            elif password:  # 用户名密码认证
                self.__ssh_transport.connect(
                    self.hostname,
                    port=self.port,
                    username=username,
                    password=password,
                    banner_timeout=60,
                    timeout=tcp_timeout,
                )
            else:
                self._init_results = (False, '认证用的参数不完整')
        except FileNotFoundError:
            self._init_results = (False, '认证私钥文件不存在')
        except TimeoutError:
            self._init_results = (False, '服务器网络连接超时')
        except ssh_exception.AuthenticationException:
            self._init_results = (False, '服务器身份认证失败')

    def _shell_cache(self):
        """输出 Shell 缓存的格式化文本"""
        result = ''
        while True:
            time.sleep(0.5)
            res = self.__ssh_session.recv(65535).decode('utf8')
            result += res
            if res.endswith('# ') or res.endswith('$ '):
                break
        result = re.sub('\\x1b.*?m', '', result)  # 移除 `\\xblah[0m` 等无效内容
        return result.strip('\\n')  # 移除换行

    def run(self, command: bytes) -> str:
        """执行 Shell 命令, 持续交互

        Args:
            command (bytes): 要发送的命令数据

        Returns:
            str: 格式化后的 Shell 缓存
        """
        channel = self.__ssh_session  # 获取通道
        channel.send(f'command\\n')  # 执行命令
        return self._shell_cache()

    @staticmethod
    def upload_callback(current_bytes: int, total_bytes: int):
        """上传回调函数 (如果需要回调，可以在继承时重写这个函数)

        Args:
            current_bytes (int): 目前传输的字节数
            total_bytes (int): 要传输的总字节数
        """
        # print(f'current_bytes/total_bytes')
        pass

    def upload(self, local_path: str, target_path: str):
        """将本地文件从本地主机上传到远程服务器

        Args:
            local_path (str): 要上传的本地文件路径
            target_path (str): 远程服务器上的存储路径, 包含文件名

        Returns:
            tuple: (result: bool, describe: str)
        """
        try:
            self.__sftp.put(
                localpath=local_path,
                remotepath=target_path,
                callback=LinuxClient.upload_callback,
                confirm=True,
            )
            # 增加权限让其他用户可读可执行, 用 0o 作为前缀表示八进制
            self.__sftp.chmod(target_path, 0o755)
        except FileNotFoundError:
            return (False, '本地或远程路径无效')
        else:
            return (True, '本地文件上传成功')

    @staticmethod
    def download_callback(current_bytes: int, total_bytes: int):
        """下载回调函数 (如果需要回调，可以在继承时重写这个函数)

        Args:
            current_bytes (int): 目前传输的字节数
            total_bytes (int): 要传输的总字节数
        """
        # print(f'current_bytes/total_bytes')
        pass

    def download(self, target_path: str, local_path: str):
        """将远程文件从远程服务器下载到本地

        Args:
            target_path (str): 要下载的远程文件路径
            local_path (str): 本地主机上的存储路径

        Returns:
            tuple: (result: bool, describe: str)
        """
        try:
            self.__sftp.get(
                remotepath=target_path,
                localpath=local_path,
                callback=LinuxClient.download_callback,
                prefetch=True,
            )
        except FileNotFoundError:
            return (False, '远程或本地路径无效')
        else:
            return (True, '远程文件下载成功')

    def __del__(self):
        """释放 LinuxClient 实例对象"""
        if self._init_results[0]:
            self.__sftp.close()
            self.__ssh_session.close()
        self.__ssh_transport.close()

然后是调用上面 LinuxClient 类的示例代码：

from linux_client import LinuxClient

lc = LinuxClient(
    hostname='123.45.67.89',
    port=34567,
    username='ecs-user',
    private_key_file='demo/pic-res-server.pem',
)

# 服务器连接结果
print(lc._init_results)
# 连接成功后, 先获取 SSH Banner
print(lc._shell_cache())

# 执行 pwd 命令并获取输出
print(lc.run('pwd'))
# 执行 ls 命令并获取输出
print(lc.run('ls'))
# 执行 cd db_back/test_database 命令并获取输出
print(lc.run('cd db_back/test_database'))

# 上传文件
upload_result, upload_describe = lc.upload(
    local_path='demo/test_upload.txt',
    target_path='/home/ecs-user/db_back/test_database/test_upload.txt',
)
print(upload_result, upload_describe)

# 下载文件
download_result, download_describe = lc.download(
    target_path='/home/ecs-user/db_back/test_database/test_upload.txt',
    local_path='demo/test_upload.txt',
)
print(download_result, download_describe)

最后看下示例代码的控制台打印效果：

$ (True, '初始化成功')
$ Welcome to Ubuntu 20.04.5 LTS (GNU/Linux 5.4.0-125-generic x86_64)
$
$  * Documentation:  https://help.ubuntu.com
$  * Management:     https://landscape.canonical.com
$  * Support:        https://ubuntu.com/advantage
$ New release '22.04.2 LTS' available.
$ Run 'do-release-upgrade' to upgrade to it.
$
$
$ Welcome to Alibaba Cloud Elastic Compute Service !
$
$ Last login: Wed Mar  1 16:46:42 2023 from 123.45.67.89
$ ecs-user@pic-res-server:~$
$ pwd
$ /home/ecs-user
$ ecs-user@pic-res-server:~$
$ ls
$ db_back  wf-api
$ ecs-user@pic-res-server:~$
$ cd db_back/test_database
$ ecs-user@pic-res-server:~/db_back/test_database$
$ True 本地文件上传成功
$ True 远程文件下载成功