files/usr/bin/invirt-novnc-wsproxy

   1 #!/usr/bin/env python
   2
   3 import socket
   4 import ssl
   5 from select import select
   6 import tempfile
   7 import urlparse
   8 from novnc.websocket import WebSocketServer
   9 import invirt.remctl
  10 from invirt.config import structs as config
  11
  12 # From Python >=2.7.9
  13 _RESTRICTED_SERVER_CIPHERS = (
  14     'ECDH+AESGCM:DH+AESGCM:ECDH+AES256:DH+AES256:ECDH+AES128:DH+AES:ECDH+HIGH:'
  15     'DH+HIGH:ECDH+3DES:DH+3DES:RSA+AESGCM:RSA+AES:RSA+HIGH:RSA+3DES:!aNULL:'
  16     '!eNULL:!MD5:!DSS:!RC4'
  17 )
  18
  19 class WebSocketProxy(WebSocketServer):
  20     """
  21     Proxy traffic from a WebSockets client to an Invirt VNC server,
  22     doing the auth handshake for the client.
  23     """
  24
  25     def new_client(self):
  26         url = urlparse.urlparse(self.path)
  27         query = urlparse.parse_qs(url.query)
  28
  29         host = query.get('host', [None])[-1]
  30         vmname = query.get('vmname', [None])[-1]
  31         token = query.get('token', [None])[-1]
  32
  33         target_host = None
  34         target_port = config.vnc.base_port
  35
  36         for h in config.hosts:
  37             if h.hostname == host:
  38                 target_host = h.ip
  39
  40         if not target_host:
  41             raise Exception("host not found")
  42         if not vmname:
  43             raise Exception("vmname not provided")
  44         if not token:
  45             raise Exception("token not provided")
  46
  47         tsock = self.socket(target_host, target_port, connect=True)
  48
  49         with tempfile.NamedTemporaryFile() as cafile:
  50             cadata = invirt.remctl.remctl(config.remote.hostname, "web", "vnccert", host)
  51             cafile.write(cadata)
  52             cafile.flush()
  53
  54             # TODO: Use ssl.create_default_context when we move to Python >=2.7.9
  55             tsock = ssl.wrap_socket(tsock, ca_certs=cafile.name, cert_reqs=ssl.CERT_REQUIRED, ssl_version=ssl.PROTOCOL_SSLv23, ciphers=_RESTRICTED_SERVER_CIPHERS)
  56
  57             # Start proxying
  58             try:
  59                 extra_data = self.do_auth_handshake(tsock, vmname, token)
  60                 self.do_proxy(tsock, extra_data)
  61             except:
  62                 if tsock:
  63                     tsock.shutdown(socket.SHUT_RDWR)
  64                     tsock.close()
  65                     self.vmsg("%s:%s: Target closed" % (
  66                         target_host, target_port))
  67                 raise
  68
  69     def do_auth_handshake(self, target, vmname, token):
  70         target.send("CONNECTVNC %s VNCProxy/1.0\r\nAuth-token: %s\r\n\r\n" % (vmname, token))
  71         data = target.recv(128)
  72         if data.startswith("VNCProxy/1.0 200 "):
  73             if "\n" in data:
  74                 return data[data.find("\n")+3:]
  75             return None
  76         else:
  77             raise Exception(data)
  78
  79     def do_proxy(self, target, extra_data):
  80         """
  81         Proxy client WebSocket to normal target socket.
  82         """
  83         cqueue = []
  84         c_pend = 0
  85         tqueue = []
  86         rlist = [self.client, target]
  87
  88         if extra_data:
  89             tqueue.append(extra_data)
  90
  91         while True:
  92             wlist = []
  93
  94             if tqueue: wlist.append(target)
  95             if cqueue or c_pend: wlist.append(self.client)
  96             ins, outs, excepts = select(rlist, wlist, [], 1)
  97             if excepts: raise Exception("Socket exception")
  98
  99             if target in outs:
 100                 # Send queued client data to the target
 101                 dat = tqueue.pop(0)
 102                 sent = target.send(dat)
 103                 if sent != len(dat):
 104                     # requeue the remaining data
 105                     tqueue.insert(0, dat[sent:])
 106
 107
 108             if target in ins:
 109                 # Receive target data, encode it and queue for client
 110                 buf = target.recv(self.buffer_size)
 111                 if len(buf) == 0: raise self.EClose("Target closed")
 112
 113                 cqueue.append(buf)
 114
 115
 116             if self.client in outs:
 117                 # Send queued target data to the client
 118                 c_pend = self.send_frames(cqueue)
 119
 120                 cqueue = []
 121
 122
 123             if self.client in ins:
 124                 # Receive client data, decode it, and queue for target
 125                 bufs, closed = self.recv_frames()
 126                 tqueue.extend(bufs)
 127
 128                 if closed:
 129                     # TODO: What about blocking on client socket?
 130                     self.send_close()
 131                     raise self.EClose(closed)
 132
 133 if __name__ == '__main__':
 134     server = WebSocketProxy(cert="/etc/apache2/ssl/server.crt",
 135                             key="/etc/apache2/ssl/server.key",
 136                             listen_port=config.vnc.novnc_port,
 137                             ssl_only=True)
 138     server.start_server()