#!/usr/bin/python3

import socket
import socketserver

LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 10004
SRS_HOST = "127.0.0.1"
SRS_PORT = 10003
SRS_DOMAIN = "srs.warenform.de"

MAX_REQUEST = 8192
TIMEOUT = 3


def read_netstring(sock):
    length = bytearray()

    while True:
        c = sock.recv(1)
        if not c:
            raise EOFError()
        if c == b":":
            break
        if not c.isdigit() or len(length) >= 8:
            raise ValueError("Invalid netstring length")
        length.extend(c)

    if not length:
        raise ValueError("Empty netstring length")

    size = int(length)
    if size > MAX_REQUEST:
        raise ValueError("Netstring too large")

    data = bytearray()
    while len(data) < size:
        chunk = sock.recv(size - len(data))
        if not chunk:
            raise EOFError()
        data.extend(chunk)

    if sock.recv(1) != b",":
        raise ValueError("Invalid netstring terminator")

    return bytes(data)


def netstring(data):
    return str(len(data)).encode() + b":" + data + b","


class Handler(socketserver.BaseRequestHandler):
    def handle(self):
        self.request.settimeout(TIMEOUT)

        while True:
            try:
                request = read_netstring(self.request)
            except (EOFError, OSError):
                return
            except ValueError:
                self.request.sendall(netstring(b"PERM Invalid request"))
                return

            try:
                operation, address = request.decode("utf-8").split(" ", 1)
            except (UnicodeError, ValueError):
                self.request.sendall(netstring(b"PERM Invalid request"))
                continue

            if operation != "reverse":
                self.request.sendall(netstring(b"NOTFOUND "))
                continue

            if "@" not in address:
                self.request.sendall(netstring(b"NOTFOUND "))
                continue

            domain = address.rsplit("@", 1)[1].lower()

            if domain != SRS_DOMAIN:
                self.request.sendall(netstring(b"NOTFOUND "))
                continue

            try:
                with socket.create_connection(
                    (SRS_HOST, SRS_PORT), timeout=TIMEOUT
                ) as upstream:
                    upstream.settimeout(TIMEOUT)
                    upstream.sendall(netstring(request))
                    response = read_netstring(upstream)

                if response.startswith((b"OK ", b"NOTFOUND ", b"TEMP ", b"PERM ")):
                    result = response
                else:
                    result = b"TEMP Invalid upstream response"

            except (OSError, EOFError, ValueError):
                result = b"TEMP SRS backend unavailable"

            self.request.sendall(netstring(result))


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


if __name__ == "__main__":
    with Server((LISTEN_HOST, LISTEN_PORT), Handler) as server:
        server.serve_forever()