import paramiko from scp import SCPClient from pathlib import Path,PurePosixPath import sys from easydict import EasyDict import json import stat from tqdm import tqdm import traceback def load_json(json_path: str): with open(json_path, "r") as f: data = json.load(f) return data def write_json(data: dict, json_path: str): with open(json_path, "w") as f: json.dump(data, f,ensure_ascii=False,indent=4) class autoFile(object): def __init__(self,cache_path='.'): self.ssh:paramiko.SSHClient=None self.scp:SCPClient=None self.sftp:paramiko.SFTPClient=None self.config:EasyDict=EasyDict(dict( connection={}, path_map={}, cache_path=cache_path, max_filename_len=120, )) def connect_with_jump(self, jump_server_ip='', jump_server_port=22, jump_server_username='root', jump_server_password='', remote_host_ip='', remote_host_port=22, remote_host_username='root', remote_host_password='', private_key_file=None ): self.config.connection.update(dict( jump_server_ip=jump_server_ip, jump_server_port=jump_server_port, jump_server_username=jump_server_username, jump_server_password=jump_server_password, remote_host_ip=remote_host_ip, remote_host_port=remote_host_port, remote_host_username=remote_host_username, remote_host_password=remote_host_password, private_key_file=private_key_file, )) self.real_connect_with_jump() def real_connect_with_jump(self): jump_server_ip=self.config.connection.jump_server_ip jump_server_port=self.config.connection.jump_server_port jump_server_username=self.config.connection.jump_server_username jump_server_password=self.config.connection.jump_server_password remote_host_ip=self.config.connection.remote_host_ip remote_host_port=self.config.connection.remote_host_port remote_host_username=self.config.connection.remote_host_username remote_host_password=self.config.connection.remote_host_password private_key_file=self.config.connection.private_key_file jump_server = paramiko.SSHClient() jump_server.set_missing_host_key_policy(paramiko.AutoAddPolicy()) # 自动接受key print(f"connecting jump server {jump_server_username}@{jump_server_ip}") jump_server.connect(hostname=jump_server_ip, port=jump_server_port, username=jump_server_username, password=jump_server_password) # 跳板机连接 print(f"connect done") jump_transport = jump_server.get_transport() # 创建Transport对象 # 建立隧道 jump_channel = jump_transport.open_channel(kind="direct-tcpip", dest_addr=(remote_host_ip, remote_host_port), src_addr=(jump_server_ip, jump_server_port)) print(f"connecting remote server {remote_host_username}@{remote_host_ip}") remote_host = paramiko.SSHClient() remote_host.set_missing_host_key_policy(paramiko.AutoAddPolicy()) # 自动接受key if private_key_file is not None: # 使用密钥连接（通过隧道） private_key = paramiko.RSAKey.from_private_key_file('/root/.ssh/id_rsa') remote_host.connect(remote_host_ip, username=remote_host_username, pkey=private_key, sock=jump_channel) else: # 使用密码连接（通过隧道） remote_host.connect(hostname=remote_host_ip, port=remote_host_port, username=remote_host_username, password=remote_host_password, sock=jump_channel,allow_agent=False,look_for_keys=False) stdin, stdout, stderr = remote_host.exec_command("cat /proc/version") print(f"connect done: {stdout.read().decode()}") self.ssh=remote_host try: self.sftp=self.ssh.open_sftp() except Exception as e: print(f"[WARNNING]SFTP connect failed: {e}\n[WARNNING]This may cause the folder cannot not be downloaded automatically") def progress_callback(filename, size, sent): if len(filename) > self.config.max_filename_len: filename=f"...{filename[(len(filename)-self.config.max_filename_len):]}" sys.stdout.write("%s\'s progress: %.2f%% \r" % (filename, float(sent)/float(size)*100) ) scp=SCPClient(remote_host.get_transport(),progress=progress_callback) self.scp=scp def getfile(self,filepath,force_override=False): mapped_path,relative_path,mapped=self.mapping(str(filepath)) real_path=mapped_path if mapped: real_path:Path=Path(self.config.cache_path) / relative_path if self.scp is not None and (not real_path.exists() or force_override): real_path.parent.mkdir(parents=True,exist_ok=True) self.download(str(mapped_path),str(real_path)) return real_path def getdir(self,dirpath,force_override=False,recursion_depth=0): mapped_path,relative_path,mapped=self.mapping(str(dirpath)) real_path=mapped_path cache_path=Path(self.config.cache_path) if mapped: real_path:Path=cache_path / relative_path if self.sftp is None: return real_path is_file=True try: is_file=stat.S_ISREG(self.sftp.stat(str(mapped_path)).st_mode) except Exception as e: print(f"Get file type failed: {e}") if is_file: return self.getfile(dirpath,force_override) real_path.mkdir(parents=True,exist_ok=True) dir=PurePosixPath(dirpath) for file in tqdm(self.sftp.listdir(str(mapped_path)), leave=False): filepath=mapped_path / file is_file_child=stat.S_ISREG(self.sftp.stat(str(filepath)).st_mode) if is_file_child: self.getfile(dir/file, force_override) elif recursion_depth>0: self.getdir(dir/file, force_override, recursion_depth-1) return real_path def get(self,dirorfile,force_override=False,recursion_depth=0): if self.sftp is None: return self.getfile(dirorfile, force_override) return self.getdir(dirorfile,force_override,recursion_depth) def download(self,filepath,local_path): self.scp.get(filepath,local_path) # sys.stdout.write('\n') #the progress bar didn't warp return local_path def set_path_mapping(self,input_prefix,target_prefix=None,**kargs): self.config.path_map[input_prefix]=target_prefix if target_prefix is not None else input_prefix def mapping(self,filepath:str): for key in self.config.path_map: if filepath.startswith(key): local_path=PurePosixPath(filepath) relative_path=local_path.relative_to(key) map_path=PurePosixPath(self.config.path_map[key]) mappedpath=map_path / relative_path return mappedpath,relative_path,True return Path(filepath),None,False def save_config(self,filepath): write_json(self.config,filepath) def load_config(self,file_path,connect=True): self.config.update(load_json(file_path)) if connect: self.real_connect_with_jump() def __del__(self): if self.scp is not None: self.scp.close() if self.sftp is not None: self.sftp.close() if self.ssh is not None: self.ssh.close()