path: root/Sshuttle VPN.app/Contents/Resources/main.py

import sys, os, pty
from AppKit import *
import my, models, askpass

def sshuttle_args(host, auto_nets, auto_hosts, dns, nets, debug,
                  no_latency_control):
    argv = [my.bundle_path('sshuttle/sshuttle', ''), '-r', host]
    assert(argv[0])
    if debug:
        argv.append('-v')
    if auto_nets:
        argv.append('--auto-nets')
    if auto_hosts:
        argv.append('--auto-hosts')
    if dns:
        argv.append('--dns')
    if no_latency_control:
        argv.append('--no-latency-control')
    argv += nets
    return argv


class _Callback(NSObject):
    def initWithFunc_(self, func):
        self = super(_Callback, self).init()
        self.func = func
        return self
    def func_(self, obj):
        return self.func(obj)


class Callback:
    def __init__(self, func):
        self.obj = _Callback.alloc().initWithFunc_(func)
        self.sel = self.obj.func_


class Runner:
    def __init__(self, argv, logfunc, promptfunc, serverobj):
        print 'in __init__'
        self.id = argv
        self.rv = None
        self.pid = None
        self.fd = None
        self.logfunc = logfunc
        self.promptfunc = promptfunc
        self.serverobj = serverobj
        self.buf = ''
        self.logfunc('\nConnecting to %s.\n' % self.serverobj.host())
        print 'will run: %r' % argv
        self.serverobj.setConnected_(False)
        pid,fd = pty.fork()
        if pid == 0:
            # child
            try:
                os.execvp(argv[0], argv)
            except Exception, e:
                sys.stderr.write('failed to start: %r\n' % e)
                raise
            finally:
                os._exit(42)
        # parent
        self.pid = pid
        self.file = NSFileHandle.alloc()\
               .initWithFileDescriptor_closeOnDealloc_(fd, True)
        self.cb = Callback(self.gotdata)
        NSNotificationCenter.defaultCenter()\
            .addObserver_selector_name_object_(self.cb.obj, self.cb.sel,
                        NSFileHandleDataAvailableNotification, self.file)
        self.file.waitForDataInBackgroundAndNotify()

    def __del__(self):
        self.wait()

    def _try_wait(self, options):
        if self.rv == None and self.pid > 0:
            pid,code = os.waitpid(self.pid, options)
            if pid == self.pid:
                if os.WIFEXITED(code):
                    self.rv = os.WEXITSTATUS(code)
                else:
                    self.rv = -os.WSTOPSIG(code)
                self.serverobj.setConnected_(False)
                self.serverobj.setError_('VPN process died')
                self.logfunc('Disconnected.\n')
        print 'wait_result: %r' % self.rv
        return self.rv

    def wait(self):
        return self._try_wait(0)
        
    def poll(self):
        return self._try_wait(os.WNOHANG)

    def kill(self):
        assert(self.pid > 0)
        print 'killing: pid=%r rv=%r' % (self.pid, self.rv)
        if self.rv == None:
            self.logfunc('Disconnecting from %s.\n' % self.serverobj.host())
            os.kill(self.pid, 15)
            self.wait()

    def gotdata(self, notification):
        print 'gotdata!'
        d = str(self.file.availableData())
        if d:
            self.logfunc(d)
            self.buf = self.buf + d
            if 'Connected.\r\n' in self.buf:
                self.serverobj.setConnected_(True)
            self.buf = self.buf[-4096:]
            if self.buf.strip().endswith(':'):
                lastline = self.buf.rstrip().split('\n')[-1]
                resp = self.promptfunc(lastline)
                add = ' (response)\n'
                self.buf += add
                self.logfunc(add)
                self.file.writeData_(my.Data(resp + '\n'))
            self.file.waitForDataInBackgroundAndNotify()
        self.poll()
        #print 'gotdata done!'


class SshuttleApp(NSObject):
    def initialize(self):
        d = my.PList('UserDefaults') 
        my.Defaults().registerDefaults_(d)


