/pentestgpt_agent/tests/support/local_target.py

from __future__ import annotations

import re
import secrets
import threading
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs


@dataclass(frozen=True)
class LocalTarget:
    url: str
    flag: str


_TEMPLATE = re.compile(r"^\{\{\s*(?P<expression>.*?)\s*\}\}$")
_ARITHMETIC = re.compile(r"^(?P<left>-?\d{1,6})\s*(?P<operator>[+*-])\s*(?P<right>-?\d{1,6})$")


def _render_template(value: str, flag: str) -> str:
    template = _TEMPLATE.fullmatch(value)
    if template is None:
        return value
    expression = template.group("expression")
    if expression == "config":
        return f"{{'FLAG': '{flag}'}}"
    if expression in {"config.FLAG", "config['FLAG']", 'config["FLAG"]'}:
        return flag
    arithmetic = _ARITHMETIC.fullmatch(expression)
    if arithmetic is None:
        return value
    left = int(arithmetic.group("left"))
    right = int(arithmetic.group("right"))
    return str(
        {
            "+": left + right,
            "-": left - right,
            "*": left * right,
        }[arithmetic.group("operator")]
    )


def _handler(flag: str, *, vulnerable: bool) -> type[BaseHTTPRequestHandler]:
    class ApplicationHandler(BaseHTTPRequestHandler):
        def log_message(self, *args: object) -> None:
            pass

        def _respond(self, body: str) -> None:
            encoded = body.encode()
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.send_header("Content-Length", str(len(encoded)))
            self.end_headers()
            self.wfile.write(encoded)

        def do_GET(self) -> None:
            if self.path == "/":
                self._respond('<a href="/apply">Driver application</a>')
            elif self.path == "/apply":
                self._respond('<form method="post"><input name="name"></form>')
            else:
                self.send_error(404)

        def do_POST(self) -> None:
            length = int(self.headers.get("Content-Length", "0"))
            form = parse_qs(self.rfile.read(length).decode())
            name = form.get("name", [""])[0]
            rendered = _render_template(name, flag) if vulnerable else name
            self._respond(f"Application received for {rendered}")

    return ApplicationHandler


@contextmanager
def local_template_target(*, vulnerable: bool = True) -> Iterator[LocalTarget]:
    flag = f"FLAG{{{secrets.token_hex(16)}}}"
    server = ThreadingHTTPServer(("127.0.0.1", 0), _handler(flag, vulnerable=vulnerable))
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        yield LocalTarget(
            url=f"http://127.0.0.1:{server.server_address[1]}",
            flag=flag,
        )
    finally:
        server.shutdown()
        server.server_close()
        thread.join(timeout=5)


__all__ = ["LocalTarget", "local_template_target"]