From 91b4e0c202d9fe530c1b4648ac34d8691c8ac50d Mon Sep 17 00:00:00 2001
From: James Bunton <jamesbunton@delx.net.au>
Date: Sat, 13 Oct 2012 00:38:31 +1100
Subject: [PATCH] auto-retry with HTTP ranged gets for command line grabber

---
 youtube.cgi | 79 ++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 60 insertions(+), 19 deletions(-)

diff --git a/youtube.cgi b/youtube.cgi
index 283e692..5e612d4 100755
--- a/youtube.cgi
+++ b/youtube.cgi
@@ -79,14 +79,27 @@ cookiejar = cookielib.CookieJar()
 urlopener = urllib2.build_opener(urllib2.HTTPCookieProcessor(cookiejar))
 referrer = ""
 
-def urlopen(url):
+def urlopen(url, offset=None):
 	global referrer
 	req = urllib2.Request(url)
 	if referrer:
 		req.add_header("Referer", referrer)
 	referrer = url
+
 	req.add_header("User-Agent", USER_AGENT)
-	return urlopener.open(req)
+
+	if offset:
+		req.add_header("Range", "bytes=%d-" % offset)
+
+	res = urlopener.open(req)
+
+	content_range = res.info().getheader("Content-Range")
+	if content_range:
+		tokens = content_range.split()
+		assert tokens[0] == "bytes"
+		start = int(tokens[1].split("-")[0])
+		assert start == offset
+	return res
 
 def parse_url(url):
 	f = urlopen(url)
@@ -215,20 +228,20 @@ def cgimain():
 		)
 		return
 
-def copy_with_progress(total_size, infile, outfile):
-	def pp_size(size):
-		suffixes = ["", "KiB", "MiB", "GiB"]
-		for i, suffix in enumerate(suffixes):
-			if size < 1024:
-				break
-			size /= 1024
-		return "%.2f %s" % (size, suffix)
+def pp_size(size):
+	suffixes = ["", "KiB", "MiB", "GiB"]
+	for i, suffix in enumerate(suffixes):
+		if size < 1024:
+			break
+		size /= 1024
+	return "%.2f %s" % (size, suffix)
 
+def copy_with_progress(content_length, infile, outfile):
 	def print_status():
 		sys.stdout.write("\33[2K\r")
 		sys.stdout.write("%s / %s (%s/sec)" % (
 			pp_size(bytes_read),
-			pp_size(total_size),
+			pp_size(content_length),
 			pp_size(bytes_read / (now - start_ts)),
 		))
 		sys.stdout.flush()
@@ -258,17 +271,45 @@ def main():
 	except:
 		print >>sys.stderr, "Usage: %s http://youtube.com/watch?v=FOOBAR" % sys.argv[0]
 		sys.exit(1)
+
 	doc = parse_url(url)
 	video_url, filename = get_video_url(doc)
-	video_data = urlopen(video_url)
-	if os.path.isfile(filename):
-		print >>sys.stderr, "Error! File exists:", filename
-		sys.exit(1)
-	outfile = open(filename, "w")
-	total_size = int(video_data.info().getheader("Content-Length"))
 	print "Downloading", filename.encode("utf-8")
-	copy_with_progress(total_size, video_data, outfile)
-	video_data.close()
+
+	outfile = open(filename, "a")
+	offset = outfile.tell()
+	if offset > 0:
+		print "Resuming download from", pp_size(offset)
+	total_size = None
+
+	while True:
+		try:
+			video_data = urlopen(video_url, offset)
+		except urllib2.HTTPError, e:
+			if e.code == 416:
+				print "File is complete!"
+				break
+			else:
+				raise
+
+		content_length = int(video_data.info().getheader("Content-Length"))
+		if total_size is None:
+			total_size = content_length
+
+		try:
+			copy_with_progress(content_length, video_data, outfile)
+		except IOError:
+			print
+
+		video_data.close()
+		if outfile.tell() != total_size:
+			old_offset = offset
+			offset = outfile.tell()
+			if old_offset == offset:
+				time.sleep(1)
+			print "Restarting download from", pp_size(offset)
+			continue
+
 	outfile.close()
 
 
-- 
2.39.2