import configparser
import fnmatch
import http.server
import ipaddress
import os
import signal
import socket
import socketserver
import ssl
import struct
import threading


CONFIG_FILE = os.environ.get("CONFIG_FILE", "/srv/nanodns.ini")
DNS_PORT = int(os.environ.get("DNS_PORT", "8053"))
HTTP_PORT = int(os.environ.get("HTTP_PORT", "8081"))
HTTPS_PORT = int(os.environ.get("HTTPS_PORT", "8443"))
CERTIFICATE = "/srv/certificates/manuals.crt"
PRIVATE_KEY = "/srv/certificates/manuals.key"

# Built-in defaults, used when the config file (or a given key) is absent. They
# mirror the reviewed production policy: the User Guide domain is spoofed to the
# local gateway, the reviewed store/update hosts are forwarded upstream, and
# every other name returns NXDOMAIN (default-deny).
DEFAULT_REDIRECT_URL = "http://192.168.20.30:8000/"
DEFAULT_ALLOWED_CLIENTS = "192.168.40.0/28, 192.168.20.0/24, 192.168.10.0/27"
DEFAULT_UPSTREAMS = "1.1.1.1"
DEFAULT_TIMEOUT_MS = "3000"
DEFAULT_OVERRIDES = (("manuals.playstation.net", "192.168.20.30"),)
DEFAULT_EXCEPTIONS = (
    "homebrew.page",
    "itsplk.github.io",
    "store.spectrumlibrary.online",
    "api.github.com",
    "github.com",
    "raw.githubusercontent.com",
    "release-assets.githubusercontent.com",
    "prosperopatches.com",
    "cdn.prosperopatches.com",
    "image.api.playstation.com",
    "gmedia.playstation.com",
)


def _split(value):
    return [item for item in value.replace(",", " ").split() if item]


def load_config(path):
    parser = configparser.ConfigParser(
        allow_no_value=True,
        inline_comment_prefixes=("#", ";"),
    )
    # Keep mask/address case verbatim; query names are lowercased before lookup.
    parser.optionxform = str
    parser.read(path)

    general = parser["general"] if parser.has_section("general") else {}
    redirect_url = general.get("redirect_url") or DEFAULT_REDIRECT_URL
    allowed_raw = general.get("allowed_clients") or DEFAULT_ALLOWED_CLIENTS
    allowed_clients = tuple(
        ipaddress.ip_network(entry, strict=False) for entry in _split(allowed_raw)
    )

    upstream = parser["upstream"] if parser.has_section("upstream") else {}
    servers_raw = (
        upstream.get("server") or upstream.get("servers") or DEFAULT_UPSTREAMS
    )
    upstreams = tuple(
        str(ipaddress.ip_address(server)) for server in _split(servers_raw)
    )
    timeout = max(0.1, int(upstream.get("timeout_ms") or DEFAULT_TIMEOUT_MS) / 1000.0)

    override_items = (
        parser.items("overrides")
        if parser.has_section("overrides")
        else DEFAULT_OVERRIDES
    )
    overrides = tuple(
        (mask.lower(), ipaddress.IPv4Address(address))
        for mask, address in override_items
        if address is not None
    )

    exception_items = (
        parser.items("exceptions")
        if parser.has_section("exceptions")
        else ((mask, None) for mask in DEFAULT_EXCEPTIONS)
    )
    exceptions = tuple(mask.lower() for mask, _ in exception_items)

    return redirect_url, allowed_clients, upstreams, timeout, overrides, exceptions


(
    REDIRECT_URL,
    ALLOWED_CLIENTS,
    UPSTREAMS,
    UPSTREAM_TIMEOUT,
    OVERRIDES,
    EXCEPTIONS,
) = load_config(CONFIG_FILE)


def client_allowed(address):
    if not ALLOWED_CLIENTS:
        return True
    try:
        client = ipaddress.ip_address(address)
    except ValueError:
        return False
    return any(client in network for network in ALLOWED_CLIENTS)


def match_any(name, masks):
    return any(fnmatch.fnmatchcase(name, mask) for mask in masks)


def override_for(name):
    for mask, address in OVERRIDES:
        if fnmatch.fnmatchcase(name, mask):
            return address
    return None


