import repository from arizona
[raven.git] / lib / arizona-lib / storkbtdownloadheadless.py
1 #!/usr/bin/python
2 #
3 #storkbtdownloadheadless.py
4
5 """ This is a modified form of btdownloadheadless.py from BitTorrent 4.0.1-1.
6     Modified by Jason Hardies for use with stork via storktorrent.py
7     
8     Changes are as follows:
9         -The original if __name__=="__main__": statement is wrapped by the btdownloadheadless method.
10         -By default, the method will start a quiet bittorrent download that will stop once the download is complete (which can be changed by passing non-default options to the method)
11         -An urllib.urlretrieve like download indicator can be used.
12         -The host and port of the tracker can be specified to override the values found in the torrent files.
13         -The download object (DL) can be directly accessed by passing an empty object to the btdownloadheadless method.  This object can be used to stop the torrent if it was set to seed forever, for example.
14         -A seed locator timeout class is added to raise an error when no seed is located within a certain timeout.
15     (Changes are outlined using the ##CHANGE:  ... ##END blocks.)
16 """
17 # The contents of this file are subject to the BitTorrent Open Source License
18 # Version 1.0 (the License).  You may not copy or use this file, in either
19 # source code or executable form, except in compliance with the License.  You
20 # may obtain a copy of the License at http://www.bittorrent.com/license/.
21 #
22 # Software distributed under the License is distributed on an AS IS basis,
23 # WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
24 # for the specific language governing rights and limitations under the
25 # License.
26
27 # Written by Bram Cohen, Uoti Urpala and John Hoffman
28
29 from __future__ import division
30
31 import sys
32 import os
33 import threading
34 from time import time, strftime
35 from signal import signal, SIGWINCH
36 from cStringIO import StringIO
37
38 from BitTorrent.download import Feedback, Multitorrent
39 from BitTorrent.defaultargs import get_defaults
40 from BitTorrent.parseargs import parseargs, printHelp
41 from BitTorrent.zurllib import urlopen
42 from BitTorrent.bencode import bdecode
43 from BitTorrent.ConvertedMetainfo import ConvertedMetainfo
44 from BitTorrent import configfile
45 from BitTorrent import BTFailure
46 from BitTorrent import version
47
48 ##CHANGE:ADDED:
49 #since time was imported from time, above...
50 from time import sleep as timesleep
51 ##END
52
53 def fmttime(n):
54     if n == 0:
55         return 'download complete!'
56     try:
57         n = int(n)
58         assert n >= 0 and n < 5184000  # 60 days
59     except:
60         return '<unknown>'
61     m, s = divmod(n, 60)
62     h, m = divmod(m, 60)
63     return 'finishing in %d:%02d:%02d' % (h, m, s)
64
65 def fmtsize(n):
66     s = str(n)
67     size = s[-3:]
68     while len(s) > 3:
69         s = s[:-3]
70         size = '%s,%s' % (s[-3:], size)
71     if n > 999:
72         unit = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']
73         i = 1
74         while i + 1 < len(unit) and (n >> 10) >= 999:
75             i += 1
76             n >>= 10
77         n /= (1 << 10)
78         size = '%s (%.0f %s)' % (size, n, unit[i])
79     return size
80
81 ##CHANGE:ADDED:
82 class SeedLocatorTimeout(threading.Thread):
83     """
84         A timer class, added to timeout if no seed is located
85     """
86     def __init__(self, dlobj, timeout=30.0):
87         # default timeout = 30 sec.
88         threading.Thread.__init__(self)
89         self.dlobj=dlobj
90         self.seedlocated=False
91         self.running=False
92         self.timeout=timeout
93     def run(self):
94         self.running=True
95         elapsed=0
96         tstart=time()
97         while self.running and not self.seedlocated:
98             timesleep(1.0)
99             elapsed=time()-tstart
100             if elapsed>=self.timeout:self.running=False
101         if not self.seedlocated:
102             #raise an error via the dl object
103             self.dlobj.error(self.dlobj.torrent,0,"Seed Location Timeout: no seed found after "+str(self.timeout)+" seconds.")
104     def stop(self):
105         self.running=False
106         self.seedlocated=True #don't want to have the run function attempting to raise an error
107 ##END
108
109 class HeadlessDisplayer(object):
110
111     ##CHANGE:
112     def __init__(self, doneflag,quiet=0,downloadindicator=None):
113         ##ORIGINAL:
114         #def __init__(self, doneflag):
115         ##END
116         self.doneflag = doneflag
117         ##CHANGE:ADDED:
118         self.quiet=quiet
119         self.downloadindicator=downloadindicator
120         ##END
121
122         self.done = False
123         self.percentDone = ''
124         self.timeEst = ''
125         self.downRate = '---'
126         self.upRate = '---'
127         self.shareRating = ''
128         self.seedStatus = ''
129         self.peerStatus = ''
130         self.errors = []
131         self.file = ''
132         self.downloadTo = ''
133         self.fileSize = ''
134         self.numpieces = 0
135
136     def set_torrent_values(self, name, path, size, numpieces):
137         self.file = name
138         self.downloadTo = path
139         self.fileSize = fmtsize(size)
140         self.numpieces = numpieces
141
142     def finished(self):
143         self.done = True
144         self.downRate = '---'
145         ##CHANGE:MODIFIED:(added if not self.quiet:)
146         if not self.quiet:self.display({'activity':'download succeeded', 'fractionDone':1})
147         ##END
148
149     def error(self, errormsg):
150         newerrmsg = strftime('[%H:%M:%S] ') + errormsg
151         self.errors.append(newerrmsg)
152         if not self.quiet:self.display({})
153
154     def display(self, statistics):
155         fractionDone = statistics.get('fractionDone')
156         activity = statistics.get('activity')
157         timeEst = statistics.get('timeEst')
158         downRate = statistics.get('downRate')
159         upRate = statistics.get('upRate')
160         spew = statistics.get('spew')
161
162         print '\n\n\n\n'
163         if spew is not None:
164             self.print_spew(spew)
165
166         if timeEst is not None:
167             self.timeEst = fmttime(timeEst)
168         elif activity is not None:
169             self.timeEst = activity
170
171         if fractionDone is not None:
172             self.percentDone = str(int(fractionDone * 1000) / 10)
173         if downRate is not None:
174             self.downRate = '%.1f KB/s' % (downRate / (1 << 10))
175         if upRate is not None:
176             self.upRate = '%.1f KB/s' % (upRate / (1 << 10))
177         downTotal = statistics.get('downTotal')
178         if downTotal is not None:
179             upTotal = statistics['upTotal']
180             if downTotal <= upTotal / 100:
181                 self.shareRating = 'oo  (%.1f MB up / %.1f MB down)' % (
182                     upTotal / (1<<20), downTotal / (1<<20))
183             else:
184                 self.shareRating = '%.3f  (%.1f MB up / %.1f MB down)' % (
185                    upTotal / downTotal, upTotal / (1<<20), downTotal / (1<<20))
186             numCopies = statistics['numCopies']
187             nextCopies = ', '.join(["%d:%.1f%%" % (a,int(b*1000)/10) for a,b in
188                     zip(xrange(numCopies+1, 1000), statistics['numCopyList'])])
189             if not self.done:
190                 self.seedStatus = '%d seen now, plus %d distributed copies ' \
191                                   '(%s)' % (statistics['numSeeds'],
192                                          statistics['numCopies'], nextCopies)
193             else:
194                 self.seedStatus = '%d distributed copies (next: %s)' % (
195                     statistics['numCopies'], nextCopies)
196             self.peerStatus = '%d seen now' % statistics['numPeers']
197
198         for err in self.errors:
199             print 'ERROR:\n' + err + '\n'
200         ##CHANGE:ADDED: (plus indentation on the following lines)
201         if not self.quiet:
202         ##END
203             print 'saving:        ', self.file
204             print 'percent done:  ', self.percentDone
205             print 'time left:     ', self.timeEst
206             print 'download to:   ', self.downloadTo
207             print 'download rate: ', self.downRate
208             print 'upload rate:   ', self.upRate
209             print 'share rating:  ', self.shareRating
210             print 'seed status:   ', self.seedStatus
211             print 'peer status:   ', self.peerStatus
212         ##CHANGE:ADDED:
213         #attempt to use the download_indicator(block_count,block_size,file_size)
214         if not self.downloadindicator is None:
215             self.downloadindicator(int(fractionDone * 1000),1,1000)
216         ##END
217
218     def print_spew(self, spew):
219         s = StringIO()
220         s.write('\n\n\n')
221         for c in spew:
222             s.write('%20s ' % c['ip'])
223             if c['initiation'] == 'L':
224                 s.write('l')
225             else:
226                 s.write('r')
227             total, rate, interested, choked = c['upload']
228             s.write(' %10s %10s ' % (str(int(total/10485.76)/100),
229                                      str(int(rate))))
230             if c['is_optimistic_unchoke']:
231                 s.write('*')
232             else:
233                 s.write(' ')
234             if interested:
235                 s.write('i')
236             else:
237                 s.write(' ')
238             if choked:
239                 s.write('c')
240             else:
241                 s.write(' ')
242
243             total, rate, interested, choked, snubbed = c['download']
244             s.write(' %10s %10s ' % (str(int(total/10485.76)/100),
245                                      str(int(rate))))
246             if interested:
247                 s.write('i')
248             else:
249                 s.write(' ')
250             if choked:
251                 s.write('c')
252             else:
253                 s.write(' ')
254             if snubbed:
255                 s.write('s')
256             else:
257                 s.write(' ')
258             s.write('\n')
259         print s.getvalue()
260
261
262 class DL(Feedback):
263
264     ##CHANGE:
265     def __init__(self, metainfo, config,host=None,port=None,quiet=0,downloadindicator=None,seed=0,sltimeout=30):
266         ##ORIGINAL:
267         #def __init__(self, metainfo, config):
268         ##END
269         self.doneflag = threading.Event()
270         self.metainfo = metainfo
271         self.config = config
272         ##CHANGE:ADDED:
273         self.success=0
274         self.quiet = quiet
275         self.downloadindicator = downloadindicator
276         self.m_host=host
277         self.m_port=port
278         self.seed=seed
279         self.sltimeout=SeedLocatorTimeout(self,sltimeout)
280         ##END
281
282     def run(self):
283         ##CHANGE:
284         self.d = HeadlessDisplayer(self.doneflag,self.quiet,self.downloadindicator)
285         ##ORIGINAL:
286         #self.d = HeadlessDisplayer(self.doneflag)
287         ##END
288         try:
289             self.multitorrent = Multitorrent(self.config, self.doneflag,
290                                              self.global_error)
291             # raises BTFailure if bad
292             metainfo = ConvertedMetainfo(bdecode(self.metainfo))
293             ##CHANGE:ADDED:
294             # set the host and port if specified
295             if not self.m_host is None:
296                 cport=6969
297                 if not self.m_port is None:cport=self.m_port
298                 metainfo.announce="http://"+self.m_host+":"+str(cport)+"/announce"
299             ##END
300             torrent_name = metainfo.name_fs
301             if config['save_as']:
302                 if config['save_in']:
303                     raise BTFailure('You cannot specify both --save_as and '
304                                     '--save_in')
305                 saveas = config['save_as']
306             elif config['save_in']:
307                 saveas = os.path.join(config['save_in'], torrent_name)
308             else:
309                 saveas = torrent_name
310
311             self.d.set_torrent_values(metainfo.name, os.path.abspath(saveas),
312                                 metainfo.total_bytes, len(metainfo.hashes))
313             self.torrent = self.multitorrent.start_torrent(metainfo,
314                                 self.config, self, saveas)
315             ##CHANGE:ADDED:
316             #don't run timeout if it's expecting to run forever
317             if not self.seed:self.sltimeout.start()
318             ##END
319         except BTFailure, e:
320             print str(e)
321             return
322         self.get_status()
323         self.multitorrent.rawserver.listen_forever()
324         ##CHANGE:MODIFIED: (added if not self.quiet)
325         if not self.quiet:self.d.display({'activity':'shutting down', 'fractionDone':0})
326         ##END
327         self.torrent.shutdown()
328
329     def reread_config(self):
330         try:
331             newvalues = configfile.get_config(self.config, 'btdownloadcurses')
332         except Exception, e:
333             self.d.error('Error reading config: ' + str(e))
334             return
335         self.config.update(newvalues)
336         # The set_option call can potentially trigger something that kills
337         # the torrent (when writing this the only possibility is a change in
338         # max_files_open causing an IOError while closing files), and so
339         # the self.failed() callback can run during this loop.
340         for option, value in newvalues.iteritems():
341             self.multitorrent.set_option(option, value)
342         for option, value in newvalues.iteritems():
343             self.torrent.set_option(option, value)
344
345     def get_status(self):
346         self.multitorrent.rawserver.add_task(self.get_status,
347                                              self.config['display_interval'])
348         status = self.torrent.get_status(self.config['spew'])
349         #if a seed has been located, stop the timer...
350         if not self.sltimeout is None and not status.get('downTotal') is None and status['numCopies']>0:self.sltimeout.stop()
351         if not self.quiet:self.d.display(status)
352
353     def global_error(self, level, text):
354         self.d.error(text)
355         
356     def error(self, torrent, level, text):
357         self.d.error(text)
358         ##CHANGE:ADDED:
359         if "Seed Location Timeout" in text:
360             self.success=-1
361             #self.sltimeout.stop()
362             self.doneflag.set()
363             #raise BTFailure(text)
364         elif "urlopen error" in text and "tracker" in text:
365             self.success=-1
366             self.sltimeout.stop()
367             self.doneflag.set()
368             raise BTFailure("Tracker connection error.")
369         ##END
370
371     def failed(self, torrent, is_external):
372         self.doneflag.set()
373         ##CHANGE:ADDED:
374         self.success=-1
375         self.sltimeout.stop()
376         ##END
377
378     def finished(self, torrent):
379         self.d.finished()
380         ##CHANGE:ADDED:
381         self.success=1
382         #forces it to stop once it's done:
383         if not self.seed:
384             self.doneflag.set()
385             self.sltimeout.stop()
386         ##END
387
388     ##CHANGE:ADDED:
389     def stop(self):
390         """attempts to stop the current process"""
391         self.doneflag.set()
392         self.sltimeout.stop()
393     ##END
394
395 ##CHANGE:MODIFIED/ADDED:
396 # (see original below -- basically, made into a method, with extra args and try blocks in __main__)
397 def btdownloadheadless(argv,host=None,port=None,quiet=0,downloadindicator=None,seed=0,dlobj=None,sltimeout=30):
398     """
399         Wraps the original if __name__=="__main__": block and allows for additional functionality
400         
401         Arguments:
402           argv: The argv as the btdownloadheadless.py would receive from sys.argv
403           host: The tracker host for the torrent. (if None, will use the torrent's settings) (default:None)
404           port: The tracker port for the torrent. (if None, will use the torrent's settings) (default: None)
405           quiet: Whether the downloader should print out download status like btdownloadheadless (default: 0)
406           downloadindicator: The download indicator (like would be passed to urllib.urlretrieve) (default: None)
407           seed: Whether the downloader should continue seeding after downloading (note: may not return set to continue seeding) (default:0)
408           dlobj: An object that can be used to get the DL object for the download.  If not None, the DL object can be accessed via dlobj.dl. (default: None)
409           sltimeout: The timeout value in seconds for the seed locator.  (Default: 30)
410           
411         Any method using this method should handle the IOError and BTFailure exceptions that this can raise.
412     """
413     global uiname
414     global defaults
415     global config
416     global args
417     uiname = 'btdownloadheadless'
418     defaults = get_defaults(uiname)
419
420     if len(argv) <= 1:
421         printHelp(uiname, defaults)
422         sys.exit(1)
423     config, args = configfile.parse_configuration_and_args(defaults,
424                                   uiname, argv[1:], 0, 1)
425     if args:
426         if config['responsefile']:
427             raise BTFailure, 'must have responsefile as arg or parameter, not both'
428         config['responsefile'] = args[0]
429     #try:
430     #    try:
431     if config['responsefile']:
432         h = file(config['responsefile'], 'rb')
433         metainfo = h.read()
434         h.close()
435     elif config['url']:
436         h = urlopen(config['url'])
437         metainfo = h.read()
438         h.close()
439     else:
440         raise BTFailure('you need to specify a .torrent file')
441         #don't exit here (not like it would be possible right after raising an error)
442         #sys.exit(1)
443     #except IOError, e:
444     #    raise BTFailure('Error reading .torrent file: ', str(e))
445     #don't handle the error here, pass it up
446     #except BTFailure, e:
447     #    print str(e)
448     #    sys.exit(1)
449
450     dl = DL(metainfo, config,host,port,quiet,downloadindicator,seed,sltimeout)
451     if not dlobj is None:dlobj.dl=dl
452     dl.run()
453     return dl.success
454
455 if __name__ == '__main__':
456     try:
457         try:
458             btdownloadheadless(sys.argv)
459         except IOError,e:
460             raise BTFailure('Error reading .torrent file: '+str(e))
461     except BTFailure, e:
462         #print str(e)
463         #sys.exit(1)
464         raise Exception(str(e))
465 ##ORIGINAL:
466 #if __name__ == '__main__':
467 #    uiname = 'btdownloadheadless'
468 #    defaults = get_defaults(uiname)
469 #
470 #    if len(sys.argv) <= 1:
471 #        printHelp(uiname, defaults)
472 #        sys.exit(1)
473 #    try:
474 #        config, args = configfile.parse_configuration_and_args(defaults,
475 #                                      uiname, sys.argv[1:], 0, 1)
476 #        if args:
477 #            if config['responsefile']:
478 #                raise BTFailure, 'must have responsefile as arg or ' \
479 #                      'parameter, not both'
480 #            config['responsefile'] = args[0]
481 #        try:
482 #            if config['responsefile']:
483 #                h = file(config['responsefile'], 'rb')
484 #                metainfo = h.read()
485 #                h.close()
486 #            elif config['url']:
487 #                h = urlopen(config['url'])
488 #                metainfo = h.read()
489 #                h.close()
490 #            else:
491 #                raise BTFailure('you need to specify a .torrent file')
492 #        except IOError, e:
493 #            raise BTFailure('Error reading .torrent file: ', str(e))
494 #    except BTFailure, e:
495 #        print str(e)
496 #        sys.exit(1)
497 #
498 #    dl = DL(metainfo, config)
499 #    dl.run()
500 ##END