/rip.py

#!/usr/bin/env python

import sys
import commands
import getopt
import re
import string
import math
import os
import dbus


def main():
	os.nice(15)
	settings = {}
	settings["title"] = "0"
	settings["angle"] = 1
	settings["path"] = "/dev/dvd"
	settings["disc_title"] = ""
	settings["chapter"] = 1
	settings["output_directory"] = "~/dvd/<disc_title>/<title>"
	settings["nosubs"] = False
	settings = read_args(settings)
	got_title_info = False

	if( not os.path.exists(settings["path"]) ) :
		usage_and_exit("no dvd found at path: " + settings["path"])
	
	title_info = {}
	if settings["disc_title"] != "" :
		title_info["disc_title"] = settings["disc_title"]
	else :
		title_info = get_title_info(settings["path"])
		got_title_info = True

	print "Disc Title: " + title_info["disc_title"]


	output_dir = settings["output_directory"].replace("<disc_title>", title_info["disc_title"])
	output_dir = os.path.normpath(os.path.expanduser(output_dir))
	suspend = SuspendInhibit("rip", "Ripping a DVD", SuspendInhibit.SUSPENDING)
	
	try :	
		if(settings["title"] == "0") :
			if not got_title_info :
				title_info = get_title_info(settings["path"])
				got_title_info = True
			
			title = parse_int(title_info["longest_title"])
		
			output_dir = output_dir.replace("<title>", str(title).zfill(2))
		
			grab_title(title, settings, title_info["disc_title"], output_dir)
		else :
			have_valid_title = False
		
			for title in settings["title"].split(",") :
				title = parse_int(title)

				output_dir1 = output_dir.replace("<title>", str(title).zfill(2))
			
				if got_title_info :
					if( not title in title_info["title_numbers"] ) :
						continue
				have_valid_title = True

				print "Title: " + str(title)
			
				grab_title(title, settings, title_info["disc_title"], output_dir1)
			
			if( not have_valid_title ) :
				usage_and_exit("You must specify at least one valid title.")
				
		print "Complete."
		os.system("eject " + settings["path"])
	finally :
		suspend.uninhibit()	
	

def save_chapters(title_num, disc_path, output_directory) :
	print "saving chapters to directory: " + output_directory
	if not os.path.exists(output_directory) :
		os.makedirs(output_directory)
	output_file = os.path.join(output_directory, "chapters.txt")
	commands.getoutput("dvdxchap -t %d %s > %s" % (title_num, disc_path, output_file))
	

class SuspendInhibit :
	LOGGING_OUT = 1
	USER_SWITCHING = 2
	SUSPENDING = 4
	IDLE_SESSION = 8
	
	def __init__(self, program, activity, flags=SUSPENDING) :
		self.dev = None
		try:
			bus = dbus.Bus(dbus.Bus.TYPE_SESSION)
			#last worked on jaunty
			#devobj = bus.get_object('org.freedesktop.PowerManagement', '/org/freedesktop/PowerManagement/Inhibit') 
			#self.dev = dbus.Interface(devobj, "org.freedesktop.PowerManagement.Inhibit")
			devobj = bus.get_object('org.gnome.SessionManager', '/org/gnome/SessionManager')
			self.dev = dbus.Interface(devobj, 'org.gnome.SessionManager')
			self.cookie = self.dev.Inhibit(program, 0, activity, flags)
		except:
			print "Can't inhibit suspend"

	def uninhibit(self) :
		if self.dev != None:
			self.dev.Uninhibit(self.cookie)


