from ftplib import error_perm from ftplib import FTP import logging from time import sleep import os import re from vector_utils.connections.connection import Connection class FtpConnection(Connection): """The connection class is used for FTP connections. Attributes: conn_obj (ConnectionInfo object) connection (any, optional) """ @Connection.handle_exceptions def __init__(self, conn_obj, logger=None): super().__init__(conn_obj) self.connection = FTP() self.connection.connect( self.conn_obj.domain_name, self.conn_obj.port) self._logger = logger or self._setup_logger() self.login() @property def logger(self): return self._logger @logger.setter def logger(self, value): self._logger = value def _setup_logger(self): """Setup logger helper.""" logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) return logger @Connection.handle_exceptions def check_connection(self, remote_initial_dir): """Check connection to the FTP server.""" self.connection.cwd(remote_initial_dir) @Connection.handle_exceptions def close_connection(self): """Close the client connection.""" self.connection.close() @Connection.handle_exceptions def delete(self, filename, remove_dir=False): pass @Connection.handle_exceptions def file_exists(self, file_name): pass @Connection.handle_exceptions def file_size(self, file_name): pass @Connection.handle_exceptions def is_dir(self, dirname): pass @Connection.handle_exceptions def login(self): """Login and setup connection to ftp server.""" self.connection.login( self.conn_obj.user_name, self.conn_obj.password) @Connection.handle_exceptions def mkdir(self, dest_dir): """Create directories recursively. Args: dest_dir (str): The full directory path. Returns: bool: True if we were able to create the path """ if dest_dir == '/': self.connection.cwd('/') return if dest_dir == '': return try: self.connection.cwd(dest_dir) except error_perm: dirname, basename = os.path.split(dest_dir.rstrip('/')) self.mkdir(dirname) self.connection.mkd(basename) self.connection.cwd(basename) return True @Connection.handle_exceptions def rmdir(self, dir_name): pass @Connection.handle_exceptions def reconnect(self): self.connection.connect( self.conn_obj.domain_name, self.conn_obj.port) self.login() @Connection.handle_exceptions def scan_dir(self, directory, exceptions=[]): """Return dict of filename or dirname mapping to its size. Args: directory (str): The path which you which to scan. exceptions (list): The files or folders to exclude. Returns: dict: Mapping of each file or directory to its size. """ dir_list = [] self.connection.cwd(directory) self.connection.retrlines('LIST', dir_list.append) dir_list = list(map(create_filesize_dict, dir_list)) dir_dict = {} for item in dir_list: for filename, filesize in item.items(): dir_dict[filename] = filesize return dir_dict @Connection.handle_exceptions def transfer_files(self, file_list, transfer_mode='upload', batch_file_id=None, pid=None, overwrite=None, check_delivered_files=True, d_job=None): """Transfer files up or down. Args: file_list (list(dict)): A list of remote/local keypairs transfer_mode (str): upload or download batch_file_id: pid: overwrite: check_delivered_files: d_job: Returns: True if succeeds the upload succeeded for entire list Raises: Exception: We are unable to transfer files. """ attempts = 1 directories = {os.path.dirname(x['remote']) for x in file_list} try: self.check_connection('/') except: self.reconnect() for directory in directories: self.mkdir(directory) while attempts <= 5: try: for file in file_list: local_file = file['local'] remote_file = file['remote'] if transfer_mode == 'download': self._download_file(remote_file, local_file) else: self._upload_file(local_file, remote_file) return True except Exception as e: attempts += 1 self.reconnect() if attempts > 5: raise e sleep(15) raise Exception('Unable to transfer files') def _upload_file(self, source_path, target_path): """Upload a file to its remote location. Args: source_path (str): The full file source path. target_path (str): The full file target path. Raises: Exception: The file was not uploaded successfully. """ filesize = os.stat(source_path).st_size with open(source_path, 'rb') as file: self.connection.storbinary( 'STOR {}'.format(target_path), file) dir_list = self.scan_dir(os.path.dirname(target_path)) if filesize != dir_list[os.path.basename(target_path)]: raise Exception( 'The file {} was not ' 'uploaded successfully'.format(source_path)) def _download_file(self, remote_path, local_path): """Download a file from its remote location. Args: remote_path (str): The remote path of a file local_path (str): The local path to save a file. Raises: Exception: The file was not downloaded successfully. Exception: The file you are attempting to download does not exist. """ remote_dir = self.scan_dir(os.path.dirname(remote_path)) try: filesize = remote_dir[os.path.basename(remote_path)] except KeyError: raise Exception( 'The file {} you are attempting ' 'to download doesn\'t exist.'.format(remote_path)) else: with open(local_path, 'wb') as local_file: self.connection.retrbinary( 'RETR {}'.format(remote_path), local_file.write) if filesize != os.stat(local_path).st_size: raise Exception( 'The file {} was not ' 'downloaded successfully'.format(remote_path)) def parse_list_format_linux(list_line): """Parse the list_line from LIST in linux format.""" parsed_line = re.split('[\\s]+', list_line) filename = parsed_line[8] filesize = int(parsed_line[4]) return { filename: filesize } def parse_list_format_windows(list_line): """Parse the list_line from LIST in windows format.""" parsed_line = re.split('[\\s]+', list_line) if parsed_line[2] == '': filename = parsed_line[3] filesize = 0 else: filename = parsed_line[3] filesize = int(parsed_line[2]) return { filename: filesize } def create_filesize_dict(list_line): """Build a dict from LIST command output. We attempt multiple LIST response formats here. If unable to parse without exception then attempt next format. Args: list_line (str): a line from the output of LIST command Returns: dict: a dict containing a filename mapping to its size """ try: return parse_list_format_linux(list_line) except: pass try: return parse_list_format_windows(list_line) except: pass raise Exception('could not parse LIST response')