def parse_question(packet):
    if len(packet) < 17:
        raise ValueError("short DNS query")

    labels = []
    offset = 12
    while True:
        if offset >= len(packet):
            raise ValueError("unterminated DNS name")
        length = packet[offset]
        offset += 1
        if length == 0:
            break
        if length & 0xC0 or length > 63 or offset + length > len(packet):
            raise ValueError("unsupported DNS name")
        labels.append(packet[offset : offset + length].decode("ascii"))
        offset += length

    if offset + 4 > len(packet):
        raise ValueError("short DNS question")
    qtype, qclass = struct.unpack_from("!HH", packet, offset)
    return ".".join(labels).lower(), qtype, qclass, offset + 4


def error_response(packet, rcode):
    transaction_id, request_flags = struct.unpack_from("!HH", packet, 0)
    response_flags = 0x8000 | (request_flags & 0x0100) | rcode
    try:
        _, _, _, question_end = parse_question(packet)
        question = packet[12:question_end]
        question_count = 1
    except (UnicodeDecodeError, ValueError):
        question = b""
        question_count = 0
    return (
        struct.pack("!HHHHHH", transaction_id, response_flags, question_count, 0, 0, 0)
        + question
    )


def receive_exactly(connection, wanted):
    chunks = []
    received = 0
    while received < wanted:
        chunk = connection.recv(wanted - received)
        if not chunk:
            raise OSError("upstream DNS connection closed early")
        chunks.append(chunk)
        received += len(chunk)
    return b"".join(chunks)


def _forward_one(packet, server, use_tcp):
    family = socket.AF_INET6 if ":" in server else socket.AF_INET
    socket_type = socket.SOCK_STREAM if use_tcp else socket.SOCK_DGRAM
    with socket.socket(family, socket_type) as upstream:
        upstream.settimeout(UPSTREAM_TIMEOUT)
        upstream.connect((server, 53))
        if use_tcp:
            upstream.sendall(struct.pack("!H", len(packet)) + packet)
            response_length = struct.unpack("!H", receive_exactly(upstream, 2))[0]
            response = receive_exactly(upstream, response_length)
        else:
            upstream.send(packet)
            response = upstream.recv(65535)

    if len(response) < 12 or response[:2] != packet[:2]:
        raise ValueError("invalid upstream DNS response")
    return response


def forward_query(packet, use_tcp=False):
    last_error = OSError("no upstream resolver configured")
    for server in UPSTREAMS:
        try:
            return _forward_one(packet, server, use_tcp)
        except (OSError, ValueError) as error:
            last_error = error
    raise last_error


def dns_response(packet, use_tcp=False):
    if len(packet) < 12:
        return b""

    transaction_id, request_flags, question_count = struct.unpack_from("!HHH", packet, 0)
    response_flags = 0x8400 | (request_flags & 0x0100)  # response, authoritative, preserve RD
    try:
        name, qtype, qclass, question_end = parse_question(packet)
        if question_count != 1:
            raise ValueError("expected one DNS question")
    except (UnicodeDecodeError, ValueError):
        return struct.pack("!HHHHHH", transaction_id, response_flags | 1, 0, 0, 0, 0)

    question = packet[12:question_end]

    # Exceptions bypass overrides and are forwarded upstream (nanoDNS semantics).
    if match_any(name, EXCEPTIONS):
        try:
            return forward_query(packet, use_tcp=use_tcp)
        except (OSError, ValueError) as error:
            print(f"DNS upstream failure for {name}: {error}", flush=True)
            return error_response(packet, 2)  # SERVFAIL

    address = override_for(name)
    if address is None:
        # Default-deny: a name matching no rule returns NXDOMAIN.
        return struct.pack("!HHHHHH", transaction_id, response_flags | 3, 1, 0, 0, 0) + question

    if qclass == 1 and qtype in (1, 255):
        answer = (
            b"\xc0\x0c"
            + struct.pack("!HHIH", 1, 1, 60, 4)
            + address.packed
        )
        return struct.pack("!HHHHHH", transaction_id, response_flags, 1, 1, 0, 0) + question + answer

    # NOERROR with no data prevents IPv6 or unsupported record fallback.
    return struct.pack("!HHHHHH", transaction_id, response_flags, 1, 0, 0, 0) + question


