
from asysocks.unicomm.common.target import UniTarget, UniProto
from asysocks.unicomm.common.packetizers import StreamPacketizer
from asysocks.unicomm.server import UniServer
from hashlib import sha1
from base64 import b64encode
import asyncio
import traceback
import struct

# https://github.com/Pithikos/python-websocket-server/

'''
+-+-+-+-+-------+-+-------------+-------------------------------+
 0                   1                   2                   3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-------+-+-------------+-------------------------------+
|F|R|R|R| opcode|M| Payload len |    Extended payload length    |
|I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
|N|V|V|V|       |S|             |   (if payload len==126/127)   |
| |1|2|3|       |K|             |                               |
+-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
|     Extended payload length continued, if payload len == 127  |
+ - - - - - - - - - - - - - - - +-------------------------------+
|                     Payload Data continued ...                |
+---------------------------------------------------------------+
'''

FIN    = 0x80
OPCODE = 0x0f
MASKED = 0x80
PAYLOAD_LEN = 0x7f
PAYLOAD_LEN_EXT16 = 0x7e
PAYLOAD_LEN_EXT64 = 0x7f

OPCODE_CONTINUATION = 0x0
OPCODE_TEXT         = 0x1
OPCODE_BINARY       = 0x2
OPCODE_CLOSE_CONN   = 0x8
OPCODE_PING         = 0x9
OPCODE_PONG         = 0xA

CLOSE_STATUS_NORMAL = 1000
DEFAULT_CLOSE_REASON = bytes('', encoding='utf-8')




class WebSocketServer:
    def __init__(self, client_handler, host, port, ssl_ctx = None, proxies = None):
        self.client_handler = client_handler
        self.host = host
        self.port = port
        self.ssl_ctx = ssl_ctx
        self.proxies = proxies
    
        self.clients = []
        self.id_counter = 0
        self.__main_task = None
        self.started_evt = asyncio.Event()

    async def __aenter__(self):
        self.__main_task = asyncio.create_task(self.serve())
        await asyncio.sleep(1)
        return self
    
    async def __aexit__(self, exc_type, exc, tb):
        await self.terminate()

    async def terminate(self):
        for client in self.clients:
            await client.terminate()
        self.clients = []
    
    async def __handle_connection(self, connection):
        client_id = self.id_counter
        self.id_counter += 1
        #print('Server: New client connected with id %s' % client_id)
        client = WebSocketClientHandler(client_id, self, connection, self.client_handler)
        self.clients.append(client)
        await client.run()
        self.clients.remove(client)

    
    async def serve(self):
        try:
            target = UniTarget(self.host, self.port, UniProto.SERVER_TCP, proxies=self.proxies)
            packetizer = StreamPacketizer()

            server = UniServer(target, packetizer)
            async for connection in server.serve():
                x = asyncio.create_task(self.__handle_connection(connection))

        except Exception as e:
            traceback.print_exc()

