#!/usr/bin/env python3
import socket, select
import subprocess
import os.path
import struct

ENDPOINT = ("", 1337)

def wtf_encode(data):
    return data.decode("latin1").encode("UTF-8")

def wtf_decode(data):
    return data.decode("UTF-8").encode("latin1")

class DisplayClient:
    def __init__(self, proxy, dcsock):
        self.docroot = "/home/matti/devel/jsx11/client/"
        self.proxy = proxy
        self.dcsock = dcsock
        self.proxy.register_sock(self.dcsock, self)
        self.xssock = None
        self.xcsocks = None
        self.wsconnected = False
        self.displaynum = None
        self.processes = None
        self.buf = bytearray()

    def die(self):
        if self.processes:
            for p in self.processes:
                p.kill()

        self.proxy.unregister_sock(self.dcsock)
        if self.xssock:
            for xc in self.xcsocks:
                self.proxy.unregister_sock(xc[1])

            self.proxy.unregister_sock(self.xssock)

        self.proxy.dclients.remove(self)
        
        if self.displaynum is not None:
            print ("Display {0} dying".format(self.displaynum))
        else:
            print ("Dying")
        

    def start_x_server(self):
        self.displaynum = self.proxy.alloc_displaynum()
        if self.displaynum is None:
            self.ws_send(0, b"\x00All display numbers in use; try again later")
            return self.die()

        self.xssock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.xssock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.xssock.bind(("localhost", 6000 + self.displaynum))
        self.xssock.listen(5)
        self.proxy.register_sock(self.xssock, self)
        
        self.newxcid = 1
        self.xcsocks = []

        env = {"DISPLAY": "localhost:{0}.0".format(self.displaynum)}

        #self.process = subprocess.Popen(["/usr/bin/xeyes"],
        #    env={"DISPLAY": "localhost:{0}.0".format(self.displaynum)})
        self.processes = [
            subprocess.Popen(["/usr/bin/xclock", "-update", "1", "-geometry", "200x200+200+10"], env=env),
            subprocess.Popen(["/usr/bin/xeyes"], env=env),
        ]

    def onread(self, sock):
        if sock == self.xssock:
            self.onread_xs_accept()
        elif sock == self.dcsock:
            if self.wsconnected:
                self.onread_ws()
            else:
                self.onread_http_or_handshake()
        else:
            self.onread_xc(sock)

    def onread_xs_accept(self):
        (xcsock, xcaddr) = self.xssock.accept()
        self.xcsocks.append((self.newxcid, xcsock))
        self.proxy.register_sock(xcsock, self)
        self.ws_send(0, b"\x01" + struct.pack("b", self.newxcid))
        self.newxcid += 1

    def onread_http_or_handshake(self):
        searchfrom = len(self.buf)
        data = self.dcsock.recv(4096)
        self.buf.extend(data)

        index = self.buf.find(b"\r\n\r\n", searchfrom)
        if index != -1:
            headers = self.buf[:index]
            self.buf = self.buf[index+4:]

            try:
                headers = headers.decode("UTF-8")
            except UnicodeError:
                return self.send_http_bad_request()

            self.handle_http_request(headers)
    
    def handle_http_request(self, headerstring):
        headers = {}
        
        lines = headerstring.splitlines()
        req = lines[0].split(" ")

        print(req)

        if len(req) != 3 or req[0] != "GET" or req[2] != "HTTP/1.1":
            return self.send_http_bad_request()
        
        for l in lines[1:]:
            parts = l.split(":", 1)
            if len(parts) != 2:
                return self.send_http_bad_request()

            headers[parts[0].strip()] = parts[1].strip()
            
        if headers.get("Connection") == "Upgrade" and headers.get("Upgrade") == "WebSocket":
            response = b"HTTP/1.1 101 Web Socket Protocol Handshake\r\nUpgrade: WebSocket\r\nConnection: Upgrade\r\n"
            response += b"WebSocket-Origin: http://tsukasa.lumpio.dy.fi:1337\r\n"
            response += b"WebSocket-Location: ws://tsukasa.lumpio.dy.fi:1337/\r\n"
            response += b"WebSocket-Protocol: JSX11\r\n"
            response += b"\r\n"
            self.dcsock.send(response)
            self.wsconnected = True
            self.start_x_server()
            return
        
        url = req[1]
        path = os.path.join(self.docroot, url.lstrip("/"))
        if not path.startswith(self.docroot):
            return self.send_http_not_found(url)
        
        f = None
        try:
            f = open(path, "r")
            self.send_http_response(contenttype=self.guess_mime(url), body=f.read())
            f.close()
            f = None
        except IOError:
            return self.send_http_not_found(url)
        finally:
            if f:
                f.close()
    
    MIME = (
        (".html", "text/html"),
        (".js", "text/javascript"),
        (".css", "text/css"),
        (".txt", "text/plain"),
        (".py", "text/plain"),
    )

    def guess_mime(self, url):
        for m in self.MIME:
            if url.endswith(m[0]):
                return m[1]

        return "application/octet-stream"

    def send_http_response(self, status="200 OK", contenttype="text/plain; charset=UTF-8", body=None):
        if body is None:
            self.dcsock.send("HTTP/1.1 {0}\r\n\r\n".format(status).encode("UTF-8"))
        else:
            if isinstance(body, str):
                body = body.encode("UTF-8")
            self.dcsock.send("HTTP/1.1 {0}\r\nContent-Type: {0}\r\nContent-Length: {1}\r\n\r\n".format(
                status, contenttype, len(body)).encode("UTF-8"))
            self.dcsock.send(body)

        return self.die()

    def send_http_not_found(self, path):
        return self.send_http_response(status="404 Not Found", body="The path '{0}' was not found on this server.".format(path))

    def send_http_bad_request(self):
        return self.send_http_response(status="400 Bad Request", body="Your client sent a request this server cannot understand.")

    def onread_ws(self):
        try:
            data = self.dcsock.recv(4096)
        except socket.error:
            return self.die()

        if len(data) == 0:
            return self.die()

        #print ("WS<<< " + repr(data))
        self.buf.extend(data)
        
        while len(self.buf) > 0:
            if self.buf[0] == 0x00:
                index = self.buf.find(b"\xff")
                if index == -1:
                    break
                
                msg = wtf_decode(self.buf[1:index])
                self.buf = self.buf[index+1:]

                if len(msg) < 1:
                    print("WS protocol error (message length < 1)")
                    return self.die()
                
                cid = msg[0]
                
                for xc in self.xcsocks:
                    if xc[0] == cid:
                        xc[1].send(msg[1:])
                        break
                else:
                    print("WS protocol error (unknown cid)")
                    return self.die()
            else:
                print("WS protocol error (unknown frame type {0})".format(self.buf[0]))
                return self.die()

    def ws_send(self, cid, data):
        try:
            self.dcsock.send(b"\x00" + struct.pack("b", cid) + wtf_encode(data) + b"\xff")
        except socket.error:
            return self.die()

    def onread_xc(self, sock):
        data = sock.recv(4096)

        if len(data) == 0:
            return self.die();

        for xc in self.xcsocks:
            if xc[1] == sock:
                self.ws_send(xc[0], data)
                break

    def __str__(self):
        r = "DisplayClient"
        
        if self.dcsock is not None:
            r += ", dc = {0}".format(self.dcsock.fileno())
        if self.xssock is not None:
            r += ", xs = {0}".format(self.xssock.fileno())

        return r

