#!/usr/bin/env python3
"""
Mini SSRF *shape* lab helper for Tech Forge curl course.
Exposes GET /fetch?url=... that retrieves a URL — but ONLY allows
httpbin.org and localhost/127.0.0.1. Binds to 127.0.0.1 only.

Usage:
  python3 mini-ssrf.py
  curl -s 'http://127.0.0.1:9091/fetch?url=https://httpbin.org/get'
  curl -s 'http://127.0.0.1:9091/fetch?url=http://127.0.0.1:9090/'
"""
from __future__ import annotations

import json
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, HTTPServer

HOST = "127.0.0.1"
PORT = 9091
ALLOW_HOSTS = {"httpbin.org", "localhost", "127.0.0.1"}
MAX_BYTES = 64_000


def allowed(url: str) -> tuple[bool, str]:
    try:
        p = urllib.parse.urlparse(url)
    except Exception:
        return False, "unparseable URL"
    if p.scheme not in ("http", "https"):
        return False, "only http/https allowed"
    host = (p.hostname or "").lower()
    if host not in ALLOW_HOSTS:
        return False, f"host '{host}' not in allowlist {sorted(ALLOW_HOSTS)}"
    return True, "ok"


class Handler(BaseHTTPRequestHandler):
    server_version = "MiniSSRF/1.0"

    def do_GET(self) -> None:  # noqa: N802
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/":
            body = (
                "mini-ssrf lab helper — Tech Forge\n"
                "GET /fetch?url=https://httpbin.org/get\n"
                "Allowlist: httpbin.org, localhost, 127.0.0.1\n"
            )
            self._send(200, body, "text/plain")
            return

        if parsed.path != "/fetch":
            self._send(404, "Not found. Use /fetch?url=...\n", "text/plain")
            return

        qs = urllib.parse.parse_qs(parsed.query)
        url = (qs.get("url") or [""])[0]
        if not url:
            self._send(400, json.dumps({"error": "missing url param"}) + "\n", "application/json")
            return

        ok, reason = allowed(url)
        if not ok:
            self._send(
                403,
                json.dumps({"error": "blocked", "reason": reason, "url": url}, indent=2) + "\n",
                "application/json",
            )
            return

        try:
            req = urllib.request.Request(url, headers={"User-Agent": "TechForge-mini-ssrf/1.0"})
            with urllib.request.urlopen(req, timeout=8) as resp:
                raw = resp.read(MAX_BYTES)
                ctype = resp.headers.get("Content-Type", "application/octet-stream")
                out = {
                    "fetched": url,
                    "status": resp.status,
                    "content_type": ctype,
                    "body_preview": raw.decode("utf-8", errors="replace")[:4000],
                }
                self._send(200, json.dumps(out, indent=2) + "\n", "application/json")
        except urllib.error.HTTPError as e:
            self._send(
                200,
                json.dumps({"fetched": url, "status": e.code, "error": str(e)}, indent=2) + "\n",
                "application/json",
            )
        except Exception as e:
            self._send(
                502,
                json.dumps({"fetched": url, "error": str(e)}, indent=2) + "\n",
                "application/json",
            )

    def log_message(self, fmt: str, *args) -> None:
        print(f"[mini-ssrf] {self.address_string()} - {fmt % args}")

    def _send(self, status: int, body: str, ctype: str) -> None:
        data = body.encode()
        self.send_response(status)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("X-Lab-Helper", "techforge-mini-ssrf")
        self.end_headers()
        self.wfile.write(data)


def main() -> None:
    if HOST not in ("127.0.0.1", "localhost", "::1"):
        raise SystemExit("Refusing to bind outside loopback.")
    httpd = HTTPServer((HOST, PORT), Handler)
    print(f"mini-ssrf listening on http://{HOST}:{PORT}/ (Ctrl+C to stop)")
    print(f"  curl -s 'http://{HOST}:{PORT}/fetch?url=https://httpbin.org/get'")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.")


if __name__ == "__main__":
    main()