class ThreadingUDPServer(socketserver.ThreadingUDPServer):
    allow_reuse_address = True
    daemon_threads = True


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


class UDPDNSHandler(socketserver.BaseRequestHandler):
    def handle(self):
        packet, connection = self.request
        if not client_allowed(self.client_address[0]):
            print(f"DNS UDP {self.client_address[0]} denied", flush=True)
            return
        try:
            name, qtype, _, _ = parse_question(packet)
            print(f"DNS UDP {self.client_address[0]} {name} type={qtype}", flush=True)
        except (UnicodeDecodeError, ValueError):
            print(f"DNS UDP {self.client_address[0]} malformed", flush=True)
        response = dns_response(packet)
        if response:
            connection.sendto(response, self.client_address)


class TCPDNSHandler(socketserver.BaseRequestHandler):
    def handle(self):
        if not client_allowed(self.client_address[0]):
            print(f"DNS TCP {self.client_address[0]} denied", flush=True)
            return
        length_bytes = self._receive_exactly(2)
        if length_bytes is None:
            return
        wanted = struct.unpack("!H", length_bytes)[0]
        packet = self._receive_exactly(wanted)
        if packet is None:
            return
        try:
            name, qtype, _, _ = parse_question(packet)
            print(f"DNS TCP {self.client_address[0]} {name} type={qtype}", flush=True)
        except (UnicodeDecodeError, ValueError):
            print(f"DNS TCP {self.client_address[0]} malformed", flush=True)
        response = dns_response(packet, use_tcp=True)
        if response:
            self.request.sendall(struct.pack("!H", len(response)) + response)

    def _receive_exactly(self, wanted):
        try:
            return receive_exactly(self.request, wanted)
        except OSError:
            return None


class RedirectHandler(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        if self.path == "/_health":
            body = b"ok\n"
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        self.send_response(302)
        self.send_header("Location", REDIRECT_URL)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_HEAD(self):
        self.send_response(302)
        self.send_header("Location", REDIRECT_URL)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, message, *args):
        user_agent = self.headers.get("User-Agent", "-")
        print(
            f"HTTP {self.client_address[0]} {message % args} user-agent={user_agent!r}",
            flush=True,
        )


class HTTPServer(http.server.ThreadingHTTPServer):
    allow_reuse_address = True
    daemon_threads = True


def serve(server, name):
    thread = threading.Thread(target=server.serve_forever, name=name, daemon=True)
    thread.start()
    return thread


if __name__ == "__main__":
    dns_udp = ThreadingUDPServer(("0.0.0.0", DNS_PORT), UDPDNSHandler)
    dns_tcp = ThreadingTCPServer(("0.0.0.0", DNS_PORT), TCPDNSHandler)
    http = HTTPServer(("0.0.0.0", HTTP_PORT), RedirectHandler)
    https = HTTPServer(("0.0.0.0", HTTPS_PORT), RedirectHandler)

    tls = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    tls.load_cert_chain(CERTIFICATE, PRIVATE_KEY)
    https.socket = tls.wrap_socket(https.socket, server_side=True)

    servers = [dns_udp, dns_tcp, http, https]
    serve(dns_udp, "dns-udp")
    serve(dns_tcp, "dns-tcp")
    serve(http, "http-redirect")
    serve(https, "https-redirect")

    overrides_summary = (
        ", ".join(f"{mask}->{address}" for mask, address in OVERRIDES) or "none"
    )
    print(
        "PS5 guide gateway (nanoDNS-style, default-deny): "
        f"redirect -> {REDIRECT_URL}; overrides {overrides_summary}; "
        f"forward {','.join(EXCEPTIONS) or 'none'} -> {','.join(UPSTREAMS) or 'none'}; "
        f"clients {','.join(str(network) for network in ALLOWED_CLIENTS) or 'any'}",
        flush=True,
    )

    stopped = threading.Event()

    def stop(*_):
        stopped.set()

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    stopped.wait()
    for server in servers:
        server.shutdown()
        server.server_close()