from app.models.template import ProxyTemplateDocument

DOCS_TEMPLATES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"


def builtin_template(
    *,
    id: str,
    name: str,
    protocol: str,
    family: str,
    tags: list[str],
    action: str,
    summary: str,
    behavior: str,
    code: str,
) -> ProxyTemplateDocument:
    category = protocol if protocol in {"http", "tcp"} else family
    return ProxyTemplateDocument(
        id=id,
        name=name,
        description=summary,
        category=category,
        protocol=protocol,
        family=family,
        tags=tags,
        action=action,
        summary=summary,
        behavior=behavior,
        source_url=DOCS_TEMPLATES_URL,
        built_in=True,
        code=code,
    )


BUILTIN_TEMPLATES = [
    builtin_template(
        id="http-block-path-traversal",
        name="Block path traversal",
        protocol="http",
        family="web-hardening",
        tags=["path-traversal", "web", "block"],
        action="block",
        summary="Stops common path traversal probes in path and query.",
        behavior=(
            "Logs the matched traversal token and immediately drops the client "
            "connection before the request reaches the upstream service."
        ),
        code="""import logging
from urllib.parse import unquote

from mitmproxy import http

TRAVERSAL_MARKERS = (
    "../",
    "..\\\\",
    "%2e%2e",
    "%252e%252e",
    "..%2f",
    "..%5c",
)


def request(flow: http.HTTPFlow) -> None:
    raw_target = (
        f"{flow.request.path}?{flow.request.query}"
        if flow.request.query
        else flow.request.path
    )
    normalized = unquote(unquote(raw_target)).lower()
    for marker in TRAVERSAL_MARKERS:
        if marker in normalized:
            logging.warning(
                "[http-block-path-traversal] blocked %s %s from %s",
                flow.request.method,
                flow.request.pretty_url,
                flow.client_conn.peername,
            )
            flow.kill()
            return
""",
    ),
    builtin_template(
        id="http-block-blacklisted-substrings",
        name="Block blacklisted substrings",
        protocol="http",
        family="web-hardening",
        tags=["blacklist-string", "web", "block"],
        action="block",
        summary="Blocks requests containing forbidden strings in path, query, or body.",
        behavior=(
            "Scans URL and request body for blacklisted markers, then serves a "
            "blank HTML page with 200 OK and logs the first hit."
        ),
        code="""import logging

from mitmproxy import http

BLOCKED_STRINGS = (
    "flag{",
    "give_me_shell",
    "debug=1",
    "/etc/passwd",
)


def request(flow: http.HTTPFlow) -> None:
    body = flow.request.get_text(strict=False)
    haystack = "\\n".join(
        [flow.request.path, str(flow.request.query), body]
    ).lower()
    for marker in BLOCKED_STRINGS:
        if marker.lower() in haystack:
            logging.warning(
                "[http-block-blacklisted-substrings] blocked marker=%r url=%s",
                marker,
                flow.request.pretty_url,
            )
            flow.response = http.Response.make(
                200,
                b"",
                {"content-type": "text/html; charset=utf-8"},
            )
            return
""",
    ),
    builtin_template(
        id="http-block-large-body",
        name="Block large request bodies",
        protocol="http",
        family="web-hardening",
        tags=["length-limit", "body-size", "web", "block"],
        action="block",
        summary="Rejects oversized HTTP request bodies before they hit the service.",
        behavior=(
            "Checks the buffered request body length, returns 413 Payload Too "
            "Large, and logs the observed size."
        ),
        code="""import logging

from mitmproxy import http

MAX_BODY_BYTES = 8192


def request(flow: http.HTTPFlow) -> None:
    size = len(flow.request.raw_content or b"")
    if size > MAX_BODY_BYTES:
        logging.warning(
            "[http-block-large-body] blocked %s bytes on %s",
            size,
            flow.request.pretty_url,
        )
        flow.response = http.Response.make(
            413,
            b"",
            {"content-type": "text/plain"},
        )
""",
    ),
    builtin_template(
        id="http-block-suspicious-methods",
        name="Block suspicious HTTP methods",
        protocol="http",
        family="web-hardening",
        tags=["method-filter", "web", "block"],
        action="block",
        summary="Denies uncommon or dangerous HTTP methods.",
        behavior=(
            "Returns 405 Method Not Allowed for methods outside the allowlist "
            "and logs the blocked method and target URL."
        ),
        code="""import logging

from mitmproxy import http

ALLOWED_METHODS = {"GET", "POST", "HEAD"}


def request(flow: http.HTTPFlow) -> None:
    if flow.request.method not in ALLOWED_METHODS:
        logging.warning(
            "[http-block-suspicious-methods] blocked method=%s url=%s",
            flow.request.method,
            flow.request.pretty_url,
        )
        flow.response = http.Response.make(
            405,
            b"",
            {"allow": "GET, POST, HEAD"},
        )
""",
    ),
    builtin_template(
        id="http-block-admin-paths",
        name="Block admin and debug paths",
        protocol="http",
        family="web-hardening",
        tags=["admin-path", "debug-path", "web", "block"],
        action="block",
        summary="Protects common admin, console, and debug endpoints.",
        behavior=(
            "Matches the request path against a small denylist and returns a "
            "blank 404 page before the request is forwarded."
        ),
        code="""import logging

from mitmproxy import http

BLOCKED_PREFIXES = (
    "/admin",
    "/debug",
    "/console",
    "/actuator",
)


def request(flow: http.HTTPFlow) -> None:
    if flow.request.path.startswith(BLOCKED_PREFIXES):
        logging.warning(
            "[http-block-admin-paths] blocked path=%s",
            flow.request.path,
        )
        flow.response = http.Response.make(
            404,
            b"",
            {"content-type": "text/html; charset=utf-8"},
        )
""",
    ),
    builtin_template(
        id="http-sanitize-headers-basic",
        name="Sanitize risky forwarding headers",
        protocol="http",
        family="web-hardening",
        tags=["rewrite", "headers", "web"],
        action="rewrite",
        summary=(
            "Removes a few risky forwarding headers before the request is "
            "sent upstream."
        ),
        behavior=(
            "Deletes selected override or forwarding headers and logs what was "
            "rewritten without blocking the request."
        ),
        code="""import logging

from mitmproxy import http

STRIP_HEADERS = (
    "x-forwarded-host",
    "x-http-method-override",
    "x-original-url",
    "x-rewrite-url",
)


def request(flow: http.HTTPFlow) -> None:
    removed = []
    for header in STRIP_HEADERS:
        if header in flow.request.headers:
            removed.append(header)
            del flow.request.headers[header]
    if removed:
        logging.info(
            "[http-sanitize-headers-basic] stripped headers=%s url=%s",
            ",".join(removed),
            flow.request.pretty_url,
        )
""",
    ),
    builtin_template(
        id="http-block-command-injection-markers",
        name="Block command injection markers",
        protocol="http",
        family="web-hardening",
        tags=["command-injection", "web", "block"],
        action="block",
        summary=(
            "Rejects obvious shell metacharacters and command-substitution "
            "markers."
        ),
        behavior=(
            "Scans URL and body for shell-style injection markers, logs the "
            "first hit, and drops the client connection."
        ),
        code="""import logging

from mitmproxy import http

BLOCKED_MARKERS = (";", "&&", "||", "$(", "`", "|sh", "|bash")


def request(flow: http.HTTPFlow) -> None:
    body = flow.request.get_text(strict=False)
    haystack = "\\n".join(
        [flow.request.pretty_url, body]
    ).lower()
    for marker in BLOCKED_MARKERS:
        if marker in haystack:
            logging.warning(
                "[http-block-command-injection-markers] blocked marker=%r url=%s",
                marker,
                flow.request.pretty_url,
            )
            flow.kill()
            return
""",
    ),
    builtin_template(
        id="http-block-sqli-markers",
        name="Block obvious SQLi markers",
        protocol="http",
        family="web-hardening",
        tags=["sqli", "web", "block"],
        action="block",
        summary="Stops common SQL injection probes before they reach the app.",
        behavior=(
            "Matches a short list of high-signal SQLi markers in path, query, "
            "and body, then serves a blank HTML page and logs the hit."
        ),
        code="""import logging

from mitmproxy import http

BLOCKED_MARKERS = (
    "' or 1=1",
    '" or 1=1',
    "union select",
    "sleep(",
    "information_schema",
    "--",
)


def request(flow: http.HTTPFlow) -> None:
    body = flow.request.get_text(strict=False)
    haystack = "\\n".join(
        [flow.request.path, str(flow.request.query), body]
    ).lower()
    for marker in BLOCKED_MARKERS:
        if marker in haystack:
            logging.warning(
                "[http-block-sqli-markers] blocked marker=%r url=%s",
                marker,
                flow.request.pretty_url,
            )
            flow.response = http.Response.make(
                200,
                b"",
                {"content-type": "text/html; charset=utf-8"},
            )
            return
""",
    ),
    builtin_template(
        id="tcp-block-repeated-byte-threshold",
        name="Block repeated-byte flood",
        protocol="tcp",
        family="tcp-hardening",
        tags=["bof", "repeated-byte", "tcp", "block"],
        action="block",
        summary="Closes connections that send suspiciously long runs of the same byte.",
        behavior=(
            "Measures the longest repeated-byte run in each client message, logs "
            "the byte and run length, and closes the TCP flow on threshold hit."
        ),
        code="""import logging

from mitmproxy import tcp

MAX_REPEAT = 128


def longest_run(data: bytes) -> tuple[int, int]:
    best_byte = -1
    best_run = 0
    current_byte = -1
    current_run = 0
    for byte in data:
        if byte == current_byte:
            current_run += 1
        else:
            current_byte = byte
            current_run = 1
        if current_run > best_run:
            best_byte = byte
            best_run = current_run
    return best_byte, best_run


def tcp_message(flow: tcp.TCPFlow) -> None:
    message = flow.messages[-1]
    if not message.from_client:
        return
    repeated_byte, run_length = longest_run(message.content)
    if run_length > MAX_REPEAT:
        logging.warning(
            "[tcp-block-repeated-byte-threshold] blocked byte=0x%02x run=%d",
            repeated_byte,
            run_length,
        )
        message.drop()
        flow.kill()
""",
    ),
    builtin_template(
        id="tcp-block-message-length-threshold",
        name="Block oversized TCP messages",
        protocol="tcp",
        family="tcp-hardening",
        tags=["length-limit", "tcp", "block"],
        action="block",
        summary="Cuts off client messages that exceed a fixed size threshold.",
        behavior=(
            "Checks every client TCP message size, logs the offending length, "
            "drops the message, and kills the flow when the threshold is hit."
        ),
        code="""import logging

from mitmproxy import tcp

MAX_MESSAGE_BYTES = 4096


def tcp_message(flow: tcp.TCPFlow) -> None:
    message = flow.messages[-1]
    if not message.from_client:
        return
    size = len(message.content)
    if size > MAX_MESSAGE_BYTES:
        logging.warning(
            "[tcp-block-message-length-threshold] blocked message size=%d",
            size,
        )
        message.drop()
        flow.kill()
""",
    ),
    builtin_template(
        id="tcp-block-blacklisted-byte-sequences",
        name="Block blacklisted TCP byte sequences",
        protocol="tcp",
        family="tcp-hardening",
        tags=["blacklist-bytes", "tcp", "block"],
        action="block",
        summary=(
            "Blocks payloads containing forbidden byte sequences or marker "
            "strings."
        ),
        behavior=(
            "Looks for dangerous byte patterns in client messages, then drops "
            "the message and closes the flow."
        ),
        code="""import logging

from mitmproxy import tcp

BLOCKED_SEQUENCES = (
    b"/bin/sh",
    b"cat flag",
    b"AAAA%p%p",
    b"\\x90\\x90\\x90\\x90",
)


def tcp_message(flow: tcp.TCPFlow) -> None:
    message = flow.messages[-1]
    if not message.from_client:
        return
    for sequence in BLOCKED_SEQUENCES:
        if sequence in message.content:
            logging.warning(
                "[tcp-block-blacklisted-byte-sequences] blocked sequence=%r",
                sequence,
            )
            message.drop()
            flow.kill()
            return
""",
    ),
    builtin_template(
        id="tcp-block-non-printable-threshold",
        name="Block binary-noise threshold",
        protocol="tcp",
        family="tcp-hardening",
        tags=["non-printable", "binary-noise", "tcp", "block"],
        action="block",
        summary=(
            "Blocks client messages with an excessive amount of non-printable "
            "bytes."
        ),
        behavior=(
            "Treats unusually binary-heavy client payloads as suspicious, logs "
            "the ratio, and closes the flow when the threshold is exceeded."
        ),
        code="""import logging

from mitmproxy import tcp

NON_PRINTABLE_RATIO = 0.35


def tcp_message(flow: tcp.TCPFlow) -> None:
    message = flow.messages[-1]
    if not message.from_client or not message.content:
        return
    non_printable = sum(
        1
        for byte in message.content
        if byte < 0x20 and byte not in {0x09, 0x0a, 0x0d}
    )
    ratio = non_printable / len(message.content)
    if ratio > NON_PRINTABLE_RATIO:
        logging.warning(
            "[tcp-block-non-printable-threshold] blocked ratio=%.2f size=%d",
            ratio,
            len(message.content),
        )
        message.drop()
        flow.kill()
""",
    ),
]