code.vuplus.com Git - vuplus_dvbapp-plugin/blob - mediadownloader/src/FTPProgressDownloader.py

   1 from twisted.internet import reactor, defer
   2 from twisted.internet.protocol import Protocol, ClientCreator
   3 from twisted.protocols.ftp import FTPClient, FTPFileListProtocol
   4
   5 from os import SEEK_END
   6
   7 # XXX: did I ever actually test supportPartial?
   8 class FTPProgressDownloader(Protocol):
   9         """Download to a file from FTP and keep track of progress."""
  10
  11         def __init__(self, host, port, path, fileOrName, username = 'anonymous', \
  12                 password = 'my@email.com', writeProgress = None, passive = True, \
  13                 supportPartial = False, *args, **kwargs):
  14
  15                 timeout = 30
  16
  17                 # We need this later
  18                 self.path = path
  19                 self.resume = supportPartial
  20
  21                 # Initialize
  22                 self.currentlength = 0
  23                 self.totallength = None
  24                 if writeProgress and type(writeProgress) is not list:
  25                         writeProgress = [ writeProgress ]
  26                 self.writeProgress = writeProgress
  27
  28                 # Output
  29                 if isinstance(fileOrName, str):
  30                         self.filename = fileOrName
  31                         self.file = None
  32                 else:
  33                         self.file = fileOrName
  34
  35                 creator = ClientCreator(reactor, FTPClient, username, password, passive = passive)
  36
  37                 creator.connectTCP(host, port, timeout).addCallback(self.controlConnectionMade).addErrback(self.connectionFailed)
  38
  39                 self.deferred = defer.Deferred()
  40
  41         def controlConnectionMade(self, ftpclient):
  42                 # We need the client locally
  43                 self.ftpclient = ftpclient
  44
  45                 # Try to fetch filesize
  46                 self.ftpFetchSize()
  47
  48         # Handle recieved msg
  49         def sizeRcvd(self, msgs):
  50                 # Split up return
  51                 code, msg = msgs[0].split()
  52                 if code == '213':
  53                         self.totallength = int(msg)
  54                         for cb in self.writeProgress or [ ]:
  55                                 cb(0, self.totallength)
  56
  57                         # We know the size, so start fetching
  58                         self.ftpFetchFile()
  59                 else:
  60                         # Error while reading size, try to list it
  61                         self.ftpFetchList()
  62
  63         def ftpFetchSize(self):
  64                 d = self.ftpclient.queueStringCommand('SIZE ' + self.path)
  65                 d.addCallback(self.sizeRcvd).addErrback(self.ftpFetchList)
  66
  67         # Handle recieved msg
  68         def listRcvd(self, *args):
  69                 # Quit if file not found
  70                 if not len(self.filelist.files):
  71                         self.connectionFailed()
  72                         return
  73
  74                 self.totallength = self.filelist.files[0]['size']
  75                 for cb in self.writeProgress or [ ]:
  76                         cb(0, self.totallength)
  77
  78                 # Invalidate list
  79                 self.filelist = None
  80
  81                 # We know the size, so start fetching
  82                 self.ftpFetchFile()
  83
  84         def ftpFetchList(self, *args, **kwargs):
  85                 self.filelist = FTPFileListProtocol()
  86                 d = self.ftpclient.list(self.path, self.filelist)
  87                 d.addCallback(self.listRcvd).addErrback(self.connectionFailed)
  88
  89         def openFile(self):
  90                 if self.resume:
  91                         file = open(self.filename, 'ab')
  92                 else:
  93                         file = open(self.filename, 'wb')
  94
  95                 return (file, file.tell())
  96
  97         def ftpFetchFile(self):
  98                 offset = 0
  99
 100                 # Finally open file
 101                 if self.file is None:
 102                         try:
 103                                 self.file, offset = self.openFile()
 104                         except IOError, ie:
 105                                 # TODO: handle exception
 106                                 raise ie
 107
 108                 offset = self.resume and offset or 0
 109
 110                 d = self.ftpclient.retrieveFile(self.path, self, offset = offset)
 111                 d.addCallback(self.ftpFinish).addErrback(self.connectionFailed)
 112
 113         def dataReceived(self, data):
 114                 if not self.file:
 115                         return
 116
 117                 if self.writeProgress:
 118                         self.currentlength += len(data)
 119                         for cb in self.writeProgress:
 120                                 cb(self.currentlength, self.totallength)
 121                 try:
 122                         # XXX: why did i always seek? do we really need this?
 123                         if self.resume:
 124                                 self.file.seek(0, SEEK_END)
 125
 126                         self.file.write(data)
 127                 except IOError, ie:
 128                         # TODO: handle exception
 129                         self.file = None
 130                         raise ie
 131
 132         def ftpFinish(self, code = 0, message = None):
 133                 self.ftpclient.quit()
 134                 if self.file is not None:
 135                         self.file.close()
 136                 self.deferred.callback(code)
 137
 138         def connectionFailed(self, reason = None):
 139                 if self.file is not None:
 140                         self.file.close()
 141                 self.deferred.errback(reason)