tivomirror/tivomirror.py

#!/usr/local/bin/python 
# -*- coding: utf8 -*-

# Download shows from the Tivo

import sys
reload(sys)
sys.setdefaultencoding('utf-8')

import anydbm
import cookielib
import datetime
import getopt
import errno
import functools
import logging
import logging.handlers
import os
import pytz
import re
import requests
import signal
import shutil
import subprocess
import sys
import threading
import time
import urllib2
import xml.dom.minidom
import yaml




class Config:
	config = '~/.tivo/config.yaml'
	cookies = "cookies.txt"
	gig = 1024.0 * 1024 * 1024
	headers = requests.utils.default_headers()
	host = "tivo.lassitu.de"
	ignoreepisodetitle = False
	mak = "7194378159"
	minfree = 10 * gig
	proxies = None
	targetdir = "/p2/media/video/TV"
	tivodecode = "tivodecode"
	useragent = 'Mozilla/5.0 (Windows NT 10.0; WOW64; rv:52.0) Gecko/20100101 Firefox/52.0'

	def __init__(self, file=None):
		self.postprocess = None
		if file:
			self.config = file
		self.load(self.config)

		self.headers.update({ 'User-Agent': self.useragent })
		requests.packages.urllib3.disable_warnings()
		self.session = requests.session()
		self.session.verify = False
		self.session.auth = requests.auth.HTTPDigestAuth("tivo", self.mak)
		self.session.keep_alive = False
		self.session.proxies = self.proxies

	def load(self, file):
		file = os.path.expanduser(file)
		with open(file, 'r') as f:
			y = yaml.load(f)

		for key in y:
			setattr(self, key, y[key])

		for show in self.shows:
			if isinstance(show, dict):
				if 'series' in show:
					IncludeShow(show['series'],
						short=show.get('short', show['series']),
						unique=show.get('unique'))
				else:
					logger.error("Need either a string, or a dict with 'series' and 'short' entries for a show, got \"{}\".".format\
						(show))
					sys.exit(1)
			else:
				IncludeShow(show)

	def __repr__(self):
		return "Config options for tivomirror (singleton)"

config = None

class IncludeShow:
	includes = dict()
	
	def __init__(self, title, short=None, unique=True):
		self.short = short
		self.title = title
		self.timestamp = False
		self.unique = unique or unique == None
		self.includes[title] = self




logger = logging.getLogger('tivomirror')
logger.setLevel(logging.INFO)


class flushfile(object):
	def __init__(self, f):
		self.f = f
	def write(self, x):
		self.f.write(x)
		self.f.flush()
sys.stdout = flushfile(sys.stdout)

tmp = "/tmp"

# prepare global requests sesssion to download the TOC and the episodes


def roundTime(dt=None, roundTo=60):
   """
   http://stackoverflow.com/questions/3463930/how-to-round-the-minute-of-a-datetime-object-python
   """
   if dt == None : dt = datetime.datetime.now()
   seconds = (dt.replace(tzinfo=None) - dt.min).seconds
   rounding = (seconds+roundTo/2) // roundTo * roundTo
   return dt + datetime.timedelta(0,rounding-seconds,-dt.microsecond)


class TimeoutError(Exception):
    pass

def timeout(seconds=10, error_message=os.strerror(errno.ETIMEDOUT)):
	def decorator(func):
		def _handle_timeout(signum, frame):
			raise TimeoutError(error_message)

		def wrapper(*args, **kwargs):
			signal.signal(signal.SIGALRM, _handle_timeout)
			signal.alarm(seconds)
			try:
				result = func(*args, **kwargs)
			finally:
				signal.alarm(0)
			return result

		return functools.wraps(func)(wrapper)

	return decorator


def trimDescription(desc):
	desc = desc.strip()
	i = desc.rfind(". Copyright Tribune Media Services, Inc.");
	if i > 0:
		desc = desc[0:i]
	i = desc.rfind(". * Copyright Rovi, Inc");
	if i > 0:
		desc = desc[0:i]
	i = desc.rfind(". Copyright Rovi, Inc");
	if i > 0:
		desc = desc[0:i]
	if len(desc) > 80:
		desc = desc[0:80]
	return desc