class WebSocketClientHandler:
    def __init__(self, client_id, server, connection, client_handler):
        self.client_id = client_id
        self.server = server
        self.connection = connection
        self.client_handler = client_handler
        self.incoming_buffer = asyncio.Queue()
    
    @staticmethod
    def make_handshake_response(key):
        return \
          'HTTP/1.1 101 Switching Protocols\r\n'\
          'Upgrade: websocket\r\n'              \
          'Connection: Upgrade\r\n'             \
          'Sec-WebSocket-Accept: %s\r\n'        \
          '\r\n' % WebSocketClientHandler.calculate_response_key(key)
    
    @staticmethod
    def calculate_response_key(key):
        GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
        hash = sha1(key.encode() + GUID.encode())
        response_key = b64encode(hash.digest()).strip()
        return response_key.decode('ASCII')

    async def __read_http_headers(self):
        headers = {}
        # first line should be HTTP GET
        http_raw = await self.connection.packetizer.readuntil(b'\r\n')
        http_get = http_raw.decode().strip()
        if http_get.upper().startswith('GET') is False:
            raise Exception('Invalid HTTP request: %s' % http_get)
        # remaining should be headers
        while True:
            http_raw = await self.connection.packetizer.readuntil(b'\r\n')
            header = http_raw.decode().strip()
            if not header:
                break
            head, value = header.split(':', 1)
            headers[head.lower().strip()] = value.strip()
        return headers
    
    async def handshake(self):
        try:
            headers = await self.__read_http_headers()

            if 'upgrade' not in headers:
                raise Exception('Client tried to connect but was missing an upgrade header')
            if headers['upgrade'].lower() != 'websocket':
                raise Exception('Client tried to connect but was missing a key')

            if 'sec-websocket-key' not in headers:
                raise Exception('Client tried to connect but was missing a key')

            key = headers['sec-websocket-key']
            response = WebSocketClientHandler.make_handshake_response(key)
            await self.connection.write(response.encode())
        except Exception as e:
            traceback.print_exc()
    
    async def __recv_internal(self):
        try:
            while True:
                b1, b2 = 0, 0
                res = await self.connection.packetizer.readexactly(2)
                b1 = res[0]
                b2 = res[1]

                fin    = b1 & FIN
                opcode = b1 & OPCODE
                masked = b2 & MASKED
                payload_length = b2 & PAYLOAD_LEN

                if opcode == OPCODE_CLOSE_CONN:
                    raise Exception("Client asked to close connection.")
                
                if not masked:
                    raise Exception("Client must always be masked.")
                
                if opcode == OPCODE_CONTINUATION:
                    raise Exception("Continuation frames are not supported.")

                if payload_length in [126, 127]:
                    payload_length_raw = await self.connection.packetizer.readexactly(2 if payload_length == 126 else 8)
                    payload_length = int.from_bytes(payload_length_raw, byteorder='big')

                masks = await self.connection.packetizer.readexactly(4)
                message_bytes = bytearray()

                raw = await self.connection.packetizer.readexactly(payload_length)
                for c in raw:
                    c ^= masks[len(message_bytes) % 4]
                    message_bytes.append(c)
                yield opcode, bytes(message_bytes)
        
        except Exception as e:
            #traceback.print_exc()
            yield None, None

    async def send_pong(self, data):
        if isinstance(data, str) is True:
            data = data.encode('utf-8')
        await self.__send_internal(OPCODE_PONG, data)

    async def send_ping(self, data):
        if isinstance(data, str) is True:
            data = data.encode('utf-8')
        await self.__send_internal(OPCODE_PING, data)

    async def start_recv(self):
        async for opcode, data in self.__recv_internal():
            if opcode is None:
                break
            #print('Received opcode: %s' % opcode)
            #print('Received data: %s' % data)
            
            if opcode == OPCODE_TEXT:
                self.incoming_buffer.put_nowait(data.decode('utf-8'))
            elif opcode == OPCODE_BINARY:
                self.incoming_buffer.put_nowait(data)
            elif opcode == OPCODE_PING:
                await self.send_pong(data)
            elif opcode == OPCODE_PONG:
                print('Received pong')
            elif opcode == OPCODE_CLOSE_CONN:
                break
            else:
                print('Unknown opcode %s' % opcode)
            
        #print('Client disconnected')

    async def __send_internal(self, opcode, data:bytes):
        header  = bytearray()
        payload_length = len(data)

        # Normal payload
        if payload_length <= 125:
            header.append(FIN | opcode)
            header.append(payload_length)

        # Extended payload
        elif payload_length >= 126 and payload_length <= 65535:
            header.append(FIN | opcode)
            header.append(PAYLOAD_LEN_EXT16)
            header.extend(struct.pack(">H", payload_length))

        # Huge extended payload
        elif payload_length < 18446744073709551616:
            header.append(FIN | opcode)
            header.append(PAYLOAD_LEN_EXT64)
            header.extend(struct.pack(">Q", payload_length))

        else:
            raise Exception("Message is too big. Consider breaking it into chunks.")
        
        await self.connection.write(header + data)
    
    async def send(self, data):
        if isinstance(data, str) is True:
            opcode = OPCODE_TEXT
            await self.__send_internal(opcode, data.encode('utf-8'))
        else:
            opcode = OPCODE_BINARY
            await self.__send_internal(opcode, data)
            
    async def run(self):
        x = asyncio.create_task(self.connection.stream())
        await self.handshake()
        x = asyncio.create_task(self.start_recv())
        conn = WebSocketClientConnection(self)
        await self.client_handler(conn)

class WebSocketClientConnection:
    def __init__(self, clienthandler):
        self.clienthandler = clienthandler
    
    def __aiter__(self):
        return self
    
    async def __anext__(self):
        return await self.clienthandler.incoming_buffer.get()
    
    async def recv(self):
        return await self.clienthandler.incoming_buffer.get()
    
    async def send(self, data):
        await self.clienthandler.send(data)

    async def ping(self, data):
        await self.clienthandler.send_ping(data)


serve = WebSocketServer