code.delx.au - webdl/blob - common.py

   1 # vim:ts=4:sts=4:sw=4:noet
   2
   3 from lxml import etree
   4 import json
   5 try:
   6         import hashlib
   7 except ImportError:
   8         import md5 as hashlib
   9 import os
  10 import shutil
  11 import signal
  12 import subprocess
  13 import sys
  14 import tempfile
  15 import time
  16 import urllib
  17
  18
  19 import autosocks
  20 autosocks.try_autosocks()
  21
  22 CACHE_DIR = os.path.expanduser("~/.cache/webdl")
  23
  24 class Node(object):
  25         def __init__(self, title, parent=None):
  26                 self.title = title
  27                 if parent:
  28                         parent.children.append(self)
  29                 self.parent = parent
  30                 self.children = []
  31                 self.can_download = False
  32
  33         def download(self):
  34                 raise NotImplemented
  35
  36
  37 def load_root_node():
  38         root_node = Node("Root")
  39
  40         print "Loading iView episode data...",
  41         sys.stdout.flush()
  42         import iview
  43         iview_node = Node("ABC iView", root_node)
  44         iview.fill_nodes(iview_node)
  45         print "done"
  46
  47         print "Loading SBS episode data...",
  48         sys.stdout.flush()
  49         import sbs
  50         sbs_node = Node("SBS", root_node)
  51         sbs.fill_nodes(sbs_node)
  52         print "done"
  53
  54         return root_node
  55
  56
  57 def urlopen(url, max_age):
  58         if not os.path.isdir(CACHE_DIR):
  59                 os.makedirs(CACHE_DIR)
  60
  61         if max_age <= 0:
  62                 return urllib.urlopen(url)
  63
  64         filename = hashlib.md5(url).hexdigest()
  65         filename = os.path.join(CACHE_DIR, filename)
  66         if os.path.exists(filename):
  67                 file_age = int(time.time()) - os.path.getmtime(filename)
  68                 if file_age < max_age:
  69                         return open(filename)
  70
  71         src = urllib.urlopen(url)
  72         dst = open(filename, "w")
  73         shutil.copyfileobj(src, dst)
  74         src.close()
  75         dst.close()
  76
  77         return open(filename)
  78
  79 def grab_xml(url, max_age):
  80         f = urlopen(url, max_age)
  81         doc = etree.parse(f)
  82         f.close()
  83         return doc
  84
  85 def grab_json(url, max_age):
  86         f = urlopen(url, max_age)
  87         doc = json.load(f)
  88         f.close()
  89         return doc
  90
  91 def download_rtmp(filename, vbase, vpath, hash_url=None):
  92         if vpath.endswith(".flv"):
  93                 vpath = vpath[:-4]
  94         cmd = [
  95                 "rtmpdump",
  96                 "-o", filename,
  97                 "-r", vbase,
  98                 "-y", vpath,
  99         ]
 100         if hash_url is not None:
 101                 cmd += ["--swfVfy", hash_url]
 102         try:
 103                 p = subprocess.Popen(cmd)
 104                 ret = p.wait()
 105                 if ret != 0:
 106                         print >>sys.stderr, "rtmpdump exited with error code:", ret
 107                         return False
 108                 else:
 109                         return True
 110         except OSError, e:
 111                 print >>sys.stderr, "Failed to run rtmpdump!", e
 112                 return False
 113         except KeyboardInterrupt:
 114                 print "Cancelled", cmd
 115                 try:
 116                         p.terminate()
 117                         p.wait()
 118                 except KeyboardInterrupt:
 119                         p.send_signal(signal.SIGKILL)
 120                         p.wait()
 121
 122 def download_urllib(filename, url):
 123         print "Downloading: %s -> %s" % (url, filename)
 124         try:
 125                 src = urllib.urlopen(url)
 126                 dst = open(filename, "w")
 127                 shutil.copyfileobj(src, dst)
 128                 return True
 129         except KeyboardInterrupt:
 130                 print "\nCancelled", url
 131         finally:
 132                 src.close()
 133                 dst.close()
 134         return False
 135