invirt-vnc-client

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