class JSX11Proxy:
    def __init__(self, endpoint):
        self.wsservsock = None
        self.socks = {}
        self.dclients = []
        self.endpoint = endpoint

    def register_sock(self, sock, owner):
        self.socks[sock] = owner

    def unregister_sock(self, sock):
        print("unregistering {0}".format(sock))
        sock.close()
        del self.socks[sock]

    def die(self):
        if self.wsservsock:
            self.unregister_sock(self.wsservsock)
        for dc in self.dclients:
            dc.die()

    #def updateallsocks(self):
    #    socks = []
    #    socks.append(self.wsservsock)
    #    for dc in self.dclients:
    #        socks.append(dc.dcsock)
    #        socks.append(dc.xssock)
    #        if dc.xcsocks is not None:
    #            socks.extend(dc.xcsocks.keys())
    
    def run(self):
        self.wsservsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.wsservsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.wsservsock.bind(self.endpoint)
        self.wsservsock.listen(5)
        self.register_sock(self.wsservsock, self)
        
        try:
            while True:
                readable, writable, exc = select.select(self.socks.keys(), [], [], 1)
                if len(readable) > 0:
                    for s in readable:
                        self.socks[s].onread(s)
        finally:
            self.die()
    
    def onread(self, sock):
        (dcsock, dcaddr) = self.wsservsock.accept()
        #print("New WSSERVSOCK connection from {0}".format(dcaddr))
        dclient = DisplayClient(self, dcsock)
        self.dclients.append(dclient)

    def alloc_displaynum(self):
        nums = set((0, 1, 2, 3, 4, 5, 6, 7))
        for dc in self.dclients:
            if dc.displaynum is not None:
                nums.remove(dc.displaynum)

        if len(nums) == 0:
            return None

        return list(nums)[0]

if __name__ == "__main__":
    proxy = None
    try:
        print("Listening on {0}".format(repr(ENDPOINT)))
        proxy = JSX11Proxy(ENDPOINT)
        proxy.run()
    finally:
        print("\nExiting...")