491 lines
16 KiB
Python
491 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Protocol-level OAuth, push, and SMTP mocks for the isolated local drill."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import signal
|
|
import socketserver
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from typing import Any
|
|
|
|
|
|
HTTP_PORT = 8080
|
|
SMTP_PORT = 2525
|
|
SMTP_RETRY_PORT = 2526
|
|
SMTP_TIMEOUT_PORT = 2527
|
|
DELAY_SECONDS = 0.35
|
|
|
|
|
|
class BoundaryState:
|
|
def __init__(self) -> None:
|
|
self.lock = threading.Lock()
|
|
self.oauth_mode = "success"
|
|
self.oauth_authorize_requests = 0
|
|
self.oauth_token_requests = 0
|
|
self.oauth_user_requests = 0
|
|
self.oauth_code_counter = 0
|
|
self.oauth_codes: dict[str, dict[str, Any]] = {}
|
|
self.oauth_tokens: set[str] = set()
|
|
self.push_mode = "success"
|
|
self.push_attempts = 0
|
|
self.push_receipts: dict[str, str] = {}
|
|
self.push_deliveries = 0
|
|
self.push_notifications: list[dict[str, Any]] = []
|
|
self.smtp_connections = {
|
|
str(SMTP_PORT): 0,
|
|
str(SMTP_RETRY_PORT): 0,
|
|
str(SMTP_TIMEOUT_PORT): 0,
|
|
}
|
|
self.smtp_messages = 0
|
|
self.smtp_rejections = 0
|
|
self.smtp_digests: list[str] = []
|
|
|
|
def control(self, component: str, mode: str) -> None:
|
|
with self.lock:
|
|
if component == "oauth":
|
|
self.oauth_mode = mode
|
|
self.oauth_authorize_requests = 0
|
|
self.oauth_token_requests = 0
|
|
self.oauth_user_requests = 0
|
|
self.oauth_codes = {}
|
|
self.oauth_tokens = set()
|
|
elif component == "push":
|
|
self.push_mode = mode
|
|
self.push_attempts = 0
|
|
self.push_receipts = {}
|
|
self.push_deliveries = 0
|
|
self.push_notifications = []
|
|
elif component == "smtp":
|
|
self.smtp_connections = {
|
|
str(SMTP_PORT): 0,
|
|
str(SMTP_RETRY_PORT): 0,
|
|
str(SMTP_TIMEOUT_PORT): 0,
|
|
}
|
|
self.smtp_messages = 0
|
|
self.smtp_rejections = 0
|
|
self.smtp_digests = []
|
|
else:
|
|
raise ValueError("unsupported component")
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
with self.lock:
|
|
return {
|
|
"oauth": {
|
|
"mode": self.oauth_mode,
|
|
"authorize_requests": self.oauth_authorize_requests,
|
|
"token_requests": self.oauth_token_requests,
|
|
"user_requests": self.oauth_user_requests,
|
|
"issued_codes": len(self.oauth_codes),
|
|
"consumed_codes": sum(
|
|
1 for code in self.oauth_codes.values() if code["used"]
|
|
),
|
|
},
|
|
"push": {
|
|
"mode": self.push_mode,
|
|
"attempts": self.push_attempts,
|
|
"deliveries": self.push_deliveries,
|
|
"idempotency_keys": len(self.push_receipts),
|
|
"notifications": list(self.push_notifications),
|
|
},
|
|
"smtp": {
|
|
"connections": dict(self.smtp_connections),
|
|
"messages": self.smtp_messages,
|
|
"rejections": self.smtp_rejections,
|
|
"message_digests": list(self.smtp_digests),
|
|
},
|
|
}
|
|
|
|
|
|
STATE = BoundaryState()
|
|
OAUTH_CLIENT_ID = os.environ["MOCK_OAUTH_CLIENT_ID"]
|
|
OAUTH_CLIENT_SECRET = os.environ["MOCK_OAUTH_CLIENT_SECRET"]
|
|
PUSH_BEARER_TOKEN = os.environ["MOCK_PUSH_BEARER_TOKEN"]
|
|
|
|
|
|
def json_bytes(value: Any) -> bytes:
|
|
return json.dumps(value, separators=(",", ":"), sort_keys=True).encode()
|
|
|
|
|
|
class BoundaryHTTPServer(ThreadingHTTPServer):
|
|
daemon_threads = True
|
|
|
|
|
|
class BoundaryHTTPHandler(BaseHTTPRequestHandler):
|
|
server_version = "WhoNeedHelpBoundaryMock/1"
|
|
|
|
def log_message(self, _format: str, *_args: Any) -> None:
|
|
return
|
|
|
|
def send_json(self, status: int, value: Any) -> None:
|
|
payload = json_bytes(value)
|
|
try:
|
|
self.send_response(status)
|
|
self.send_header("content-type", "application/json")
|
|
self.send_header("content-length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
return
|
|
|
|
def read_body(self) -> bytes:
|
|
length = int(self.headers.get("content-length", "0"))
|
|
if length < 0 or length > 65_536:
|
|
raise ValueError("invalid body length")
|
|
return self.rfile.read(length)
|
|
|
|
def do_GET(self) -> None: # noqa: N802
|
|
parsed = urllib.parse.urlsplit(self.path)
|
|
|
|
if parsed.path == "/healthz":
|
|
self.send_json(200, {"status": "ok"})
|
|
elif parsed.path == "/state":
|
|
self.send_json(200, STATE.snapshot())
|
|
elif parsed.path == "/oauth/authorize":
|
|
self.oauth_authorize(parsed)
|
|
elif parsed.path == "/oauth/user":
|
|
self.oauth_user()
|
|
else:
|
|
self.send_json(404, {"error": "not_found"})
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
parsed = urllib.parse.urlsplit(self.path)
|
|
|
|
try:
|
|
body = self.read_body()
|
|
except ValueError:
|
|
self.send_json(413, {"error": "invalid_body"})
|
|
return
|
|
|
|
if parsed.path == "/control":
|
|
self.control(body)
|
|
elif parsed.path == "/oauth/token":
|
|
self.oauth_token(body)
|
|
elif parsed.path == "/push":
|
|
self.push(body)
|
|
else:
|
|
self.send_json(404, {"error": "not_found"})
|
|
|
|
def control(self, body: bytes) -> None:
|
|
try:
|
|
command = json.loads(body)
|
|
component = command["component"]
|
|
mode = command.get("mode", "success")
|
|
STATE.control(component, mode)
|
|
except (json.JSONDecodeError, KeyError, TypeError, ValueError):
|
|
self.send_json(400, {"error": "invalid_control"})
|
|
return
|
|
|
|
self.send_json(200, {"status": "reset", "component": component, "mode": mode})
|
|
|
|
def oauth_authorize(self, parsed: urllib.parse.SplitResult) -> None:
|
|
params = urllib.parse.parse_qs(parsed.query)
|
|
required = {
|
|
"client_id",
|
|
"redirect_uri",
|
|
"state",
|
|
"code_challenge",
|
|
"code_challenge_method",
|
|
}
|
|
|
|
if not required.issubset(params) or params["client_id"][0] != OAUTH_CLIENT_ID:
|
|
self.send_json(400, {"error": "invalid_authorization_request"})
|
|
return
|
|
|
|
if params["code_challenge_method"][0] != "S256":
|
|
self.send_json(400, {"error": "unsupported_challenge_method"})
|
|
return
|
|
|
|
redirect_uri = params["redirect_uri"][0]
|
|
state = params["state"][0]
|
|
|
|
with STATE.lock:
|
|
STATE.oauth_authorize_requests += 1
|
|
mode = STATE.oauth_mode
|
|
|
|
if mode != "deny_authorize":
|
|
STATE.oauth_code_counter += 1
|
|
code = f"boundary-code-{STATE.oauth_code_counter}"
|
|
STATE.oauth_codes[code] = {
|
|
"challenge": params["code_challenge"][0],
|
|
"redirect_uri": redirect_uri,
|
|
"used": False,
|
|
}
|
|
|
|
if mode == "deny_authorize":
|
|
callback_params = {
|
|
"error": "access_denied",
|
|
"error_description": "local mock denial",
|
|
"state": state,
|
|
}
|
|
else:
|
|
callback_params = {"code": code, "state": state}
|
|
|
|
separator = "&" if urllib.parse.urlsplit(redirect_uri).query else "?"
|
|
location = f"{redirect_uri}{separator}{urllib.parse.urlencode(callback_params)}"
|
|
self.send_response(302)
|
|
self.send_header("location", location)
|
|
self.send_header("content-length", "0")
|
|
self.end_headers()
|
|
|
|
def oauth_token(self, body: bytes) -> None:
|
|
params = {
|
|
key: values[0]
|
|
for key, values in urllib.parse.parse_qs(body.decode()).items()
|
|
}
|
|
|
|
with STATE.lock:
|
|
STATE.oauth_token_requests += 1
|
|
attempt = STATE.oauth_token_requests
|
|
mode = STATE.oauth_mode
|
|
|
|
if mode == "token_temporary_once" and attempt == 1:
|
|
self.send_json(503, {"error": "temporarily_unavailable"})
|
|
return
|
|
|
|
if mode == "token_timeout_once" and attempt == 1:
|
|
time.sleep(DELAY_SECONDS)
|
|
|
|
if (
|
|
params.get("client_id") != OAUTH_CLIENT_ID
|
|
or params.get("client_secret") != OAUTH_CLIENT_SECRET
|
|
or params.get("grant_type") != "authorization_code"
|
|
):
|
|
self.send_json(401, {"error": "invalid_client"})
|
|
return
|
|
|
|
code_value = params.get("code")
|
|
verifier = params.get("code_verifier", "")
|
|
challenge = base64.urlsafe_b64encode(
|
|
hashlib.sha256(verifier.encode()).digest()
|
|
).rstrip(b"=").decode()
|
|
|
|
with STATE.lock:
|
|
code = STATE.oauth_codes.get(code_value or "")
|
|
|
|
if (
|
|
code is None
|
|
or code["used"]
|
|
or code["redirect_uri"] != params.get("redirect_uri")
|
|
or code["challenge"] != challenge
|
|
):
|
|
valid = False
|
|
else:
|
|
valid = True
|
|
code["used"] = True
|
|
token = f"boundary-access-{STATE.oauth_code_counter}-{attempt}"
|
|
STATE.oauth_tokens.add(token)
|
|
|
|
if not valid:
|
|
self.send_json(400, {"error": "invalid_grant"})
|
|
return
|
|
|
|
self.send_json(
|
|
200,
|
|
{"access_token": token, "scope": "", "token_type": "bearer"},
|
|
)
|
|
|
|
def oauth_user(self) -> None:
|
|
authorization = self.headers.get("authorization", "")
|
|
token = authorization.removeprefix("Bearer ")
|
|
|
|
with STATE.lock:
|
|
STATE.oauth_user_requests += 1
|
|
valid = token in STATE.oauth_tokens
|
|
|
|
if not valid:
|
|
self.send_json(401, {"error": "unauthorized"})
|
|
return
|
|
|
|
self.send_json(
|
|
200,
|
|
{
|
|
"id": 4242,
|
|
"login": "local-neighbor",
|
|
"name": "Local Neighbor",
|
|
"html_url": "https://github.com/local-neighbor",
|
|
"avatar_url": "https://avatars.example.invalid/4242",
|
|
},
|
|
)
|
|
|
|
def push(self, body: bytes) -> None:
|
|
try:
|
|
notification = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
self.send_json(400, {"error": "invalid_json"})
|
|
return
|
|
|
|
idempotency_key = self.headers.get("idempotency-key", "")
|
|
authorization_ok = (
|
|
self.headers.get("authorization") == f"Bearer {PUSH_BEARER_TOKEN}"
|
|
)
|
|
|
|
if not authorization_ok or not idempotency_key:
|
|
self.send_json(401, {"error": "unauthorized"})
|
|
return
|
|
|
|
if not all(
|
|
isinstance(notification.get(key), str) and notification[key]
|
|
for key in ("recipient", "title", "body", "idempotency_key")
|
|
) or notification["idempotency_key"] != idempotency_key:
|
|
self.send_json(400, {"error": "invalid_notification"})
|
|
return
|
|
|
|
with STATE.lock:
|
|
STATE.push_attempts += 1
|
|
attempt = STATE.push_attempts
|
|
mode = STATE.push_mode
|
|
|
|
if mode == "reject":
|
|
self.send_json(400, {"error": "invalid_recipient"})
|
|
return
|
|
|
|
if mode == "temporary_once" and attempt == 1:
|
|
self.send_json(503, {"error": "temporarily_unavailable"})
|
|
return
|
|
|
|
with STATE.lock:
|
|
receipt = STATE.push_receipts.get(idempotency_key)
|
|
duplicate = receipt is not None
|
|
|
|
if receipt is None:
|
|
receipt = f"push-receipt-{len(STATE.push_receipts) + 1}"
|
|
STATE.push_receipts[idempotency_key] = receipt
|
|
STATE.push_deliveries += 1
|
|
STATE.push_notifications.append(
|
|
{
|
|
"idempotency_key": idempotency_key,
|
|
"recipient": notification["recipient"],
|
|
"title": notification["title"],
|
|
"body": notification["body"],
|
|
"data": notification.get("data", {}),
|
|
}
|
|
)
|
|
|
|
if mode == "timeout_after_accept" and attempt == 1:
|
|
time.sleep(DELAY_SECONDS)
|
|
|
|
self.send_json(202, {"id": receipt, "duplicate": duplicate})
|
|
|
|
|
|
class BoundarySMTPServer(socketserver.ThreadingTCPServer):
|
|
allow_reuse_address = True
|
|
daemon_threads = True
|
|
|
|
def __init__(self, address: tuple[str, int], mode: str):
|
|
self.mode = mode
|
|
super().__init__(address, BoundarySMTPHandler)
|
|
|
|
|
|
class BoundarySMTPHandler(socketserver.StreamRequestHandler):
|
|
def send_line(self, line: bytes) -> None:
|
|
self.wfile.write(line + b"\r\n")
|
|
self.wfile.flush()
|
|
|
|
def handle(self) -> None:
|
|
port = self.server.server_address[1]
|
|
|
|
with STATE.lock:
|
|
key = str(port)
|
|
STATE.smtp_connections[key] += 1
|
|
connection_number = STATE.smtp_connections[key]
|
|
|
|
if self.server.mode == "timeout":
|
|
time.sleep(DELAY_SECONDS)
|
|
return
|
|
|
|
if self.server.mode == "retry" and connection_number == 1:
|
|
self.send_line(b"421 4.3.0 temporary local mock failure")
|
|
return
|
|
|
|
self.send_line(b"220 boundary.local ESMTP")
|
|
recipients: list[str] = []
|
|
|
|
while True:
|
|
raw = self.rfile.readline(8192)
|
|
if not raw:
|
|
return
|
|
|
|
command = raw.decode(errors="replace").strip()
|
|
upper = command.upper()
|
|
|
|
if upper.startswith(("EHLO ", "HELO ")):
|
|
self.send_line(b"250-boundary.local")
|
|
self.send_line(b"250 SIZE 1048576")
|
|
elif upper.startswith("MAIL FROM:"):
|
|
self.send_line(b"250 2.1.0 sender ok")
|
|
elif upper.startswith("RCPT TO:"):
|
|
recipient = command.split(":", 1)[1].strip("<>")
|
|
|
|
if "reject@" in recipient:
|
|
with STATE.lock:
|
|
STATE.smtp_rejections += 1
|
|
self.send_line(b"550 5.1.1 recipient rejected")
|
|
else:
|
|
recipients.append(recipient)
|
|
self.send_line(b"250 2.1.5 recipient ok")
|
|
elif upper == "DATA":
|
|
self.send_line(b"354 end with <CRLF>.<CRLF>")
|
|
chunks: list[bytes] = []
|
|
|
|
while True:
|
|
line = self.rfile.readline(65_536)
|
|
if not line or line == b".\r\n":
|
|
break
|
|
chunks.append(line)
|
|
|
|
digest = hashlib.sha256(b"".join(chunks)).hexdigest()
|
|
|
|
with STATE.lock:
|
|
STATE.smtp_messages += 1
|
|
STATE.smtp_digests.append(digest)
|
|
|
|
self.send_line(b"250 2.0.0 queued-boundary")
|
|
elif upper == "RSET":
|
|
recipients = []
|
|
self.send_line(b"250 2.0.0 reset")
|
|
elif upper == "QUIT":
|
|
self.send_line(b"221 2.0.0 bye")
|
|
return
|
|
else:
|
|
self.send_line(b"500 5.5.2 unsupported command")
|
|
|
|
|
|
def serve() -> None:
|
|
http_server = BoundaryHTTPServer(("0.0.0.0", HTTP_PORT), BoundaryHTTPHandler)
|
|
smtp_servers = [
|
|
BoundarySMTPServer(("0.0.0.0", SMTP_PORT), "normal"),
|
|
BoundarySMTPServer(("0.0.0.0", SMTP_RETRY_PORT), "retry"),
|
|
BoundarySMTPServer(("0.0.0.0", SMTP_TIMEOUT_PORT), "timeout"),
|
|
]
|
|
servers = [http_server, *smtp_servers]
|
|
threads = [
|
|
threading.Thread(target=server.serve_forever, daemon=True) for server in servers
|
|
]
|
|
|
|
for thread in threads:
|
|
thread.start()
|
|
|
|
stopped = threading.Event()
|
|
|
|
def stop(_signum: int, _frame: Any) -> None:
|
|
stopped.set()
|
|
|
|
signal.signal(signal.SIGTERM, stop)
|
|
signal.signal(signal.SIGINT, stop)
|
|
stopped.wait()
|
|
|
|
for server in servers:
|
|
server.shutdown()
|
|
server.server_close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
serve()
|