#!/usr/bin/python
from twisted.internet import reactor, ssl, protocol
from OpenSSL import SSL
import base64, pickle
import getopt, sys

verbose = False

def usage():
    print """%s [-v] [-l [HOST:]PORT] {-a AUTHTOKEN|VMNAME}
 -l, --listen [HOST:]PORT  port (and optionally host) to listen on for
                           connections (default is 127.0.0.1 and a randomly
                           chosen port)
 -a, --authtoken AUTHTOKEN Authentication token for connecting to the VNC server
 VMNAME                    VM name to connect to (automatically fetches an
                           authentication token using remctl)
 -v                        verbose status messages""" % (sys.argv[0])

class ClientContextFactory(ssl.ClientContextFactory):

    def _verify(self, connection, x509, errnum, errdepth, ok):
        print '_verify (ok=%d):' % ok
        print '  subject:', x509.get_subject()
        print '  issuer:', x509.get_issuer()
        print '  errnum %s, errdepth %d' % (errnum, errdepth)
        return ok

    def getContext(self):
        ctx = ssl.ClientContextFactory.getContext(self)

        certFile = '/mit/xvm/vnc/servers.cert'
        if verbose: print "Loading certificates from %s" % certFile
        ctx.load_verify_locations(certFile)
        ctx.set_verify(SSL.VERIFY_PEER|SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
                       self._verify)

        return ctx

class Proxy(protocol.Protocol):
    peer = None

    def setPeer(self, peer):
        self.peer = peer

    def connectionLost(self, reason):
        if self.peer is not None:
            self.peer.transport.loseConnection()
            self.peer = None

    def dataReceived(self, data):
        self.peer.transport.write(data)

class ProxyClient(Proxy):
    ready = False

    def connectionMade(self):
        self.peer.setPeer(self)
        data = "CONNECTVNC %s VNCProxy/1.0\r\nAuth-token: %s\r\n\r\n" % (self.factory.machine, self.factory.authtoken)
        self.transport.write(data)
        if verbose: print "ProxyClient: connection made"
    def dataReceived(self, data):
        if not ready:
            if verbose: print 'ProxyClient: received data "%s"' % data
            if data.startswith("VNCProxy/1.0 200 "):
                ready = True
                if "\n" in data:
                    self.peer.transport.write(data[data.find("\n")+1:])
                self.peer.transport.resumeProducing() # Allow reading
            else:
                print "Failed to connect: %s" % data
                self.transport.loseConnection()

class ProxyClientFactory(protocol.ClientFactory):
    protocol = ProxyClient
    
    def __init__(self, authtoken, machine):
        self.authtoken = authtoken
        self.machine = machine

    def setServer(self, server):
        self.server = server

    def buildProtocol(self, *args, **kw):
        prot = protocol.ClientFactory.buildProtocol(self, *args, **kw)
        prot.setPeer(self.server)
        return prot

    def clientConnectionFailed(self, connector, reason):
        self.server.transport.loseConnection()


class ProxyServer(Proxy):
    clientProtocolFactory = ProxyClientFactory
    authtoken = None
    machine = None

    def connectionMade(self):
        # Don't read anything from the connecting client until we have
        # somewhere to send it to.
        self.transport.pauseProducing()
        
        if verbose: print "ProxyServer: connection made"

        client = self.clientProtocolFactory(self.factory.authtoken, self.factory.machine)
        client.setServer(self)

        reactor.connectSSL(self.factory.host, self.factory.port, client, ClientContextFactory())
        

class ProxyFactory(protocol.Factory):
    protocol = ProxyServer

    def __init__(self, host, port, authtoken, machine):
        self.host = host
        self.port = port
        self.authtoken = authtoken
        self.machine = machine

def main():
    global verbose
    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "hl:a:v",
                                       ["help", "listen=", "authtoken="])
    except getopt.GetoptError, err:
        print str(err) # will print something like "option -a not recognized"
        usage()
        sys.exit(2)
    listen = ["127.0.0.1", None]
    authtoken = None
    for o, a in opts:
        if o == "-v":
            verbose = True
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-l", "--listen"):
            if ":" in a:
                listen = a.split(":", 2)
            else:
                listen[1] = a
        elif o in ("-a", "--authtoken"):
            authtoken = a
        else:
            assert False, "unhandled option"

    # Get authentication token
    if authtoken is None:
        # User didn't give us an authentication token, so we need to get one
        if len(args) != 1:
            print "VMNAME not given or too many arguments"
            usage()
            sys.exit(2)
        from subprocess import PIPE, Popen
        try:
            p = Popen(["remctl", "remote", "control", args[0], "vnctoken"],
                      stdout=PIPE)
        except OSError:
            if verbose: print "remctl not found in path. Trying remctl locker."
            p = Popen(["athrun", "remctl", "remctl",
                       "remote", "control", args[0], "vnctoken"],
                      stdout=PIPE)
        authtoken = p.communicate()[0]
        if p.returncode != 0:
            print "Unable to get authentication token"
            sys.exit(1)
        if verbose: print 'Got authentication token "%s" for VM %s' % \
                          (authtoken, args[0])

    # Unpack authentication token
    try:
        token_outer = base64.urlsafe_b64decode(authtoken)
        token_outer = pickle.loads(token_outer)
        token_inner = pickle.loads(token_outer["data"])
        machine = token_inner["machine"]
        connect_host = token_inner["connect_host"]
        connect_port = token_inner["connect_port"]
        token_expires = token_inner["expires"]
        if verbose: print "Unpacked authentication token:\n%s" % \
                          repr(token_inner)
    except:
        print "Invalid authentication token"
        sys.exit(1)
    
    if verbose: print "Will connect to %s:%s" % (connect_host, connect_port) 
    
    listen[1] = 10003
    reactor.listenTCP(listen[1], ProxyFactory(connect_host, connect_port, authtoken, machine))
    
    print "Ready to connect. Connect to %s:%s now with your VNC client. The password is 'moocow'." % (listen[0], listen[1])
    
    reactor.run()

if '__main__' == __name__:
    main()