def saveCookies(session, filename):
	cj = cookielib.MozillaCookieJar(filename)
	for cookie in session.cookies:
		logger.debug("storing cookie {}".format(cookie))
		cj.set_cookie(cookie)
	logger.debug("Saving cookies to {}".format(cj))
	cj.save(ignore_discard=True, ignore_expires=True)


class TivoException(Exception):
	def __init__(self, value):
		self.value = value
	def __str__(self):
		return repr(self.value)

class TivoItem:
	def __init__(self, i):
		self.title = getTagText(i, "Title")
		self.episode = getTagText(i, "EpisodeTitle")
		self.episodeNumber = getTagText(i, "EpisodeNumber")
		self.description = trimDescription(getTagText(i, "Description"))
		d = getTagText(i, "CaptureDate")
		self.date = datetime.datetime.fromtimestamp(int(d, 16), pytz.utc)
		self.time = int(d, base=0)
		est = pytz.timezone('US/Eastern')
		eastern = roundTime(self.date, 15*60).astimezone(est)
		self.datestr = self.date.strftime("%Y%m%d-%H%M")
		self.shortdate = eastern.strftime("%m%d-%H%M")
		self.url = getTagText(i, "Url")
		self.url = self.url + "&Format=video/x-tivo-mpeg-ts"
		self.inprogress = getTagText(i, "InProgress")
		self.available = getTagText(i, "Available")
		self.sourcesize = int(getTagText(i, "SourceSize"))
		self.highdef = getTagText(i, "HighDefinition")
		self.unique = True
		if config.ignoreepisodetitle:
			self.episode = self.datestr
		if self.episode == "":
			if self.description != "":
				self.episode = self.description
			else:
				self.episode = self.datestr
		self.formatnames()
	def makeNotUnique(self):
		self.unique = False
		self.formatnames()
	def formatnames(self):
		if self.episodeNumber and self.episodeNumber != u'0':
			en = int(self.episodeNumber)
			if en >= 100:
				self.name = "{} S{:02d}E{:02d} {}".format(self.title, en / 100, en % 100, self.episode)
			else:
				self.name = "{} E{} {}".format(self.title, self.episodeNumber, self.episode)
		elif self.unique:
			self.name = "{} - {}".format(self.title, self.episode)
		else:
			self.name = "{} - {} - {}".format(self.title, self.datestr, self.episode)
		self.dir  = "{}/{}".format(config.targetdir, re.sub("[:/]", "-", self.title))
		self.file = "{}/{}".format(self.dir, re.sub("[:/]", "-", self.name))
		self.name = self.name.encode("utf-8");
		self.dir  = self.dir.encode("utf-8");
		self.file = self.file.encode("utf-8");
	def getPath(self, options):
		title = self.title
		if options.short:
			title = options.short
		if self.episodeNumber and self.episodeNumber != u'0':
			en = int(self.episodeNumber)
			if en >= 100:
				name = "{} S{:02d}E{:02d} {}".format(title, en / 100, en % 100, self.episode)
			else:
				name = "{} E{} {}".format(title, self.episodeNumber, self.episode)
		elif self.unique:
			name = "{} - {}".format(title, self.episode)
		else:
			name = "{} - {} {}".format(title, self.shortdate, self.episode)
		path = "{}/{}".format(self.dir, re.sub("[:/]", "-", name))
		return path.encode("utf-8");
	def __str__(self):
		return repr(self.title)