class TitleInfo :
	_file = ""
	_subtitles = None
	_text = None
	_audio = None
	title = 0
	working_directory = ""
	disc_path = ""
	
	def __init__(self, working_directory, disc_path = None, title_num = None) :
		self.working_directory = working_directory
		if not os.path.exists(working_directory) :
			os.makedirs(working_directory)
		self.title = title_num
		self.disc_path = disc_path
		self._file = os.path.join(working_directory, "disc_info.txt")
		
	def save_info(self) :
		self.get_text()
	
	def get_text(self) :
		if self._text is not None :
			return self._text
		elif os.path.exists(self._file) :
			return read_text_file(self._file)
		else :
			if self.title == None:
				print "No disc title and file not found at", self._file
				sys.exit(1)
			self._text = commands.getoutput(
				"mplayer dvd://%d -dvd-device %s -vo dummy -ao dummy -identify 2>/dev/null" 
				% (self.title, self.disc_path))
			write_text_file(self._file, self._text)
			return self._text
	
	def get_subtitles(self, lang = None) :
		if self._subtitles == None :
			matches = re.findall(r"subtitle\b.+(\d+) language: (\w+)", self.get_text())
			self._subtitles = [{"id" : int(m[0]), "lang" : m[1].strip().lower()} for m in matches]
		if lang is not None :
			ret = []
			for sub in self._subtitles :
				if sub["lang"] in lang :
					ret.append(sub)
			return ret
		return self._subtitles or []

	def get_audio_by_id(self, id) :
		id = int(id)
		for track in self.get_audio() :
			if int(track["id"]) == id :
				return track
		raise Exception, "track %d not found" % id
	
	def get_video_seconds(self) :
		m = re.search(r"^ID_LENGTH=(\d+)", self.get_text(), re.M)
		if m :
			return int(m.group(1))
		else :
			raise Exception, "Couldn't determine the number of seconds for the video."
		
	def get_default_audio_id(self) :
		pattern = r"^ID_AUDIO_ID=(\d+)"
		matches = re.findall(pattern, self.get_text(), re.M)
		
		print "default audio: ", matches[-1]
		
		return int(matches[-1])
		
	def get_audio(self, lang = None) :
		if self._audio == None :
			pattern = r"audio stream: (\d+) format: .+ \((.+)\) language: (\w+) aid: (\d+)."
			matches = re.findall(pattern, self.get_text())
			
			self._audio = map(
				lambda m :
					{"stream" : int(m[0]), "format" : m[1], "lang" : m[2], "id" : m[3]}
				, matches)
			
		id = self.get_default_audio_id()
		
		for track in self._audio :
			if track["id"] == id :
				yield track
				break
		
		for track in self._audio :
			if (lang is None or track["lang"] in lang) and track["id"] != id :
				yield track
				
		retval = self._audio
		
		"""
		if lang is not None :
			retval = filter(
				lambda track :
					track["lang"] in lang
				, retval)
		return retval
		"""


def save_subtitles(title_info, chapter, lang) :
	print "saving subtitles to directory: " + title_info.working_directory
	subtitles = title_info.get_subtitles(lang)
	if len(subtitles) == 0 :
		subtitles = title_info.get_subtitles()
	for sub in subtitles :
		output_file = os.path.join(title_info.working_directory, "subs%s" % sub["id"])
		cmd = "mencoder dvd://%d -chapter %s -dvd-device %s -nosound -ovc frameno -o /dev/null -sid %s -vobsubout %s" 
		cmd = cmd % (title_info.title, chapter, title_info.disc_path, sub["id"], output_file)
		commands.getoutput(cmd)

def save_audio(title_info, lang) :
	print "saving audio to directory: " + title_info.working_directory

	audio_tracks = list(title_info.get_audio(lang))
	try:
		if not any(track["id"] == 128 for track in audio_tracks):
			#When 128 doesn't exist as a track, an error is thrown, so nothing's added.
			audio_tracks.insert(0, title_info.get_audio_by_id(128))
	except: pass
	
	if len(audio_tracks) == 0 :
		audio_tracks = list(title_info.get_audio())
	if len(audio_tracks) == 0 :
		raise Exception, "still found zero audio tracks!"
	
	output_file = os.path.join(title_info.working_directory, "audiostreamlist")
	streamlist = [track["id"] for track in audio_tracks]
	write_text_file(output_file, "\n".join(streamlist))
	

