Media-Manager/Universal/Network/sftpservice.py

import time

class SFTPService:
    def __init__(self, logger):
        self.logger = logger
        self.report(2, "Initializing")
        try:
            import paramiko
            self.enabled = True
            self.paramiko = paramiko
            self.report(2, "API Found")
        except ImportError:
            self.enabled = False
            self.report(3, "Module 'paramiko' was not found")
            self.report(3, "Disabling SFTP features")
        self.connected = False

    def isEnabled(self):
        self.report(2, "isEnabled")
        return self.enabled

    def isConnected(self):
        self.report(2, "isConnected")
        return self.connected

    def connectSFTPServer(self, host, port, username, password, key_file, key_type):
        self.report(2, "connectSFTPServer")
        self.host = host
        self.port = int(port)
        self.username = username
        self.password = password
        self.key_file = key_file
        self.key_type = key_type
        tmp_key = None
        if key_file != None:
            if key_file.lower().find("dsa") != -1:
                tmp_key = self.paramiko.DSAKey.from_private_key_file(key_file)
            elif key_file.lower().find("rsa") != -1:
                tmp_key = self.paramiko.RSAKey.from_private_key_file(key_file)
        try:
            tmp = self.paramiko.Transport((host, int(port)))
            tmp.connect(None, username, password, tmp_key)
        except:
            self.connected = False
        try:
            self.sftp_socket = self.paramiko.SFTPClient.from_transport(tmp)
            self.connected = True
            return True
        except:
            self.connected = False
        return False

    def attemptReconnect(self, retries=1):
        self.report(2, "attemptReconnect")
        self.connected = False
        if retries == 1:
            self.report(1, "Lost connect to SFTP Server...")
        if retries <= 5:
            self.report(1, "Attempt #" + str(retries) + "... Attemping to reconnect in 15 seconds...")
            self.closeConnection()
            time.sleep(15)
            if self.connectSFTPServer(self.host, self.port, self.username, self.password, self.key_file, self.key_type):
                self.logger(0, "Successfully reconnected to SFTP Server")
                return True
            else:
                return self.attemptReconnect(retries + 1)
        else:
            self.report(3, "Could not reconnect to SFTP Server...")
            return False

    def closeConnection(self):
        self.report(2, "closeConnection")
        if self.sftp_socket is not None:
            self.sftp_socket.close()
        self.connected = False

    def getFileSize(self, path, retried=0):
        self.report(2, "getFileSize")
        try:
            return int(str(self.sftp_socket.lstat(path)).split()[4])
        except IOError:
            return None
        except:
            if retried == 0:
                if self.attemptReconnect():
                    return self.getFileSize(path, 1)
            else:
                return None

    def doesFileExists(self, path):
        self.report(2, "doesFileExists")
        if self.getFileSize(path):
            return True
        else:
            return False

    def createFolder(self, path):
        self.report(2, "createFolder")
        tmp_full_path = path.split("/")
        tmp_path = tmp_full_path[0]
        for current_path in tmp_full_path:
            if tmp_path != current_path:
                tmp_path = tmp_path + "/" + current_path
            try:
                self.sftp_socket.mkdir(tmp_path)
            except:
                pass

    def deleteFile(self, path):
        self.report(2, "deleteFile")
        try:
            self.sftp_socket.remove(path)
            return True
        except:
            self.report(3, "Failed to delete file")

    def uploadFile(self, input_path, output_path, retried=0):
        self.report(2, "uploadFile")
        try:
            self.createFolder(output_path.replace(output_path.split("/")[-1], ""))
            if self.doesFileExists(output_path):
                self.deleteFile(output_path)
            return self.sftp_socket.put(input_path, output_path)
        except:
            if retried == 0:
                if self.attemptReconnect():
                    return self.uploadFile(input_path, output_path, 1)
            else:
                return False

    def report(self, level, message):
        match level:
            case 0:
                self.logger.infoReport("[SFTP] " + message)
            case 1:
                self.logger.warningReport("[SFTP] " + message)
            case 2:
                self.logger.debugReport("[SFTP] " + message)
            case 3:
                self.logger.errorReport("[SFTP] " + message)