class TivoToc:
	def __init__(self):
		self.dom = None
		self.filename = "toc.xml"
		self.uniquedb = anydbm.open("unique.db", "c")
		self.items = []
		pass

	def load(self):
		fd = open(self.filename, "r")
		self.dom = xml.dom.minidom.parseString(fd.read())
		fd.close()
		return self.dom

	def save(self):
		fd = open(self.filename, "w")
		fd.write(self.dom.toprettyxml())
		fd.close()

	def download_chunk(self, offset):
		global config

		params = {
			'Command': 'QueryContainer',
			'Container': '/NowPlaying',
			'Recurse': 'Yes',
			'ItemCount': '50',
			'AnchorOffset': offset
		}
		url = "https://{}/TiVoConnect".format(config.host)
		logger.debug("    offset {}".format(offset))
		r = config.session.get(url, params=params, timeout=30, verify=False, proxies=config.proxies, headers=config.headers)
		if r.status_code != 200:
			r.raise_for_status()
		return r.text

	def download(self):
		global config
		offset = 0
		itemCount = 1
		self.dom = None
		root = None
		logger.info("*** Getting listing")
		while itemCount > 0:
			dom = xml.dom.minidom.parseString(self.download_chunk(offset))
			if self.dom == None:
				self.dom = dom
				root = self.dom.childNodes.item(0)
			else:
				for child in dom.childNodes.item(0).childNodes:
					if child.nodeName == "Item":
						root.appendChild(child.cloneNode(True))
			itemCount = int(getElementText(dom.documentElement.childNodes, "ItemCount"))
			offset += itemCount
		saveCookies(config.session, config.cookies)
		return self.dom

	def getItems(self):
		self.titles = {}
		for node in self.dom.getElementsByTagName("Item"):
			item = TivoItem(node)
			self.items.append(item)
			if item.title not in self.titles:
				self.titles[item.title] = []
			self.titles[item.title].append(item)
		# see if we have items that end up having an identical name; mark
		# the program title in uniquedb if that's the case
		for title in self.titles:
			names = {}
			for item in self.titles[title]:
				if item.name not in names:
					names[item.name] = []
				names[item.name].append(item)
			for name in names:
				utf8title = title.encode("utf-8")
				if len(names[name]) > 1 and not self.uniquedb.has_key(utf8title):
					self.uniquedb[utf8title] = "1"
		if getattr(self.uniquedb, "sync", None) and callable(self.uniquedb.sync):
			self.uniquedb.sync()
		# update all items based on config and uniquedb
		for item in self.items:
			multiple = None
			options = IncludeShow.includes.get(title)
			if options and not options.unique:
				item.makeNotUnique()
		return self.items


def getText(nodelist):
	rc = ""
	for node in nodelist:
		if node.nodeType == node.TEXT_NODE:
			rc = rc + node.data
	return rc

def getTagText(element, tagname):
	try:
		return getText(element.getElementsByTagName(tagname)[0].childNodes)
	except IndexError:
		return ""

def getElementText(nodes, name):
	for node in nodes:
		if node.nodeType == xml.dom.Node.ELEMENT_NODE and node.nodeName == name:
			return getText(node.childNodes)
	return None

def getAvail(dir):
	s = os.statvfs(dir)
	return s.f_bsize * s.f_bavail


class FdLogger(threading.Thread):
	def __init__(self, logger, lvl, fd):
		self.logger = logger
		self.lvl = lvl
		self.fd = fd
		threading.Thread.__init__(self)
		self.daemon = True
		self.start()

	def run(self):
		try:
			# for line in fd buffers, so use this instead
			for line in iter(self.fd.readline, b''):
				self.logger.log(self.lvl, ":   %s", line.strip('\n'))
			self.fd.close()
		except Exception:
			self.logger.exception("")