class SshuttleController(NSObject):
    # Interface builder outlets
    startAtLoginField = objc.IBOutlet()
    autoReconnectField = objc.IBOutlet()
    debugField = objc.IBOutlet()
    routingField = objc.IBOutlet()
    prefsWindow = objc.IBOutlet()
    serversController = objc.IBOutlet()
    logField = objc.IBOutlet()
    latencyControlField = objc.IBOutlet()
    
    servers = []
    conns = {}

    def _connect(self, server):
        host = server.host()
        print 'connecting %r' % host
        self.fill_menu()
        def logfunc(msg):
            print 'log! (%d bytes)' % len(msg)
            self.logField.textStorage()\
                .appendAttributedString_(NSAttributedString.alloc()\
                                         .initWithString_(msg))
            self.logField.didChangeText()
        def promptfunc(prompt):
            print 'prompt! %r' % prompt
            return askpass.askpass(prompt)
        nets_mode = server.autoNets()
        if nets_mode == models.NET_MANUAL:
            manual_nets = ["%s/%d" % (i.subnet(), i.width())
                           for i in server.nets()]
        elif nets_mode == models.NET_ALL:
            manual_nets = ['0/0']
        else:
            manual_nets = []
        noLatencyControl = (server.latencyControl() != models.LAT_INTERACTIVE)
        conn = Runner(sshuttle_args(host,
                                    auto_nets = nets_mode == models.NET_AUTO,
                                    auto_hosts = server.autoHosts(),
                                    dns = server.useDns(),
                                    nets = manual_nets,
                                    debug = self.debugField.state(),
                                    no_latency_control = noLatencyControl),
                      logfunc=logfunc, promptfunc=promptfunc,
                      serverobj=server)
        self.conns[host] = conn

    def _disconnect(self, server):
        host = server.host()
        print 'disconnecting %r' % host
        conn = self.conns.get(host)
        if conn:
            conn.kill()
        self.fill_menu()
        self.logField.textStorage().setAttributedString_(
                        NSAttributedString.alloc().initWithString_(''))
    
    @objc.IBAction
    def cmd_connect(self, sender):
        server = sender.representedObject()
        server.setWantConnect_(True)

    @objc.IBAction
    def cmd_disconnect(self, sender):
        server = sender.representedObject()
        server.setWantConnect_(False)

    @objc.IBAction
    def cmd_show(self, sender):
        self.prefsWindow.makeKeyAndOrderFront_(self)
        NSApp.activateIgnoringOtherApps_(True)

    @objc.IBAction
    def cmd_quit(self, sender):
        NSApp.performSelector_withObject_afterDelay_(NSApp.terminate_,
                                                     None, 0.0)

    def fill_menu(self):
        menu = self.menu
        menu.removeAllItems()

        def additem(name, func, obj):
            it = menu.addItemWithTitle_action_keyEquivalent_(name, None, "")
            it.setRepresentedObject_(obj)
            it.setTarget_(self)
            it.setAction_(func)
        def addnote(name):
            additem(name, None, None)

        any_inprogress = None
        any_conn = None
        any_err = None
        if len(self.servers):
            for i in self.servers:
                host = i.host()
                title = i.title()
                want = i.wantConnect()
                connected = i.connected()
                numnets = len(list(i.nets()))
                if not host:
                    additem('Connect Untitled', None, i)
                elif i.autoNets() == models.NET_MANUAL and not numnets:
                    additem('Connect %s (no routes)' % host, None, i)
                elif want:
                    any_conn = i
                    additem('Disconnect %s' % title, self.cmd_disconnect, i)
                else:
                    additem('Connect %s' % title, self.cmd_connect, i)
                if not want:
                    msg = 'Off'
                elif i.error():
                    msg = 'ERROR - try reconnecting'
                    any_err = i
                elif connected:
                    msg = 'Connected'
                else:
                    msg = 'Connecting...'
                    any_inprogress = i
                addnote('   State: %s' % msg)
        else:
            addnote('No servers defined yet')

        menu.addItem_(NSMenuItem.separatorItem())
        additem('Preferences...', self.cmd_show, None)
        additem('Quit Sshuttle VPN', self.cmd_quit, None)

        if any_err:
            self.statusitem.setImage_(self.img_err)
            self.statusitem.setTitle_('Error!')
        elif any_conn:
            self.statusitem.setImage_(self.img_running