[MerlinMusicPlayer] fixed google cover display
[enigma2-plugins.git] / 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 as 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 as 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)