import http.server
import os
import posixpath
import ssl
import urllib.parse


class UAPrefixRoutingHandler(http.server.SimpleHTTPRequestHandler):
    # Match on the first N characters of User-Agent.
    UA_PREFIX_LEN = 5

    # Edit this map as needed:
    # prefix -> {"subdir": directory_to_serve, "status": HTTP_status_for_success}
    UA_ROUTE_MAP = {
        "Mozil": {"subdir": "site_mozilla", "status": 266},
        "Claud": {"subdir": "site_claude", "status": 266},
        "curl/": {"subdir": "site_curl", "status": 200},
     }

    DEFAULT_SUBDIR = "site_default"
    DEFAULT_STATUS = 200

    def _selected_route(self):
        user_agent = self.headers.get("User-Agent", "")
        prefix = user_agent[: self.UA_PREFIX_LEN]
        return self.UA_ROUTE_MAP.get(
            prefix,
            {"subdir": self.DEFAULT_SUBDIR, "status": self.DEFAULT_STATUS},
        )

    def _selected_subdir(self):
        return self._selected_route()["subdir"]

    def _selected_status(self):
        return self._selected_route()["status"]

    def send_response(self, code, message=None):
        if code == http.HTTPStatus.OK:
            code = self._selected_status()
        super().send_response(code, message)

    def translate_path(self, path):
        # Same idea as SimpleHTTPRequestHandler.translate_path, but rooted
        # in a user-agent-specific subdirectory.
        path = path.split("?", 1)[0]
        path = path.split("#", 1)[0]
        trailing_slash = path.endswith("/")
        path = posixpath.normpath(urllib.parse.unquote(path))
        parts = [part for part in path.split("/") if part]

        translated = os.path.join(os.getcwd(), self._selected_subdir())
        for part in parts:
            part = os.path.basename(part)
            if part in (os.curdir, os.pardir):
                continue
            translated = os.path.join(translated, part)

        if trailing_slash:
            translated += "/"
        return translated

    def log_request(self, code="-", size="-"):
        super().log_request(code, size)
        try:
            status_code = int(code)
        except (TypeError, ValueError):
            return

        if 200 <= status_code < 300:
            user_agent = self.headers.get("User-Agent", "")
            print(f"Successful request User-Agent: {user_agent}", flush=True)

# Define server address and handler
server_address = ('0.0.0.0', 443)

handler = UAPrefixRoutingHandler

# Create the HTTP server
httpd = http.server.HTTPServer(server_address, handler)

# Set up SSL context (modern way)
context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
context.load_cert_chain(certfile='ctf28.com.crt', keyfile='ctf28.com.key')

# Wrap the socket
httpd.socket = context.wrap_socket(httpd.socket, server_side=True)

print("Serving on port 443")
httpd.serve_forever()