#!/usr/bin/env python

import gettext
gettext.install('torbrowser-launcher', '/usr/share/torbrowser-launcher/locale')

from twisted.internet import gtk2reactor
gtk2reactor.install()
from twisted.internet import reactor

import pygtk
pygtk.require('2.0')
import gtk

import os, sys, subprocess, locale, urllib2, gobject, time, pickle, json, tarfile, psutil

from twisted.web.client import Agent, ResponseDone
from twisted.web.http_headers import Headers
from twisted.internet.protocol import Protocol
from twisted.internet.ssl import ClientContextFactory

from OpenSSL.SSL import Context, VERIFY_PEER, VERIFY_FAIL_IF_NO_PEER_CERT
from OpenSSL.crypto import load_certificate, FILETYPE_PEM

class VerifyTorProjectCert(ClientContextFactory):

    def __init__(self, torproject_pem):
        self.torproject_ca = load_certificate(FILETYPE_PEM, open(torproject_pem, 'r').read())

    def getContext(self, host, port):
        ctx = ClientContextFactory.getContext(self)
        ctx.set_verify_depth(0)
        ctx.set_verify(VERIFY_PEER | VERIFY_FAIL_IF_NO_PEER_CERT, self.verifyHostname)
        return ctx

    def verifyHostname(self, connection, cert, errno, depth, preverifyOK):
        return cert.digest('sha256') == self.torproject_ca.digest('sha256')


