Standalone VNC client in Python, for people who don't want to run Java

svn path=/trunk/scripts/vnc-client/; revision=1969

diff --git a/invirt-vnc-client b/invirt-vnc-client
new file mode 100755 (executable)
index 0000000..a4a9570
--- /dev/null
+++ b/invirt-vnc-client
@@ -0,0 +1,194 @@
+#!/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()