def grab_title(title_num, settings, disc_title, output_directory) :
	disc_path = settings["path"]
	angle = settings["angle"]
	chapter = settings["chapter"]
	
	print "ripping title to directory: " + output_directory
	if not os.path.exists(output_directory) :
		os.makedirs(output_directory)
		
	output_file = os.path.join(output_directory, "%s-%d-%d.vob" % (disc_title, title_num, angle))
	
	LANG = ["en", "unknown", "eng"]
	info = TitleInfo(output_directory, disc_path, title_num)
	
	info.save_info()

	save_audio(info, LANG)
	
	save_chapters(title_num, disc_path, output_directory)

	cmd = "mplayer dvd://%d -v -chapter %s -dumpstream -dumpfile '%s' -dvd-device %s 2> /dev/null" \
		% (title_num, chapter, output_file, disc_path)
	print cmd

	pipe = os.popen(cmd)
	while 1:
		line = pipe.readline()
		if not line:
			break
		if line.find("DVD next cell") >= 0 :
			print line,
			

	if not settings["nosubs"] :
		save_subtitles(info, chapter, LANG)


	
def write_text_file(path, to_write) :
	f = open(path, 'w')
	f.write(to_write)
	f.close()
	
def read_text_file(path) :
	if os.path.isfile(path) :
		f=open(path, 'r')
		line = f.read()
		f.close()
		return line
	else:
		return ""

def get_title_info(dvd_path) :
	output = commands.getoutput("lsdvd " + dvd_path)
	disc_title = re.search("Disc Title: (\w+)", output, re.I).group(1).lower()
	re_title = re.compile("Title: (\d+),", re.I | re.M)
	title_numbers = re_title.findall(output)
	title_numbers = [parse_int(title_num) for title_num in title_numbers]
	longest_title = re.search(r"Longest track: (\d+)", output, re.I).group(1)
	
	return {"disc_title": disc_title,  "longest_title": longest_title, "title_numbers": title_numbers}
	
	
def parse_int(s) :
	m = re.search(r"-?\d+", s)
	if( not m ) :
		return 0
	return int(m.group())
	
	

def read_args(settings) :
	try:
		opts, args = getopt.gnu_getopt(sys.argv[1:], "?c:t:l:a:d:o",
			["help",
			"title=",
			"angle=",
			"chapter=",
			"dvd-path=",
			"output-directory=",
			"disc-title=",
			"nosubs",
			])
	except getopt.GetoptError, msg:
		usage_and_exit(msg)
		
	for opt, arg in opts:
		if opt in ("-?", "--help"):
			usage_and_exit()
		elif opt in ("-t", "--title"):
			settings["title"] = arg
		elif opt in ("-a", "--angle"):
			settings["angle"] = arg
		elif opt in ("-c", "--chapter"):
			settings["chapter"] = arg
		elif opt in ("-d", "--dvd-path"):
			settings["path"] = arg
		elif opt in ("-o", "--output-directory"):
			settings["output_directory"] = arg
		elif opt in ("-l", "--disc-title"):
			settings["disc_title"] = arg
		elif opt in ("--nosubs") :
			settings["nosubs"] = True
		else :
			print "unknown arg: " + opt

	if len(args) > 0 :
		usage_and_exit("unknown arguments: " + str(args))

	return settings

	
def usage_and_exit(errmsg):
	"""Print a usage message, plus an ERRMSG (if provided), then exit.
	If ERRMSG is provided, the usage message is printed to stderr and
	the script exits with a non-zero error code.  Otherwise, the usage
	message goes to stdout, and the script exits with a zero
	errorcode."""
	if errmsg is None:
		stream = sys.stdout
	else:
		stream = sys.stderr
	print >> stream, __doc__
	if errmsg:
		print >> stream, "\nError: %s" % (errmsg)
		sys.exit(2)
	sys.exit(0)

	
if __name__ == "__main__":
	main()