使用方法

  1. (py-ssh) ➜  py-ssh python tmp.py -h
  2. usage: tmp.py [-h] [--user USER] [--password PASSWORD] [--port PORT] [--addresses IPS [IPS ...]] [--cmd CMD] [--sem SEM] [--down_host DOWN_HOST] [--files FILES [FILES ...]] [--dst_path DST_PATH]
  4. optional arguments:
  5.   -h, --help            show this help message and exit
  6.   --user USER, -u USER  Specify ssh username
  7.   --password PASSWORD, -p PASSWORD
  8.                         Specify ssh password
  9.   --port PORT, -P PORT  Specify ssh port;default 22
  10.   --addresses IPS [IPS ...], -a IPS [IPS ...]
  11.                         Specify ssh hosts
  12.   --cmd CMD, -c CMD     Specify cmd for execute
  13.   --sem SEM, -s SEM     Number of threads; default 5
  14.   --down_host DOWN_HOST
  15.                         Download file server address
  17.   --files FILES [FILES ...], -f FILES [FILES ...]
  18.                         Specify files use for copy
  19.   --dst_path DST_PATH, -d DST_PATH
  20.                         File storage directory
  21.   --dir DIR             Specify storage files directory

参数解释

  • user ssh登陆用户
  • password ssh登陆密码
  • port ssh登陆端口号
  • addresses 需要登陆的主机清单，支持单主机和多主机， —addresses 192.168.50.163 192.168.50.160:192.168.50.163
  • cmd 执行的命令
  • sem 并发的信号量
  • down_host 用于下载文件的地址, 示例: 192.168.1.10:8000
  • files 需要进行拷贝到远程主机的文件, 必须搭配 dst_path 使用
  • dst_path 拷贝到远程主机的目录
  • dir 用于存放需要被拷贝的文件，先拷贝 files 然后 dir 里面的文件，统一放到 dstz_path下面

用法示例

  1. python tmp.py --user cortex \
  2.                             --addresses 192.168.50.163 192.168.50.160:192.168.51.163 \
  3.               --password cortex \
  4.               --files /Users/wuxing/install.sh /Users/wuxing/minicom.log \
  5.               --dst_path /tmp \
  6.               --down_host 192.168.50.101:8000 \
  7.               --cmd 'ls -l /tmp' \
  8.               --dir files \
  9.               --sem 10

