RustProxy/tools/dev/mock_auth.py

#!/usr/bin/env python3
"""LOOPBACK-ONLY fake Tinyauth protocol peer. NEVER use this as real authentication.
A cookie named mock_auth with value allow approves a request. Any client can mint it.
This exists to test RustProxy auth_request, identity forwarding and fail-closed paths.
"""
from __future__ import annotations
import argparse
from collections import deque
from http.cookies import SimpleCookie
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import ipaddress
import json
import threading
import time
from urllib.parse import urlsplit

class State:
    def __init__(self):
        self.lock = threading.Lock()
        self.status = None
        self.delay_ms = 0
        self.location = None
        self.calls = 0
        self.recent = deque(maxlen=16)

class Server(ThreadingHTTPServer):
    daemon_threads = True
    def __init__(self, address, state=None):
        if not ipaddress.ip_address(address[0]).is_loopback:
            raise ValueError("The fake auth server may bind only a numeric loopback address")
        self.state = state or State()
        super().__init__(address, Handler)

class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "RustProxy-Mock-Auth/0.4.0"
    sys_version = ""
    def log_message(self, *_):
        pass
    def setup(self):
        super().setup()
        self.connection.settimeout(5)
    def reply(self, status, value=None, headers=()):
        data = b"" if value is None else json.dumps(value, ensure_ascii=True).encode()
        self.send_response(status)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Type", "application/json")
        for key, value in headers:
            self.send_header(key, value)
        self.end_headers()
        if self.command != "HEAD":
            try:
                self.wfile.write(data)
            except (BrokenPipeError, ConnectionResetError):
                pass
    def do_GET(self):
        path = urlsplit(self.path).path
        if self.headers.get("Transfer-Encoding") or self.headers.get("Content-Length", "0") != "0":
            self.close_connection = True
            return self.reply(400, {"error":"GET verification requests must have no body"})
        state = self.server.state
        if path == "/api/auth/nginx":
            cookie = SimpleCookie()
            try:
                cookie.load(self.headers.get("Cookie", ""))
            except Exception:
                pass
            allowed = "mock_auth" in cookie and cookie["mock_auth"].value == "allow"
            with state.lock:
                state.calls += 1
                # Do not store Cookie, Authorization, API keys, or query strings.
                state.recent.append({"method":self.headers.get("X-Original-Method"),
                    "client":self.headers.get("X-Real-IP"), "host":self.headers.get("X-Forwarded-Host"),
                    "proto":self.headers.get("X-Forwarded-Proto"),
                    "forwarded":self.headers.get("Forwarded"),"xff":self.headers.get("X-Forwarded-For")})
                status = state.status if state.status is not None else (204 if allowed else 401)
                delay_ms, location = state.delay_ms, state.location
            time.sleep(delay_ms / 1000)
            headers = []
            if 200 <= status < 300:
                headers = [("Remote-User", "mock-user"), ("Remote-Email", "mock@example.test"), ("Remote-Groups", "testers")]
            if location:
                headers.append(("X-Tinyauth-Location", location))
            return self.reply(status, headers=headers)
        if path == "/state":
            with state.lock:
                return self.reply(200, {"calls":state.calls,"status":state.status,"delay_ms":state.delay_ms,"recent":list(state.recent)})
        if path == "/login":
            return self.reply(200, {"warning":"Fake login fixture, not authentication", "test_cookie":"mock_auth=allow", "instructions":"Send this cookie to the protected local test vhost."})
        self.reply(404, {"error":"not found"})
    def do_POST(self):
        if urlsplit(self.path).path != "/control":
            self.close_connection = True
            return self.reply(405, {"error":"only POST /control is supported"})
        try:
            if self.headers.get("Transfer-Encoding") or len(self.headers.get_all("Content-Length", [])) != 1:
                raise ValueError("one Content-Length is required")
            size = int(self.headers["Content-Length"])
            if not 0 <= size <= 4096:
                raise ValueError("body limit is 4096 bytes")
            data = self.rfile.read(size)
            if len(data) != size:
                raise ValueError("incomplete request")
            payload = json.loads(data)
            if not isinstance(payload, dict) or set(payload) - {"status", "delay_ms", "location", "reset"}:
                raise ValueError("invalid fields")
            status = payload.get("status")
            if status is not None and (type(status) is not int or status not in (200, 204, 401, 403, 500, 503)):
                raise ValueError("unsupported verification status")
            delay = payload.get("delay_ms", 0)
            if type(delay) is not int or not 0 <= delay <= 15000:
                raise ValueError("delay must be 0..15000 ms")
            location = payload.get("location")
            if location is not None and (not isinstance(location, str) or len(location) > 2048
                    or any(ord(c) < 32 or ord(c) > 126 for c in location)):
                raise ValueError("invalid redirect header")
            with self.server.state.lock:
                state = self.server.state
                if payload.get("reset"):
                    state.calls = 0
                    state.recent.clear()
                state.status, state.delay_ms, state.location = status, delay, location
        except (ValueError, KeyError, TimeoutError, OSError) as error:
            self.close_connection = True
            return self.reply(400, {"error":str(error)})
        self.reply(200, {"updated":True})
    def do_HEAD(self):
        self.do_GET()

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", type=int, default=9300)
    args = parser.parse_args()
    if not 1 <= args.port <= 65535:
        parser.error("port must be 1..65535")
    server = Server(("127.0.0.1", args.port))
    print(f"FAKE AUTH - LOCAL TESTING ONLY: http://127.0.0.1:{args.port}/api/auth/nginx", flush=True)
    print("Approve with Cookie: mock_auth=allow. /control configures failures/delays. Ctrl+C stops.", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
if __name__ == "__main__":
    main()