import socket

from threading import Thread
from enum import Enum


class Direction(Enum):
    CLIENT_TO_PROXY = '->'
    PROXY_TO_SERVER = '<-'


getPacketsSize = {
    0xC1: lambda packet: packet[1],
    0xC2: lambda packet: packet[1] << 8 | packet[2],
    0xC3: lambda packet: packet[1],
    0xC4: lambda packet: packet[1] << 8 | packet[2]
}


class PacketReader(Thread):
    def __init__(self, direction, port):
        super(PacketReader, self).__init__()
        self.input = None
        self.output = None
        self.port = port
        self.direction = direction

    def run(self):
        queue = []
        while True:
            data = self.input.recv(4096)
            if data:
                if self.direction == Direction.CLIENT_TO_PROXY:
                     print(f"BUFFER [{self.port}] {self.direction.value} {data.hex()}")
                self.parse(queue, data)
                while len(queue) > 0:
                    packet = queue.pop(0)
                    if self.direction == Direction.CLIENT_TO_PROXY:
                        print(f"PACKET [{self.port}] {self.direction.value} {packet.hex()}")
                    self.output.sendall(packet)
                print("\n")

    def parse(self, queue, data):
        while len(data) > 2:
            header = data[0:3]
            size = getPacketsSize[header[0]](data)
            if len(data) < size:
                break
            queue.append(data[0:size])
            data = data[size:]


class Proxy2Server(PacketReader):
    def __init__(self, host, port):
        super(Proxy2Server, self).__init__(Direction.CLIENT_TO_PROXY, port)
        self.input = None  # game client socket
        self.output = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.output.connect((host, port))


class Client2Proxy(PacketReader):
    def __init__(self, host, port):
        super(Client2Proxy, self).__init__(Direction.PROXY_TO_SERVER, port)
        self.input = None  # real server socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((host, port))
        sock.listen(1)
        # waiting for a connection
        self.output, addr = sock.accept()


class Proxy(Thread):
    def __init__(self, from_host, to_host, port):
        super(Proxy, self).__init__()
        self.from_host = from_host
        self.to_host = to_host
        self.port = port

    def run(self):
        while True:
            print("[proxy({}) started]".format(self.port))
            self.client2Proxy = Client2Proxy(self.from_host, self.port)
            self.proxy2Server = Proxy2Server(self.to_host, 55901)
            print('[proxy({}) connection established]'.format(self.port))
            self.client2Proxy.input = self.proxy2Server.output
            self.proxy2Server.input = self.client2Proxy.output

            self.client2Proxy.start()
            self.proxy2Server.start()


proxy = Proxy('127.0.0.1', 'serverIp', 20439)
proxy.start()