Add fallback for urllib3 versions somehow lacking drain_conn, fixes #35 and #36

waybackproxy.py

@@ -49,7 +49,7 @@ class Handler(socketserver.BaseRequestHandler):
 
 		# readline is pretty convenient
 		f = self.request.makefile()
-		
+
 		# read request line
 		reqline = line = f.readline()
 		split = line.rstrip().split(' ')
@@ -207,7 +207,7 @@ class Handler(socketserver.BaseRequestHandler):
 				# Check for redirects.
 				destination = conn.get_redirect_location()
 				if destination:
-					conn.drain_conn()
+					self.drain_conn(conn)
 					conn.release_conn()
 
 					# Check if the redirect goes to a different Wayback URL.
@@ -245,7 +245,7 @@ class Handler(socketserver.BaseRequestHandler):
 				if 'javascript' in guessed_content_type:
 					match = re.match('''(https?://web\\.archive\\.org/web/[0-9]+)([^/]*)(.+)''', request_url)
 					if match and match.group(2) != 'im_':
-						conn.drain_conn()
+						self.drain_conn(conn)
 						conn.release_conn()
 						request_url = match.group(1) + 'im_' + match.group(3)
 						continue
@@ -265,11 +265,11 @@ class Handler(socketserver.BaseRequestHandler):
 		if conn.status != 200:
 			if conn.status in (403, 404): # not found
 				if self.guess_and_send_redirect(http_version, archived_url):
-					conn.drain_conn()
+					self.drain_conn(conn)
 					conn.release_conn()
 					return
 			#elif conn.status in (301, 302): # redirect loop detection currently unused
-			#	conn.drain_conn()
+			#	self.drain_conn(conn)
 			#	conn.release_conn()
 			#	return self.send_error_page(http_version, 508, 'Infinite Redirect Loop')
 
@@ -279,7 +279,7 @@ class Handler(socketserver.BaseRequestHandler):
 			# If the memento Link header is present, this is a website error
 			# instead of a Wayback error. Pass it along if that's the case.
 			if 'Link' not in conn.headers:
-				conn.drain_conn()
+				self.drain_conn(conn)
 				conn.release_conn()
 				return self.send_error_page(http_version, conn.status, conn.reason)
 
@@ -305,7 +305,7 @@ class Handler(socketserver.BaseRequestHandler):
 			# portion of the URL away if it ends up being HTML consumed
 			# through the QUICK_IMAGES interface.
 			if hostname == 'web.archive.org':
-				conn.drain_conn()
+				self.drain_conn(conn)
 				conn.release_conn()
 				archived_url = '/'.join(request_url.split('/')[5:])
 				_print('[r] [QI]', archived_url)
@@ -317,7 +317,7 @@ class Handler(socketserver.BaseRequestHandler):
 				if match:
 					requested_date = match.group(1)
 					if self.wayback_to_datetime(requested_date) > self.wayback_to_datetime(original_date) + datetime.timedelta(int(DATE_TOLERANCE)):
-						conn.drain_conn()
+						self.drain_conn(conn)
 						conn.release_conn()
 						_print('[!]', requested_date, 'is outside the configured tolerance of', DATE_TOLERANCE, 'days')
 						if not self.guess_and_send_redirect(http_version, archived_url):
@@ -350,14 +350,14 @@ class Handler(socketserver.BaseRequestHandler):
 						# Start fetching the URL.
 						_print('[f]', archived_url)
 						conn = self.shared_state.http.urlopen('GET', request_url, retries=retry, preload_content=False)
-						
+
 						if conn.status != 200:
 							_print('[!]', conn.status, conn.reason)
 
 							# If the memento Link header is present, this is a website error
 							# instead of a Wayback error. Pass it along if that's the case.
 							if 'Link' not in conn.headers:
-								conn.drain_conn()
+								self.drain_conn(conn)
 								conn.release_conn()
 								return self.send_error_page(http_version, conn.status, conn.reason)
 
@@ -496,7 +496,7 @@ class Handler(socketserver.BaseRequestHandler):
 		# Finish and send the request.
 		response += '\r\n\r\n'
 		self.request.sendall(response.encode('utf8', 'ignore'))
-	
+
 	def send_error_page(self, http_version, code, reason):
 		"""Generate an error page."""
 
@@ -574,7 +574,7 @@ class Handler(socketserver.BaseRequestHandler):
 			self.send_redirect_page(http_version, new_url)
 			return True
 		return False
-	
+
 	def handle_settings(self, query):
 		"""Generate the settings page."""
 
@@ -644,6 +644,9 @@ class Handler(socketserver.BaseRequestHandler):
 				fmt = fmt[:-2]
 				fmt_len -= 2
 
+	def drain_conn(self, conn):
+		getattr(conn, 'drain_conn', conn.read)()
+
 print_lock = threading.Lock()
 def _print(*args, **kwargs):
 	"""Logging function."""