class TorBrowserLauncher:
    def __init__(self):
        # initialize the app
        self.set_gui(None, '', [])
        self.discover_arch_lang()
        self.build_paths()
        self.mkdir(self.paths['dir']['download'])
        self.mkdir(self.paths['dir']['tbb'])
        self.init_gnupg()

        # allow buttons to have icons
        try:
            settings = gtk.settings_get_default()
            settings.props.gtk_button_images = True
        except:
            pass

        self.launch_gui = True

        # if we haven't already hit an error
        if self.gui != 'error':
            # load settings
            if self.load_settings():
                self.build_paths(self.settings['latest_version'])

                # is tbb already running and we just need to open a new firefox?
                if self.settings['installed_version']:
                    vidalia_pid = None
                    firefox_pid = None
                    for p in psutil.process_iter():
                        try:
                            exe = None

                            # old versions of psutil don't have exe
                            if hasattr(p, 'exe'):
                                exe = p.exe
                            # need to rely on cmdline instead
                            else:
                                if len(p.cmdline) > 0:
                                    exe = p.cmdline[0]
                            
                            if exe == self.paths['file']['vidalia_bin'] or exe == './App/vidalia':
                                vidalia_pid = p.pid
                            if exe == self.paths['file']['firefox_bin']:
                                firefox_pid = p.pid

                        except:
                            pass

                    if vidalia_pid and not firefox_pid:
                        print _('Vidalia is already open, but Firefox is closed. Launching new Firefox.')
                        subprocess.Popen([self.paths['file']['firefox_bin'], '-no-remote', '-profile', self.paths['file']['firefox_profile']])
                        return
                    elif vidalia_pid and firefox_pid:
                        print _('Vidalia and Firefox are already open, bringing them to focus')

                        # figure out the window ids of vidalia and firefox
                        vidalia_win_id = None
                        firefox_win_id = None
                        p = subprocess.Popen(['wmctrl', '-l', '-p'], stdout=subprocess.PIPE)
                        for line in p.stdout.readlines():
                            line_split = line.split()
                            win_id = line_split[0]
                            win_pid = int(line_split[2])
                            if win_pid == vidalia_pid:
                                vidalia_win_id = win_id
                            if win_pid == firefox_pid:
                                firefox_win_id = win_id

                        # bring firefox to front, then vidalia
                        if firefox_win_id:
                            subprocess.call(['wmctrl', '-i', '-a', firefox_win_id])
                        if vidalia_win_id:
                            subprocess.call(['wmctrl', '-i', '-a', vidalia_win_id])

                        return

                # how long was it since the last update check?
                # 86400 seconds = 24 hours
                current_timestamp = int(time.time())
                if current_timestamp - self.settings['last_update_check_timestamp'] >= 86400:
                    # check for update
                    print 'Checking for update'
                    self.set_gui('task', _("Checking for Tor Browser update."), 
                        ['download_update_check', 
                         'attempt_update'])

                else:
                    # no need to check for update
                    print _('Checked for update within 24 hours, skipping')
                    self.start_launcher()

            else:
                self.set_gui('error', _("Error loading settings. Delete ~/.torbrowser and try again."), [])

        if self.launch_gui:
            # set up the window
            self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
            self.window.set_title(_("Tor Browser"))
            self.window.set_icon_from_file(self.paths['file']['icon'])
            self.window.set_position(gtk.WIN_POS_CENTER)
            self.window.set_border_width(10)
            self.window.connect("delete_event", self.delete_event)
            self.window.connect("destroy", self.destroy)

            # build the rest of the UI
            self.build_ui()

    # download or run TBB
    def start_launcher(self):
      # is TBB already installed?
      if os.path.isfile(self.paths['file']['start']) and os.access(self.paths['file']['start'], os.X_OK):
        if self.settings['installed_version'] == self.settings['latest_version']:
          # current version of tbb is installed, launch it
          self.run(False)
          self.launch_gui = False
        elif self.settings['installed_version'] < self.settings['latest_version']:
          # there is a tbb upgrade available
          self.set_gui('task', _("Your Tor Browser is out of date."), 
            ['download_tarball', 
             'download_tarball_sig', 
             'verify', 
             'extract', 
             'run'])
        else:
          # for some reason the installed tbb is newer than the current version?
          self.set_gui('error', _("Something is wrong. The version of Tor Browser Bundle you have installed is newer than the current version?"), [])

      # not installed
      else:
          # are the tarball and sig already downloaded?
          if os.path.isfile(self.paths['file']['tarball']) and os.path.isfile(self.paths['file']['tarball_sig']):
              # start the gui with verify
              self.set_gui('task', _("Installing Tor Browser."), 
                  ['verify', 
                   'extract', 
                   'run'])

          # first run
          else:
              self.set_gui('task', _("Downloading and installing Tor Browser."), 
                  ['download_tarball', 
                   'download_tarball_sig', 
                   'verify', 
                   'extract', 
                   'run'])
   
    # discover the architecture and language
    def discover_arch_lang(self):
        # figure out the architecture
        (sysname, nodename, release, version, machine) = os.uname()
        self.architecture = machine

        # figure out the language
        available_languages = ['en-US', 'ar', 'de', 'es-ES', 'fa', 'fr', 'it', 'ko', 'nl', 'pl', 'pt-PT', 'ru', 'vi', 'zh-CN']
        default_locale = locale.getdefaultlocale()[0]
        if default_locale == None:
            self.language = 'en-US'
        else:
            self.language = default_locale.replace('_', '-')
            if self.language not in available_languages:
                self.language = self.language.split('-')[0]
                if self.language not in available_languages:
                    for l in available_languages:
                        if l[0:2] == self.language:
                            self.language = l
            # if language isn't available, default to english
            if self.language not in available_languages:
                self.language = 'en-US'

    # build all relevant paths
    def build_paths(self, tbb_version = None):
        homedir = os.getenv('HOME')
        if not homedir:
            homedir = '/tmp/.torbrowser-'+os.getenv('USER')
            if os.path.exists(homedir) == False:
                try:
                    os.mkdir(homedir, 0700)
                except:
                    self.set_gui('error', _("Error creating {0}").format(homedir), [], False)
        if not os.access(homedir, os.W_OK):
            self.set_gui('error', _("{0} is not writable").format(homedir), [], False)

        tbb_data = '%s/.torbrowser' % homedir

        if tbb_version:
            tarball_filename = 'tor-browser-gnu-linux-'+self.architecture+'-'+tbb_version+'-dev-'+self.language+'.tar.gz'
            self.paths['file']['tarball'] = tbb_data+'/download/'+tarball_filename
            self.paths['file']['tarball_sig'] = tbb_data+'/download/'+tarball_filename+'.asc'
            self.paths['url']['tarball'] = 'https://www.torproject.org/dist/torbrowser/linux/'+tarball_filename
            self.paths['url']['tarball_sig'] = 'https://www.torproject.org/dist/torbrowser/linux/'+tarball_filename+'.asc'
            self.paths['filename']['tarball'] = tarball_filename
            self.paths['filename']['tarball_sig'] = tarball_filename+'.asc'

        else:
            self.paths = {
                'dir': {
                    'data': tbb_data,
                    'download': tbb_data+'/download',
                    'tbb': tbb_data+'/tbb/'+self.architecture,
                    'gnupg_homedir': tbb_data+'/gnupg_homedir'
                },
                'file': {
                    'settings': tbb_data+'/settings',
                    'version': tbb_data+'/version',
                    'start': tbb_data+'/tbb/'+self.architecture+'/tor-browser_'+self.language+'/start-tor-browser',
                    'vidalia_bin': tbb_data+'/tbb/'+self.architecture+'/tor-browser_'+self.language+'/App/vidalia',
                    'firefox_bin': tbb_data+'/tbb/'+self.architecture+'/tor-browser_'+self.language+'/App/Firefox/firefox',
                    'firefox_profile': tbb_data+'/tbb/'+self.architecture+'/tor-browser_'+self.language+'/Data/profile',
                    'update_check': tbb_data+'/download/RecommendedTBBVersions',
                    'icon': '/usr/share/pixmaps/torbrowser80.xpm',
                    'torproject_pem': '/usr/share/torbrowser-launcher/torproject.pem',
                    'erinn_key': '/usr/share/torbrowser-launcher/erinn.asc',
                    'sebastian_key': '/usr/share/torbrowser-launcher/sebastian.asc'
                },
                'url': {
                    'update_check': 'https://check.torproject.org/RecommendedTBBVersions'
                },
                'filename': {}
            }

    # create a directory
    def mkdir(self, path):
        try:
            if os.path.exists(path) == False:
                os.makedirs(path, 0700)
                return True
        except:
            self.set_gui('error', _("Cannot create directory {0}").format(path), [], False)
            return False
        if not os.access(path, os.W_OK):
            self.set_gui('error', _("{0} is not writable").format(path), [], False)
            return False
        return True

    # if gnupg_homedir isn't set up, set it up
    def init_gnupg(self):
        if not os.path.exists(self.paths['dir']['gnupg_homedir']):
            print _('Creating GnuPG homedir'), self.paths['dir']['gnupg_homedir']
            if self.mkdir(self.paths['dir']['gnupg_homedir']):
                # import keys
                print _('Importing keys')
                p1 = subprocess.Popen(['/usr/bin/gpg', '--homedir', self.paths['dir']['gnupg_homedir'], '--import', self.paths['file']['erinn_key']])
                p2 = subprocess.Popen(['/usr/bin/gpg', '--homedir', self.paths['dir']['gnupg_homedir'], '--import', self.paths['file']['sebastian_key']])
                # wait for keys to import before moving on
                p1.wait()
                p2.wait()

    # there are different GUIs that might appear, this sets which one we want
    def set_gui(self, gui, message, tasks, autostart=True):
        self.gui = gui
        self.gui_message = message
        self.gui_tasks = tasks
        self.gui_task_i = 0
        self.gui_autostart = autostart

    # set all gtk variables to False
    def clear_ui(self):
        if hasattr(self, 'box'):
            self.box.destroy()
        self.box = False

        self.label = False
        self.progressbar = False
        self.button_box = False
        self.start_button = False
        self.exit_button = False

    # build the application's UI
    def build_ui(self):
        self.box = gtk.VBox(False, 20)
        self.window.add(self.box)

        if self.gui == 'error':
            # labels
            self.label = gtk.Label( self.gui_message ) 
            self.label.set_line_wrap(True)
            self.box.pack_start(self.label, True, True, 0)
            self.label.show()

            #self.label2 = gtk.Label("You can fix the problem by deleting:\n"+self.paths['dir']['data']+"\n\nHowever, you will lose all your bookmarks and other Tor Browser preferences.") 
            #self.label2.set_line_wrap(True)
            #self.box.pack_start(self.label2, True, True, 0)
            #self.label2.show()

            # exit button
            exit_image = gtk.Image()
            exit_image.set_from_stock(gtk.STOCK_CANCEL, gtk.ICON_SIZE_BUTTON)
            self.exit_button = gtk.Button("Exit")
            self.exit_button.set_image(exit_image)
            self.exit_button.connect("clicked", self.destroy, None)
            self.box.add(self.exit_button)
            self.exit_button.show()

        elif self.gui == 'task':
            # label
            self.label = gtk.Label( self.gui_message ) 
            self.label.set_line_wrap(True)
            self.box.pack_start(self.label, True, True, 0)
            self.label.show()
            
            # progress bar
            self.progressbar = gtk.ProgressBar(adjustment=None)
            self.progressbar.set_orientation(gtk.PROGRESS_LEFT_TO_RIGHT)
            self.progressbar.set_pulse_step(0.01)
            self.box.pack_start(self.progressbar, True, True, 0)

            # button box
            self.button_box = gtk.HButtonBox()
            self.button_box.set_layout(gtk.BUTTONBOX_SPREAD)
            self.box.pack_start(self.button_box, True, True, 0)
            self.button_box.show()

            # start button
            start_image = gtk.Image()
            start_image.set_from_stock(gtk.STOCK_APPLY, gtk.ICON_SIZE_BUTTON)
            self.start_button = gtk.Button("Start")
            self.start_button.set_image(start_image)
            self.start_button.connect("clicked", self.start, None)
            self.button_box.add(self.start_button)
            if not self.gui_autostart:
              self.start_button.show()

            # exit button
            exit_image = gtk.Image()
            exit_image.set_from_stock(gtk.STOCK_CANCEL, gtk.ICON_SIZE_BUTTON)
            self.exit_button = gtk.Button("Exit")
            self.exit_button.set_image(exit_image)
            self.exit_button.connect("clicked", self.destroy, None)
            self.button_box.add(self.exit_button)
            self.exit_button.show()

        self.box.show()
        self.window.show()

        if self.gui_autostart:
            self.start(None)

    # start button clicked, begin tasks
    def start(self, widget, data=None):
        # disable the start button
        if self.start_button:
            self.start_button.set_sensitive(False)

        # start running tasks
        self.run_task()
      
    # run the next task in the task list
    def run_task(self):
        self.refresh_gtk()

        if self.gui_task_i >= len(self.gui_tasks):
            self.destroy(False)
            return

        task = self.gui_tasks[self.gui_task_i]
        
        # get ready for the next task
        self.gui_task_i += 1

        if task == 'download_update_check':
            print _('Downloading'), self.paths['url']['update_check']
            self.download('update check', self.paths['url']['update_check'], self.paths['file']['update_check'])
        
        if task == 'attempt_update':
            print _('Checking to see if update it needed')
            self.attempt_update()

        elif task == 'download_tarball':
            print _('Downloading'), self.paths['url']['tarball']
            self.download('tarball', self.paths['url']['tarball'], self.paths['file']['tarball'])

        elif task == 'download_tarball_sig':
            print _('Downloading'), self.paths['url']['tarball_sig']
            self.download('signature', self.paths['url']['tarball_sig'], self.paths['file']['tarball_sig'])

        elif task == 'verify':
            print _('Verifying signature')
            self.verify()

        elif task == 'extract':
            print _('Extracting'), self.paths['filename']['tarball']
            self.extract()

        elif task == 'run':
            print _('Running'), self.paths['file']['start']
            self.run()
        
        elif task == 'start_over':
            print _('Starting download over again')
            self.start_over()

    def response_received(self, response):
        class FileDownloader(Protocol):
            def __init__(self, file, total, progress, done_cb):
                self.file = file
                self.total = total
                self.so_far = 0
                self.progress = progress
                self.all_done = done_cb

            def dataReceived(self, bytes):
                self.file.write(bytes)
                self.so_far += len(bytes)
                percent = float(self.so_far) / float(self.total)
                self.progress.set_fraction(percent)
                amount = float(self.so_far)
                units = "bytes"
                for (size, unit) in [(1024 * 1024, "MiB"), (1024, "KiB")]:
                    if amount > size:
                        units = unit
                        amount = amount / float(size)
                        break

                self.progress.set_text(_('Downloaded')+(' %2.1f%% (%2.1f %s)' % ((percent * 100.0), amount, units)))

            def connectionLost(self, reason):
                print _('Finished receiving body:'), reason.getErrorMessage()
                self.all_done(reason)

        dl = FileDownloader(self.file_download, response.length, self.progressbar, self.response_finished)
        response.deliverBody(dl)

    def response_finished(self, msg):
        if msg.check(ResponseDone):
            self.file_download.close()
            # next task!
            self.run_task()

        else:
            print "FINISHED", msg
            ## FIXME handle errors

    def download_error(self, f):
        print _("Download error"), f
        self.set_gui('error', _("Error starting download:\n\n{0}\n\nAre you connected to the internet?").format(f.value), [], False)
        self.clear_ui()
        self.build_ui()

    def download(self, name, url, path):
        # initialize the progress bar
        self.progressbar.set_fraction(0) 
        self.progressbar.set_text(_('Downloading {0}').format(name))
        self.progressbar.show()
        self.refresh_gtk()

        agent = Agent(reactor, VerifyTorProjectCert(self.paths['file']['torproject_pem']))
        d = agent.request('GET', url,
                          Headers({'User-Agent': ['torbrowser-launcher']}),
                          None)

        self.file_download = open(path, 'w')
        d.addCallback(self.response_received).addErrback(self.download_error)
        
        if not reactor.running:
            reactor.run()

    def attempt_update(self):
        # load the update check file
        try:
            versions = json.load(open(self.paths['file']['update_check']))
            latest_version = None

            end = '-Linux'
            for version in versions:
                if str(version).find(end) != -1:
                    latest_version = str(version)

            if latest_version:
                self.settings['latest_version'] = latest_version[:-len(end)]
                self.settings['last_update_check_timestamp'] = int(time.time())
                self.save_settings()
                self.build_paths(self.settings['latest_version'])
                self.start_launcher()

            else:
                # failed to find the latest version
                self.set_gui('error', _("Error checking for updates."), [], False)
        
        except:
            # not a valid JSON object
            self.set_gui('error', _("Error checking for updates."), [], False)

        # now start over
        self.clear_ui()
        self.build_ui()

    def verify(self):
        # initialize the progress bar
        self.progressbar.set_fraction(0) 
        self.progressbar.set_text(_('Verifying Signature'))
        self.progressbar.show()

        p = subprocess.Popen(['/usr/bin/gpg', '--homedir', self.paths['dir']['gnupg_homedir'], '--verify', self.paths['file']['tarball_sig']])
        self.pulse_until_process_exits(p)
        
        if p.returncode == 0:
            self.run_task()
        else:
            self.set_gui('task', _("SIGNATURE VERIFICATION FAILED!\n\nYou might be under attack, or there might just be a networking problem. Click Start try the download again."), ['start_over'], False)
            self.clear_ui()
            self.build_ui()

            if not reactor.running:
                reactor.run()

    def extract(self):
        # initialize the progress bar
        self.progressbar.set_fraction(0) 
        self.progressbar.set_text(_('Installing'))
        self.progressbar.show()
        self.refresh_gtk()

        # make sure this file is a tarfile
        if tarfile.is_tarfile(self.paths['file']['tarball']):
          tf = tarfile.open(self.paths['file']['tarball'])
          tf.extractall(self.paths['dir']['tbb'])
        else:
            self.set_gui('task', _("Tor Browser Launcher doesn't understand the file format of {0}"), ['start_over'], False)
            self.clear_ui()
            self.build_ui()

        # installation is finished, so save installed_version
        self.settings['installed_version'] = self.settings['latest_version']
        self.save_settings()

        self.run_task()

    def run(self, run_next_task = True):
        subprocess.Popen([self.paths['file']['start']])
        if run_next_task:
            self.run_task()

    # make the progress bar pulse until process p (a Popen object) finishes
    def pulse_until_process_exits(self, p):
        while p.poll() == None:
            time.sleep(0.01)
            self.progressbar.pulse()
            self.refresh_gtk()

    # start over and download TBB again
    def start_over(self):
        self.label.set_text(_("Downloading Tor Browser Bundle over again."))
        self.gui_tasks = ['download_tarball', 'download_tarball_sig', 'verify', 'extract', 'run']
        self.gui_task_i = 0
        self.start(None)

    # load settings
    def load_settings(self):
        if os.path.isfile(self.paths['file']['settings']):
            self.settings = pickle.load(open(self.paths['file']['settings']))
            # sanity checks
            if not 'installed_version' in self.settings:
                return False
            if not 'latest_version' in self.settings:
                return False
            if not 'last_update_check_timestamp' in self.settings:
                return False
        else:
            self.settings = {
                'installed_version': False,
                'latest_version': '0',
                'last_update_check_timestamp': 0
            }
            self.save_settings()
        return True

    # save settings
    def save_settings(self):
        pickle.dump(self.settings, open(self.paths['file']['settings'], 'w'))
        return True
    
    # refresh gtk
    def refresh_gtk(self):
        while gtk.events_pending():
            gtk.main_iteration(False)

    # exit
    def delete_event(self, widget, event, data=None):
        return False
    def destroy(self, widget, data=None):
        if hasattr(self, 'file_download'):
            self.file_download.close()
        if reactor.running:
            reactor.stop()

if __name__ == "__main__":
    tor_browser_launcher_version = '0.1-alpha'

    print _('Tor Browser Launcher')
    print _('By Micah Lee, licensed under GPLv3')
    print _('version {0}').format(tor_browser_launcher_version)
    print 'https://github.com/micahflee/torbrowser-launcher'

    app = TorBrowserLauncher()