invirt-vnc-client

   1 #!/usr/bin/python
   2 from twisted.internet import reactor, ssl, protocol
   3 from OpenSSL import SSL
   4 import base64, pickle
   5 import getopt, sys
   6
   7 verbose = False
   8
   9 def usage():
  10     print """%s [-v] [-l [HOST:]PORT] {-a AUTHTOKEN|VMNAME}
  11  -l, --listen [HOST:]PORT  port (and optionally host) to listen on for
  12                            connections (default is 127.0.0.1 and a randomly
  13                            chosen port)
  14  -a, --authtoken AUTHTOKEN Authentication token for connecting to the VNC server
  15  VMNAME                    VM name to connect to (automatically fetches an
  16                            authentication token using remctl)
  17  -v                        verbose status messages""" % (sys.argv[0])
  18
  19 class ClientContextFactory(ssl.ClientContextFactory):
  20
  21     def _verify(self, connection, x509, errnum, errdepth, ok):
  22         print '_verify (ok=%d):' % ok
  23         print '  subject:', x509.get_subject()
  24         print '  issuer:', x509.get_issuer()
  25         print '  errnum %s, errdepth %d' % (errnum, errdepth)
  26         return ok
  27
  28     def getContext(self):
  29         ctx = ssl.ClientContextFactory.getContext(self)
  30
  31         certFile = '/mit/xvm/vnc/servers.cert'
  32         if verbose: print "Loading certificates from %s" % certFile
  33         ctx.load_verify_locations(certFile)
  34         ctx.set_verify(SSL.VERIFY_PEER|SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
  35                        self._verify)
  36
  37         return ctx
  38
  39 class Proxy(protocol.Protocol):
  40     peer = None
  41
  42     def setPeer(self, peer):
  43         self.peer = peer
  44
  45     def connectionLost(self, reason):
  46         if self.peer is not None:
  47             self.peer.transport.loseConnection()
  48             self.peer = None
  49
  50     def dataReceived(self, data):
  51         self.peer.transport.write(data)
  52
  53 class ProxyClient(Proxy):
  54     ready = False
  55
  56     def connectionMade(self):
  57         self.peer.setPeer(self)
  58         data = "CONNECTVNC %s VNCProxy/1.0\r\nAuth-token: %s\r\n\r\n" % (self.factory.machine, self.factory.authtoken)
  59         self.transport.write(data)
  60         if verbose: print "ProxyClient: connection made"
  61     def dataReceived(self, data):
  62         if not ready:
  63             if verbose: print 'ProxyClient: received data "%s"' % data
  64             if data.startswith("VNCProxy/1.0 200 "):
  65                 ready = True
  66                 if "\n" in data:
  67                     self.peer.transport.write(data[data.find("\n")+1:])
  68                 self.peer.transport.resumeProducing() # Allow reading
  69             else:
  70                 print "Failed to connect: %s" % data
  71                 self.transport.loseConnection()
  72
  73 class ProxyClientFactory(protocol.ClientFactory):
  74     protocol = ProxyClient
  75
  76     def __init__(self, authtoken, machine):
  77         self.authtoken = authtoken
  78         self.machine = machine
  79
  80     def setServer(self, server):
  81         self.server = server
  82
  83     def buildProtocol(self, *args, **kw):
  84         prot = protocol.ClientFactory.buildProtocol(self, *args, **kw)
  85         prot.setPeer(self.server)
  86         return prot
  87
  88     def clientConnectionFailed(self, connector, reason):
  89         self.server.transport.loseConnection()
  90
  91
  92 class ProxyServer(Proxy):
  93     clientProtocolFactory = ProxyClientFactory
  94     authtoken = None
  95     machine = None
  96
  97     def connectionMade(self):
  98         # Don't read anything from the connecting client until we have
  99         # somewhere to send it to.
 100         self.transport.pauseProducing()
 101
 102         if verbose: print "ProxyServer: connection made"
 103
 104         client = self.clientProtocolFactory(self.factory.authtoken, self.factory.machine)
 105         client.setServer(self)
 106
 107         reactor.connectSSL(self.factory.host, self.factory.port, client, ClientContextFactory())
 108
 109
 110 class ProxyFactory(protocol.Factory):
 111     protocol = ProxyServer
 112
 113     def __init__(self, host, port, authtoken, machine):
 114         self.host = host
 115         self.port = port
 116         self.authtoken = authtoken
 117         self.machine = machine
 118
 119 def main():
 120     global verbose
 121     try:
 122         opts, args = getopt.gnu_getopt(sys.argv[1:], "hl:a:v",
 123                                        ["help", "listen=", "authtoken="])
 124     except getopt.GetoptError, err:
 125         print str(err) # will print something like "option -a not recognized"
 126         usage()
 127         sys.exit(2)
 128     listen = ["127.0.0.1", None]
 129     authtoken = None
 130     for o, a in opts:
 131         if o == "-v":
 132             verbose = True
 133         elif o in ("-h", "--help"):
 134             usage()
 135             sys.exit()
 136         elif o in ("-l", "--listen"):
 137             if ":" in a:
 138                 listen = a.split(":", 2)
 139             else:
 140                 listen[1] = a
 141         elif o in ("-a", "--authtoken"):
 142             authtoken = a
 143         else:
 144             assert False, "unhandled option"
 145
 146     # Get authentication token
 147     if authtoken is None:
 148         # User didn't give us an authentication token, so we need to get one
 149         if len(args) != 1:
 150             print "VMNAME not given or too many arguments"
 151             usage()
 152             sys.exit(2)
 153         from subprocess import PIPE, Popen
 154         try:
 155             p = Popen(["remctl", "remote", "control", args[0], "vnctoken"],
 156                       stdout=PIPE)
 157         except OSError:
 158             if verbose: print "remctl not found in path. Trying remctl locker."
 159             p = Popen(["athrun", "remctl", "remctl",
 160                        "remote", "control", args[0], "vnctoken"],
 161                       stdout=PIPE)
 162         authtoken = p.communicate()[0]
 163         if p.returncode != 0:
 164             print "Unable to get authentication token"
 165             sys.exit(1)
 166         if verbose: print 'Got authentication token "%s" for VM %s' % \
 167                           (authtoken, args[0])
 168
 169     # Unpack authentication token
 170     try:
 171         token_outer = base64.urlsafe_b64decode(authtoken)
 172         token_outer = pickle.loads(token_outer)
 173         token_inner = pickle.loads(token_outer["data"])
 174         machine = token_inner["machine"]
 175         connect_host = token_inner["connect_host"]
 176         connect_port = token_inner["connect_port"]
 177         token_expires = token_inner["expires"]
 178         if verbose: print "Unpacked authentication token:\n%s" % \
 179                           repr(token_inner)
 180     except:
 181         print "Invalid authentication token"
 182         sys.exit(1)
 183
 184     if verbose: print "Will connect to %s:%s" % (connect_host, connect_port)
 185
 186     listen[1] = 10003
 187     reactor.listenTCP(listen[1], ProxyFactory(connect_host, connect_port, authtoken, machine))
 188
 189     print "Ready to connect. Connect to %s:%s now with your VNC client. The password is 'moocow'." % (listen[0], listen[1])
 190
 191     reactor.run()
 192
 193 if '__main__' == __name__:
 194     main()