#!/usr/bin/env python

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

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


class TorBrowserLauncher:
  def __init__(self, current_tbb_version):
    # initialize the app
    self.current_tbb_version = current_tbb_version
    self.discover_arch_lang()
    self.build_paths()
    self.mkdirs()

    launch_gui = True

    # is TBB already installed?
    if os.path.isfile(self.paths['file']['start']) and os.access(self.paths['file']['start'], os.X_OK):
      # does the version file exist?
      if os.path.isfile(self.paths['file']['version']):
        installed_tbb_version = open(self.paths['file']['version']).read().strip()

        if installed_tbb_version == current_tbb_version:
          # current version is tbb is installed, launch it
          self.run(False)
          launch_gui = False
        elif installed_tbb_version < self.current_tbb_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?", [])

      else:
        # if tbb is installed but the version file doesn't exist, something is wrong
        self.set_gui('error', "Something is wrong. You have the Tor Browser Bundle installed, but the version file is missing.", [])

    # not installed
    else:
      # save the current version to the file
      open(self.paths['file']['version'], 'w').write(self.current_tbb_version)

      # 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'])

    if launch_gui:
      self.build_ui()
      #gtk.main()
      reactor.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_data = os.getenv('HOME')+'/.torbrowser'
    tarball_filename = 'tor-browser-gnu-linux-'+self.architecture+'-'+self.current_tbb_version+'-dev-'+self.language+'.tar.gz'

    self.paths = {
      'dir': {
        'data': tbb_data,
        'download': tbb_data+'/download',
        'tbb': tbb_data+'/tbb/'+self.architecture,
        'gpg': tbb_data+'/gpgtmp'
      },
      'file': {
        'version': tbb_data+'/version',
        'start': tbb_data+'/tbb/'+self.architecture+'/tor-browser_'+self.language+'/start-tor-browser',
        'tarball': tbb_data+'/download/'+tarball_filename,
        'tarball_sig': tbb_data+'/download/'+tarball_filename+'.asc',
        'verify': '/usr/share/torbrowser-launcher/verify.sh'
      },
      'url': {
        'tarball': 'https://www.torproject.org/dist/torbrowser/linux/'+tarball_filename,
        'tarball_sig': 'https://www.torproject.org/dist/torbrowser/linux/'+tarball_filename+'.asc'
      },
      'filename': {
        'tarball': tarball_filename,
        'tarball_sig': tarball_filename+'.asc'
      }
    }

  # create directories that don't exist
  def mkdirs(self):
    if os.path.exists(self.paths['dir']['download']) == False:
      os.makedirs(self.paths['dir']['download'])
    if os.path.exists(self.paths['dir']['tbb']) == False:
      os.makedirs(self.paths['dir']['tbb'])

  # 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_autostart = autostart

  # build the application's UI
  def build_ui(self):
    self.timer = False

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

    # set up the window
    self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
    self.window.set_title("Tor Browser")
    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)

    self.box = gtk.VBox(False, 20)
    self.window.add(self.box)

    if self.gui == 'error':
      # labels
      self.label1 = gtk.Label( self.gui_message ) 
      self.label1.set_line_wrap(True)
      self.box.pack_start(self.label1, True, True, 0)
      self.label1.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
    self.start_button.set_sensitive(False)

    # start running tasks
    self.gui_task_i = 0
    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_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)
        self.progress.set_text('Downloaded %2.1f%%' % (percent * 100.0))

      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):
    print dir(msg)
    if msg.check(ResponseDone):
      self.file_download.close()
      # next task!
      self.run_task()

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


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

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

    self.file_download = open(path, 'w')
    d.addCallback(self.response_received)

  def download_chunk(self, name):
    # download 10kb a time
    chunk = self.dl_response.read(10240)
    self.dl_bytes_so_far += len(chunk)
    self.file_download.write(chunk)

    if not chunk:
      self.file_download.close()
      # next task!
      self.run_task()
      return False

    percent = float(self.dl_bytes_so_far) / self.dl_total_size
    self.progressbar.set_fraction(percent)
    percent = round(percent*100, 2)
    self.progressbar.set_text("Downloaded %d%% of %s" % (percent, name))
    self.refresh_gtk()
    
    sys.stdout.write("Downloaded %d of %d bytes (%0.2f%%)\r" % (self.dl_bytes_so_far, self.dl_total_size, percent))

    if self.dl_bytes_so_far >= self.dl_total_size:
      sys.stdout.write('\n')

    return True

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

    p = subprocess.Popen([self.paths['file']['verify'], self.paths['dir']['gpg'], self.paths['file']['tarball_sig']], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    self.pulse_until_process_exits(p)

    output = p.stdout.read()
    
    if 'Good signature' in output:
      self.run_task()
    else:
      self.progressbar.hide()
      self.label.set_text("SIGNATURE VERIFICATION FAILED!\n\nYou might be under attack, or there might just be a networking problem. Click Start try the download again.")
      self.gui_tasks = ['start_over']
      self.gui_task_i = 0
      self.start_button.show()
      self.start_button.set_sensitive(True)

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

    p = subprocess.Popen(['tar', '-xf', self.paths['file']['tarball'], '-C', self.paths['dir']['tbb']], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    self.pulse_until_process_exits(p)

    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)
  
  # 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 self.timer:
      gobject.source_remove(self.timer)
    self.timer = False

    gtk.main_quit()

if __name__ == "__main__":
  current_tbl_version = '0.1'
  current_tbb_version = '2.3.25-2'

  print 'Tor Browser Launcher'
  print 'version %s' % (current_tbl_version)
  print 'https://github.com/micahflee/torbrowser-launcher'

  app = TorBrowserLauncher(current_tbb_version)