#!/usr/bin/env python3

# uget-chrome-wrapper is a tool to integrate uGet Download manager
# with Google Chrome, Chromium, Vivaldi and Opera in Linux and Windows.

# Copyright (C) 2016  Gobinath

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# sudo apt install python3-urllib3
import struct, sys, threading, logging, json, urllib, subprocess, tempfile
from urllib.parse import urlparse
from urllib.parse import unquote
from urllib.request import urlopen
from os.path import splitext, basename, join, expanduser
from mimetypes import guess_extension

cookie_filepath = join(tempfile.gettempdir(), 'uget_cookie')
urls_filepath = join(tempfile.gettempdir(), 'uget_urls')
UGET_COMMAND = "uget-gtk"
VERSION = "2.0.6"

logger = logging.getLogger()
# log_file_path = join(expanduser('~'), 'uget-chrome-wrapper.log')
# logging.basicConfig(format='%(asctime)s [%(levelname)s]: %(message)s', filename=log_file_path, filemode='a', level=logging.DEBUG)
logger.propagate = False

# Platform specific configuration
if sys.platform == 'win32':
	# Set the default I/O mode to O_BINARY in windows
	import os, msvcrt
	msvcrt.setmode(sys.stdin.fileno(), os.O_BINARY)
	msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY)
	UGET_COMMAND = 'uget'

def extract_file_name(url):
	fileName = ''
	if 'googlevideo.com/' in url and '&title=' in url:
		# Youtube video
		logger.debug('Found Youtube video')
		url = urllib.parse.unquote(url)
		url_components = url.split('&')
		for smt in url_components:
			if smt.startswith('title='):
				title = smt.replace('title=', '').replace('+', ' ')
			elif smt.startswith('mime='):
				mime = smt.replace('mime=', '')

		if title:
			fileName = title
			if mime:
				fileName += guess_extension(mime)
			logger.debug('Youtube filename: ' + fileName)
	else:
		fileName = basename(urlparse(url).path).replace('%20', ' ')

	return fileName

# Send a message to the webapp.
def send_message(message):

	logger.info('Sending message: ' + str(message))
	try:
		# Write message size.
		sys.stdout.buffer.write(struct.pack('I', len(message)))
		# Write the message itself.
		sys.stdout.write(message)
		sys.stdout.flush()
	except Exception as e:
		logger.error('Error in sending message: ' + str(e))

def download_all_links(data):
	"""
	Extract all links from the html page and download them.
	"""
	urls = data['url']
	cookie = data['cookies']
	try:
		logger.debug('Writing links to file')
		with open(urls_filepath, 'w') as url_file:
			url_file.write(urls)
		
		command = [UGET_COMMAND, "--input-file=" + urls_filepath]
		use_cookie_file = False
		if cookie:
			try:
				with open(cookie_filepath, 'w') as cookie_file:
					cookie_file.write(cookie)
				use_cookie_file = True
			except Exception as e:
				pass
		
		if use_cookie_file:
			command.append("--http-cookie-file=" + cookie_filepath)
		# Pass the parameters to uGet
		subprocess.call(command)

		if use_cookie_file:
			try:
				os.remove(cookie_filepath)
			except Exception as e:
				pass
		
		try:
			os.remove(urls_filepath)
		except Exception as e:
			pass
	except Exception as e:
		pass

# Read messages from the webapp.
def read_message():

	logger.info('uget-chrome-wrapper is reading the message')

	while 1:
		# Read the message length (first 4 bytes).
		text_length_bytes = sys.stdin.buffer.read(4)

		# Unpack message length as 4 byte integer.
		text_length = struct.unpack('i', text_length_bytes)[0]

		logger.debug('Message length: ' + str(text_length))

		# Read the text (JSON object) of the message.
		text = sys.stdin.buffer.read(text_length).decode('utf-8')

		logger.debug('Received message: ' + str(text))

		if text:
			if not 'url' in text:
				uget_version = subprocess.Popen([UGET_COMMAND + " --version"], shell=True, stdout=subprocess.PIPE, universal_newlines=True).communicate()[0]
				uget_version = uget_version.replace('uGet', '').replace('for GTK+', '').strip()
				logger.debug('uGet version: ' + uget_version)
				send_message('{"enable": true, "version": "' + VERSION + '", "uget": "' + uget_version + '"}')
				return

			send_message('{"state": "accepted"}')
			data = json.loads(text)
			url = data['url']

			if url != "":
				if data['batch']:
					download_all_links(data)
					return
				use_cookie_file = False
				filename = data['filename']
				cookie = data['cookies']
				referer = data['referrer']
				command = [UGET_COMMAND]
				if cookie:
					try:
						with open(cookie_filepath, 'w') as cookie_file:
							cookie_file.write(cookie)
						use_cookie_file = True
					except Exception as e:
						pass
				if not filename:
					filename = extract_file_name(url)
				else:
					# Sometimes Firefox sends complete path
					filename = basename(filename)

				# Add the referer url
				if referer:
					command.append("--http-referer=" + referer)

				# Add the file name
				if filename:
					command.append("--filename=" + unquote(filename))

				# Add the cookie file
				if use_cookie_file:
					command.append("--http-cookie-file=" + cookie_filepath)

				# Add the url
				command.append(url)

				logger.debug('Execute command: ' + str(command))
				# Pass the parameters to uGet
				subprocess.call(command)

				if use_cookie_file:
					try:
						os.remove(cookie_filepath)
					except Exception as e:
						pass

			sys.exit(0)


if __name__ == '__main__':
	read_message()
