sftpUtility.py

#!/usr/bin/env python

import threading
import copy
import stat
import pysftp
from dm.common.utility.timeUtility import TimeUtility
from dm.common.utility.loggingManager import LoggingManager
import urlparse

class SftpUtility:

    DEFAULT_PORT = 22

    def __init__(self, host, port=DEFAULT_PORT, username=None, password=None, privateKey=None):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.privateKey = privateKey
        self.sftpClient = None
        self.lock = threading.RLock()

    @classmethod
    def parseFtpUrl(cls, url, defaultHost=None, defaultPort=None):
        host = defaultHost
        port = defaultPort
        scheme = None
        dirPath = url
        if url.startswith('ftp://'):
            parseResult = urlparse.urlparse(url)
            scheme = parseResult.scheme
            netlocTokens = parseResult.netloc.split(':')
            host = netlocTokens[0]
            if len(netlocTokens) > 1:
                port = int(netlocTokens[1])
            dirPath = parseResult.path
        return (scheme, host, port, dirPath)

    @classmethod
    def getSftpClient(cls, host, port=DEFAULT_PORT, username=None, password=None, privateKey=None):
        sftp = pysftp.Connection(host, username=username, password=password, port=port, private_key=privateKey)
        return sftp

    @classmethod
    def getLogger(cls):
        logger = LoggingManager.getInstance().getLogger(cls.__name__)
        return logger


    def __parseKeyValue(cls, keyValue, outputDict={}):
        key,value = keyValue.split('=')
        value = value.strip()
        outputDict[key] = value
        return outputDict
        
    def getFiles(self, dirPath, fileDict={}, replacementDirPath=None):
        self.lock.acquire()
        try:
            if not self.sftpClient:
                self.sftpClient = self.getSftpClient(self.host, self.port, self.username, self.password, self.privateKey)
            if not replacementDirPath:
                replacementDirPath = dirPath
            try:
                attrs = self.sftpClient.listdir_attr(dirPath)
            except Exception, ex:
                self.getLogger().error('Could not retrieve files from %s: %s' % (dirPath,ex))
                self.closeConnection()
                raise
        finally:
            self.lock.release()

        for attr in attrs:
            fileName = attr.filename
            mode = attr.st_mode
            if stat.S_ISDIR(mode):
                dirPath2 = '%s/%s' % (dirPath, fileName)
                replacementDirPath2 = '%s/%s' % (replacementDirPath, fileName)
                self.getFiles(dirPath2, fileDict, replacementDirPath2)
            elif stat.S_ISREG(mode):
                filePath = '%s/%s' % (replacementDirPath, fileName)
                fileInfo = {'fileSize' : attr.st_size, 
                    'fileModificationTime' : attr.st_mtime }
                fileDict[filePath] = fileInfo
        return fileDict

    def getMd5Sum(self, filePath, fileInfo={}):
        self.lock.acquire()
        try:
            if not self.sftpClient:
                self.sftpClient = self.getSftpClient(self.host, self.port, self.username, self.password, self.privateKey)
            try:
                md5Sum = self.sftpClient.execute('md5sum %s' % filePath)[0].split()[0]
                fileInfo['md5Sum'] = md5Sum
            except Exception, ex:
                self.getLogger().error('Could not get md5sum for file %s: %s' % (filePath,ex))
                self.closeConnection()
                raise
            return md5Sum
        finally:
            self.lock.release()

    def statFile(self, filePath, fileInfo={}):
        self.lock.acquire()
        try:
            if not self.sftpClient:
                self.sftpClient = self.getSftpClient(self.host, self.port, self.username, self.password, self.privateKey)
            try:
                attr = self.sftpClient.stat(filePath)
                fileInfo['fileSize'] = attr.st_size 
                fileInfo['fileModificationTime'] = attr.st_mtime
            except Exception, ex:
                self.getLogger().error('Could not get stat file %s: %s' % (filePath,ex))
                self.closeConnection()
                raise
            return fileInfo
        finally:
            self.lock.release()

    def closeConnection(self):
        logger = self.getLogger()
        self.lock.acquire()
        try:
            try:
                if self.sftpClient:
                    logger.warn('Closing SFTP connection to host %s' % self.host)
                    self.sftpClient.close()
            except Exception, ex:
                logger.error('Could not close SFTP connection to host %s: %s' % (self.host, ex))
            self.sftpClient = None
        finally:
            self.lock.release()

#######################################################################
# Testing.

if __name__ == '__main__':
    #sftpUtility = SftpUtility('s1dserv', username='dmadmin', password='theKey12')
    sftpUtility = SftpUtility('s1dserv',privateKey='/home/beams/DMADMIN/.ssh/id_dsa')
    files = sftpUtility.getFiles('/export/dm/test')
    print files
    print sftpUtility.getMd5Sum('/export/dm/test/testfile01')
    print 'Closing connection'
    sftpUtility.closeConnection()
    print sftpUtility.statFile('/export/dm/test/testfile01')