@timeout(43200)
def download_item(item, mak, target):
	global config
	count = 0
	start = time.time()
	upd = start
	url = item.url
	#url = re.sub("tivo.lassitu.de:80", "wavehh.lassitu.de:30080", url)
	logger.info("--- downloading \"{}\"".format(url))
	logger.info("    {}".format(target))
	start = time.time()
	r = config.session.get(url, stream=True, verify=False, proxies=config.proxies, headers=config.headers)
	r.raise_for_status()

	try:
		p_decode = subprocess.Popen([config.tivodecode, "--mak", mak, \
			"--no-verify", "--out", target, "-"], stdin=subprocess.PIPE,
			stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		FdLogger(logger, logging.INFO, p_decode.stdout)
		FdLogger(logger, logging.INFO, p_decode.stderr)
		def info(signum, frame):
			upd = time.time()
			dur = now - start
			mb = count / 1e6
			print "{:5.1f}% {:5.3f} GB downloaded in {:.0f} min, {.3f} MB/s".format(
				100.0 * count / item.sourcesize,
				mb / 1e3, dur / 60, mb / dur)
		try:
			signal.signal(signal.SIGINFO, info)
		except Exception:
			pass
		while True:
			time.sleep(0) # yield to logger threads
			chunk = r.raw.read(256*1024)
			if not chunk:
				break
			p_decode.stdin.write(chunk)
			count += len(chunk)
			now = time.time()
			if (now - upd) > 60:
				upd = now
				dur = now - start
				mb = count / 1e6
				logger.debug("    {:5.1f}% {:5.3f} GB downloaded in {:.0f} min, {:.3f} MB/s".format(
					100.0 * count / item.sourcesize,
					mb / 1e3, dur / 60, mb / dur))
	except Exception as e:
		logger.error("problem decoding: {}".format(e))
		raise
	finally:
		try:
			signal.signal(signal.SIGINFO, signal.SIG_IGN)
		except Exception:
			pass
		elapsed = time.time() - start
		throughput = count / elapsed
		logger.info("{:5.3f} GB transferred in {:d}:{:02d}, {:.1f} MB/s".format(
			count/1e9, int(elapsed/3600), int(elapsed / 60) % 60, throughput/1e6))
		try:
			p_decode.stdin.close()
			p_decode.poll()
			if p_decode.returncode == None:
				time.sleep(1)
				p_decode.poll()
				if p_decode.returncode == None:
					logger.debug("terminating tivodecode")
					p_decode.terminate()
		except Exception, e:
			pass
		p_decode.wait()
		logger.info("tivodecode exited with {}".format(p_decode.returncode))
	size = os.path.getsize(target)
	if size < 1024 or size < item.sourcesize * 0.8:
		logger.error("error downloading file: too small")
		os.remove(target)
		raise TivoException("downloaded file is too small")


def download_decode(item, options, mak):
	item.target = "{}.mpg".format(item.getPath(options))
	try:
		os.makedirs(item.dir)
	except OSError:
		pass
	try:
		download_item(item, mak, item.target)
	except Exception, e:
		exc_info = sys.exc_info()
		try:
			os.remove(item.target)
		except Exception, e2:
			pass
		raise exc_info[1], None, exc_info[2]
	try:
		os.utime(item.target, (item.time, item.time))
	except Exception, e:
		logger.error("Problem setting timestamp: {}".format(e))


def download_one(item, downloaddb, options):
	global config, logger
	logger.info("*** downloading \"{}\": {:.3f} GB".format(item.name, item.sourcesize / 1e9))
	try:
		download_decode(item, options, config.mak)
		downloaddb[item.name] = item.datestr
		if getattr(downloaddb, "sync", None) and callable(downloaddb.sync):
			downloaddb.sync()
		if config.postprocess:
			cmd = config.postprocess
			try:
				cmd = cmd.format(item=item, options=options, config=config)
				r = subprocess.check_output(cmd, shell=True, stderr=subprocess.STDOUT)
				logger.debug("Post-process  {}: {}".format(cmd, r))
			except Exception, e:
				logger.warn("Error running postprocess command '{}' for item {}: {}".format(cmd, item, e))
		logger.debug("Sleeping 30 seconds before moving on...")
		time.sleep(30)
	except TivoException, e:
		logger.info("Error processing \"{}\": {}".format(item.name, e))


def wantitem(item, downloaddb):
	if item.inprogress == "Yes":
		return "recording"
	if item.available == "No":
		return "not available"
	if downloaddb.has_key(item.name):
		return "already downloaded"
	for i in (item.title, item.episode, item.name):
		if IncludeShow.includes.has_key(i):
			return IncludeShow.includes[i]
	return "not included"


def mirror(toc, downloaddb, one=False):
	avail = getAvail(config.targetdir)
	if avail < config.minfree:
		logger.error("{}: {:.1f} GB available, at least {:.1f} GB needed, stopping".format\
			(config.targetdir, avail / config.gig, config.minfree / config.gig))
		sys.exit(1)
	
	items = toc.getItems()
	logger.info("*** {} shows listed".format(len(items)))
	for item in items:
		options = wantitem(item, downloaddb)
		if isinstance(options, basestring):
			logger.debug("*** skipping \"{}\": {}".format(item.name, options))
		else:
			download_one(item, downloaddb, options)
			if one:
				break


def download_episode(toc, downloaddb, episode):
	items = toc.getItems()
	options = {}
	for item in items:
		if item.title == episode or item.name == episode or item.episode == episode:
			for i in (item.title, item.episode, item.name):
				if IncludeShow.includes.has_key(i):
					options = IncludeShow.includes[i]
					download_one(item, downloaddb, options)
					return


def printtoc(toc, downloaddb):
	items = toc.getItems()
	print "*** {} shows listed".format(len(items))
	shows = {}
	for item in items:
		if item.title not in shows:
			shows[item.title] = []
		shows[item.title].append(item)
	for title in sorted(shows):
		for item in sorted(shows[title], key=lambda i: i.name):
			options = wantitem(item, downloaddb)
			if isinstance(options, basestring):
				print "{:>7.7s}: {}".format(options, item.name)
				continue
			print "*** downloading {} ({:.3f} GB)".format(item.name, item.sourcesize / 1e9)
	print "*** {} shows listed".format(len(items))


def usage():
	print >>sys.stderr, 'usage: tivomirror -dvuT [-c config] cmd'
	print >>sys.stderr, '    cmd is one of download, list, mirror, mirrorone'
	sys.exit(64)


def main():
	global config, logger
	curdir = os.getcwd()
	os.chdir(os.path.expanduser("~/.tivo"))
	handler = logging.handlers.RotatingFileHandler("tivomirror.log", maxBytes=2*1024*1024, backupCount=5)
	handler.setFormatter(logging.Formatter(fmt='tivomirror[{}] %(asctime)s %(levelname)6.6s %(message)s'.format(os.getpid()),
		datefmt='%d-%m %H:%M:%S'))
	logger.addHandler(handler)
	downloaddb = anydbm.open("downloads.db", "c")
	toc = TivoToc()
	cmd = "list"
	updateToc = False
	conffile = None

	try:
		options, remainder = getopt.getopt(sys.argv[1:], 'c:dhvuT?',
			['config', 'ignoreepisodetitle', 'debug', 'verbose', 'update', help])

		for opt, arg in options:
			if opt in ('-c', '--config'):
				conffile = arg
			if opt in ('-d', '--debug'):
				logger.setLevel(logging.DEBUG)
			if opt in ('-v', '--verbose'):
				handler = logging.StreamHandler()
				logger.addHandler(handler)
			if opt in ('-u', '--update'):
				updateToc = True
			if opt in ('-T', '--ignoreepisodetitle'):
				config.ignoreepisodetitle = True
			if opt in ('-h', '-?', '--help'):
				usage()

		config = Config(conffile)

		if len(remainder) >= 1:
			cmd = remainder[0]

		if updateToc or cmd == "mirror":
			toc.download()
			toc.save()
		else:
			toc.load()

		if cmd == "mirror":
			mirror(toc, downloaddb)
		elif cmd == "mirrorone":
			mirror(toc, downloaddb, True)
		elif cmd == "list":
			printtoc(toc, downloaddb)
		elif cmd == "download":
			download_episode(toc, downloaddb, remainder[1])
		else:
			logger.error("invalid command {}".format(cmd))
			print >>sys.stderr, "invalid command {}".format(cmd)
			usage()

		downloaddb.close()
	except getopt.GetoptError as e:
		print >>sys.stderr, 'Error parsing options: {}'.format(e)
		usage()
	except Exception:
		logger.exception("")
	logger.info("*** Completed")

if __name__ == "__main__":
	main()