脚本

  1. import getpass
  2. import hashlib
  3. import os
  4. from concurrent.futures import ThreadPoolExecutor
  5. from time import perf_counter
  6. import sys
  7. import threading
  8. import paramiko
  9. from time import sleep
  10. from loguru import logger as log
  11. import argparse
  12. from ipaddress import ip_address
  14. log.remove()
  15. log.add(sys.stderr, level="INFO")
  16. log.add("run.log", level="INFO", mode="w")
  17. log.add("error.log", level="ERROR", mode="w")
  20. class SSHConnection(threading.Thread):
  22.     def __init__(self, **kwargs):
  23.         super(SSHConnection, self).__init__()
  24.         self.username = kwargs.get('user')
  25.         self.host = kwargs.get('host')
  26.         self.port = kwargs.get('port')
  27.         self.password = kwargs.get('password')
  28.         self.keyfile = self.get_keyfile()
  29.         self.sem = kwargs.get('sem')
  30.         self.files = kwargs.get('files')
  31.         self.down_host = kwargs.get('down_host')
  32.         self.dst_path = kwargs.get('dst_path')
  33.         self.cmd = kwargs.get('cmd')
  34.         self.dir = kwargs.get('dir')
  35.         log.info(f"directory: {self.dir}")
  36.         if self.dir:
  37.             if not os.path.exists(self.dir):
  38.                 log.error(f"{self.dir} is not exists")
  39.                 sys.exit(0)
  41.     def get_keyfile(self, path=os.getcwd()):
  43.         default_keyfile = os.path.join(
  44.             os.environ.get('HOME', 'D:/'), '.ssh', 'id_rsa')
  46.         if 'id_rsa' in os.listdir(path):
  47.             keyfile = os.path.join(path, 'id_rsa')
  48.         elif os.path.isfile(default_keyfile):
  49.             keyfile = default_keyfile
  50.         else:
  51.             keyfile = ''
  53.         return keyfile
  55.     def connect(self):
  56.         try:
  57.             transport = paramiko.Transport((self.host, self.port))
  58.         except Exception:
  59.             log.error(f"{self.host} is down!")
  60.             exit(0)
  62.         if self.password:
  63.             transport.connect(username=self.username, password=self.password)
  64.         elif self.keyfile:
  65.             transport.connect(
  66.                 username=self.username,
  67.                 pkey=paramiko.RSAKey.from_private_key_file(self.keyfile))
  68.         else:
  69.             password = getpass.getpass("Password for %s@%s: " % (self.username, self.host))
  70.             transport.connect(username=self.username, password=password)
  72.         self._transport = transport
  74.         log.info("Connected to %s as %s" % (self.host, self.username))
  76.     def close(self):
  77.         try:
  78.             self._transport.close()
  79.         except AttributeError:
  80.             pass
  82.     def run_cmd(self, command):
  83.         ssh = paramiko.SSHClient()
  84.         ssh._transport = self._transport
  86.         stdin, stdout, stderr = ssh.exec_command(command)
  87.         res = stdout.read().decode()
  88.         error = stderr.read().decode()
  90.         if error.strip():
  91.             # return error
  92.             raise Exception(error)
  93.         else:
  94.             return res
  96.     def trans_file(self, localpath, remotepath, method=''):
  97.         sftp = paramiko.SFTPClient.from_transport(self._transport)
  98.         if method == 'put':
  99.             sftp.put(localpath, remotepath)
  100.             print("File %s has uploaded to %s" % (localpath, remotepath))
  101.         elif method == 'get':
  102.             sftp.get(remotepath, localpath)
  103.             print("File %s has saved as %s" % (remotepath, localpath))
  104.         else:
  105.             print('usage: trans_file(localpath, remotepath, method="get/put"')
  107.     def generter_md5(self, filepath):
  108.         with open(filepath, 'rb') as f:
  109.             data = f.read()
  111.         return hashlib.md5(data).hexdigest()
  113.     def load_directory(self):
  114.         files = []
  115.         for root, dirs, filenames in os.walk(self.dir):
  116.             for f in filenames:
  117.                 filepath = os.path.join(root, f)
  118.                 files.append(filepath)
  120.         return files
  122.     def donwload_files(self, filepath, down_host):
  123.         key = os.path.basename(filepath)
  124.         download_script = """
  125.                             if [ -f {path}/{name} ];then
  126.                                 echo "{path}/{name}  exists!"
  127.                                 rm -rf {path}/{name}
  128.                             else
  129.                                 echo "{path}/{name} not exists!" 
  130.                             fi
  131.                             wget -q http://{host}/{name} -O {path}/{name}
  132.                             """.format(path=self.dst_path, name=key, host=down_host)
  133.         self.run_cmd(download_script)
  134.         dst_md5 = self.run_cmd(
  135.             "md5sum {path}/{name} | cut -d ' ' -f 1".format(path=self.dst_path, name=key)).strip()
  136.         src_md5 = self.generter_md5(filepath)
  137.         log.debug(f"remote ssh file md5:{dst_md5}, local file md5:{src_md5}")
  138.         if dst_md5 != src_md5:
  139.             log.error(f"md5校验失败, {self.dst_path}/{key} md5: {dst_md5}, local md5: {src_md5}")
  140.             sys.exit(0)
  141.         log.info(f"{self.host}, {self.dst_path}/{key} 下载完成，md5校验成功!")
  143.     @staticmethod
  144.     def find_ips(start, end):
  145.         start = ip_address(start)
  146.         end = ip_address(end) + 1
  147.         result = []
  148.         while start < end:
  149.             result.append(str(start))
  150.             start += 1
  151.         return result
  153.     def run(self) -> None:
  154.         try:
  155.             self.sem.acquire()
  156.             self.connect()
  157.             log.debug(f"files: {self.files}")
  158.             if self.files:
  159.                 for filepath in self.files:
  160.                     self.donwload_files(filepath, self.down_host)
  162.             log.info(f"directory: {self.dir}")
  163.             if self.dir:
  164.                 log.info(f"run dir: {self.dir}")
  165.                 self.files = self.load_directory()
  166.                 for filepath in self.files:
  167.                     self.donwload_files(filepath, f"{self.down_host}/{self.dir}")
  169.             if self.cmd:
  170.                 log.info(self.run_cmd(self.cmd))
  171.         finally:
  172.             self.sem.release()
  174.     def __del__(self):
  175.         self.close()
  178. def run(**kwargs):
  179.     begin = perf_counter()
  180.     ips = kwargs.get('ips', None)
  181.     if not ips:
  182.         log.error("Must specify hosts, use '-a' or '--addresses'")
  183.         sys.exit(0)
  185.     user, password = kwargs.get('user', None), kwargs.get('password', None)
  186.     threads = []
  187.     for ip in ips:
  188.         log.debug(f"ip:{ip}")
  189.         if ':' in ip:
  190.             start, end = ip.split(':')
  191.             for i in SSHConnection.find_ips(start, end):
  192.                 kwargs['host'] = i
  193.                 threads.append(SSHConnection(**kwargs))
  194.         else:
  195.             kwargs['host'] = ip
  196.             threads.append(SSHConnection(**kwargs))
  198.     for t in threads:
  199.         t.start()
  201.     for t in threads:
  202.         t.join()
  204.     end = perf_counter()
  206.     print("Cost time:", end - begin)
  209. if __name__ == '__main__':
  210.     # run()
  211.     parser = argparse.ArgumentParser()
  212.     parser.add_argument('--user', '-u', dest="user", action='store', help='Specify ssh username')
  213.     parser.add_argument('--password', '-p', dest="password", action="store", help='Specify ssh password')
  214.     parser.add_argument('--port', '-P', dest="port", action='store', default=22, type=int,
  215.                         help='Specify ssh port;default 22')
  216.     parser.add_argument('--addresses', '-a', dest="ips", action='store', nargs='+', type=str, help='Specify ssh hosts')
  217.     parser.add_argument('--cmd', '-c', dest="cmd", action='store', type=str, help="Specify cmd for execute")
  218.     parser.add_argument('--sem', '-s', dest="sem", action='store', type=int, help="Number of threads; default 5",
  219.                         default=5)
  220.     parser.add_argument('--down_host', dest='down_host', action='store', help="Download file server address")
  222.     copy_group = parser.add_argument_group()
  223.     copy_group.add_argument('--files', '-f', dest="files", action='store', nargs='+', type=str,
  224.                             help="Specify files use for copy")
  225.     copy_group.add_argument('--dst_path', '-d', dest="dst_path", action="store", help="File storage directory")
  226.     copy_group.add_argument('--dir', dest="dir", action='store', help='Specify storage files directory')
  228.     args = parser.parse_args()
  229.     print(args)
  230.     sem = threading.Semaphore(args.sem)
  231.     params = {
  232.         "user": args.user,
  233.         "password": args.password,
  234.         "ips": args.ips,
  235.         "port": args.port,
  236.         "sem": sem,
  237.         "files": args.files,
  238.         "dst_path": args.dst_path,
  239.         "cmd": args.cmd,
  240.         "down_host": args.down_host,
  241.         "dir": args.dir
  242.     }
  243.     run(**params)