720 lines
25 KiB
Python
720 lines
25 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.google_mode = "success"
|
|
self.google_discovery_requests = 0
|
|
self.google_authorize_requests = 0
|
|
self.google_token_requests = 0
|
|
self.google_jwks_requests = 0
|
|
self.google_code_counter = 0
|
|
self.google_codes: dict[str, dict[str, Any]] = {}
|
|
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 == "google":
|
|
self.google_mode = mode
|
|
self.google_discovery_requests = 0
|
|
self.google_authorize_requests = 0
|
|
self.google_token_requests = 0
|
|
self.google_jwks_requests = 0
|
|
self.google_codes = {}
|
|
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"]
|
|
),
|
|
},
|
|
"google": {
|
|
"mode": self.google_mode,
|
|
"discovery_requests": self.google_discovery_requests,
|
|
"authorize_requests": self.google_authorize_requests,
|
|
"token_requests": self.google_token_requests,
|
|
"jwks_requests": self.google_jwks_requests,
|
|
"issued_codes": len(self.google_codes),
|
|
"consumed_codes": sum(
|
|
1 for code in self.google_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"]
|
|
GOOGLE_ISSUER = "http://external-mock:8080"
|
|
GOOGLE_KEY_ID = "local-google-rs256"
|
|
GOOGLE_RSA_MODULUS = int(
|
|
"a42d39e6e3244bdb67afce39612a6c54c27c88d3b730eaabdb70615c59d005e7"
|
|
"cdb22585d196bd337b4c9d80feb5bdc04046ccdba2c523bfbb567cb8165bfa33"
|
|
"f9df6f88828a7acbe6586a003d105b709bdfdd6bbfee36167582f9daa410d792"
|
|
"bcc1c7d3f9bfc8964d8c58250a35540dc8424cf661ec76f76326162d518be5fa"
|
|
"f345922951f0f1f805c34a1a83c2c4f3805677e0faf37f3850fcbbaf4f6e90db"
|
|
"b0f90f5b55548f56a43e69a2f805b257dec650d8ad7417188e379df5050214"
|
|
"fc52f05886b8b5407775793afd6145fae83b1e9728c272e71b80d87ed5d0bec0"
|
|
"63c3703996c8dae1e672803855d59f561fdbb5ee64464d7d09146f530fc03259c5",
|
|
16,
|
|
)
|
|
GOOGLE_RSA_PRIVATE_EXPONENT = int(
|
|
"14b1d2bba4e41d5fc1b92a70972be6cde45a18513fa53ddf7de0b39515892045"
|
|
"70eb44c9927ac2ccab7d23d96fc1eef23de7eec8bcc2c6d7d3407aa625c3604d"
|
|
"8ef0b83967e316c97ef6a41df5948b422d93d17054982d5f355ed629d6467d35"
|
|
"f4ef2446371412afc784aa53b8eeb1f2aecc94b0f5f4fda5ff6c7c9d27cb4fa8"
|
|
"d78f0902df05a4c068e180b1a345ee46607dc4cdf920ee31405a7931dcb61b40"
|
|
"65e1e03963df352553a485a1e31cc072adc46eda1a4dcce71bbe3947c1bc5d7f"
|
|
"a8134d2e58d93792b0776958abf3bc7376846c1cac4c0b634371c80476d847b4"
|
|
"11cd607add231e66341d3def88005d84707ef7715fe288bd881ea5d8430eb001",
|
|
16,
|
|
)
|
|
GOOGLE_RSA_EXPONENT = 65537
|
|
|
|
|
|
def json_bytes(value: Any) -> bytes:
|
|
return json.dumps(value, separators=(",", ":"), sort_keys=True).encode()
|
|
|
|
def base64url(value: bytes) -> str:
|
|
return base64.urlsafe_b64encode(value).rstrip(b"=").decode()
|
|
|
|
|
|
def unsigned_bytes(value: int) -> bytes:
|
|
return value.to_bytes((value.bit_length() + 7) // 8, "big")
|
|
|
|
|
|
def google_id_token(code: dict[str, Any], mode: str) -> str:
|
|
now = int(time.time())
|
|
header = {"alg": "RS256", "kid": GOOGLE_KEY_ID, "typ": "JWT"}
|
|
claims = {
|
|
"iss": GOOGLE_ISSUER,
|
|
"sub": "google-local-subject",
|
|
"aud": OAUTH_CLIENT_ID,
|
|
"exp": now + 300,
|
|
"iat": now,
|
|
"nonce": (
|
|
"invalid-nonce" if mode == "nonce_mismatch" else code["nonce"]
|
|
),
|
|
"email": "google-helper@example.invalid",
|
|
"email_verified": mode != "unverified_email",
|
|
"name": "Google Local Helper",
|
|
}
|
|
signing_input = (
|
|
f"{base64url(json_bytes(header))}.{base64url(json_bytes(claims))}"
|
|
).encode()
|
|
digest_info = bytes.fromhex("3031300d060960864801650304020105000420")
|
|
digest_info += hashlib.sha256(signing_input).digest()
|
|
padding = b"\xff" * (256 - len(digest_info) - 3)
|
|
encoded_message = b"\x00\x01" + padding + b"\x00" + digest_info
|
|
signature = pow(
|
|
int.from_bytes(encoded_message, "big"),
|
|
GOOGLE_RSA_PRIVATE_EXPONENT,
|
|
GOOGLE_RSA_MODULUS,
|
|
).to_bytes(256, "big")
|
|
return f"{signing_input.decode()}.{base64url(signature)}"
|
|
|
|
|
|
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()
|
|
elif parsed.path == "/.well-known/openid-configuration":
|
|
self.google_discovery()
|
|
elif parsed.path == "/google/authorize":
|
|
self.google_authorize(parsed)
|
|
elif parsed.path == "/google/jwks":
|
|
self.google_jwks()
|
|
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 == "/google/token":
|
|
self.google_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 google_discovery(self) -> None:
|
|
with STATE.lock:
|
|
STATE.google_discovery_requests += 1
|
|
|
|
self.send_json(
|
|
200,
|
|
{
|
|
"issuer": GOOGLE_ISSUER,
|
|
"authorization_endpoint": f"{GOOGLE_ISSUER}/google/authorize",
|
|
"token_endpoint": f"{GOOGLE_ISSUER}/google/token",
|
|
"jwks_uri": f"{GOOGLE_ISSUER}/google/jwks",
|
|
"token_endpoint_auth_methods_supported": ["client_secret_post"],
|
|
},
|
|
)
|
|
|
|
def google_authorize(self, parsed: urllib.parse.SplitResult) -> None:
|
|
params = urllib.parse.parse_qs(parsed.query)
|
|
required = {
|
|
"client_id",
|
|
"redirect_uri",
|
|
"state",
|
|
"nonce",
|
|
"scope",
|
|
"code_challenge",
|
|
"code_challenge_method",
|
|
}
|
|
|
|
if (
|
|
not required.issubset(params)
|
|
or params["client_id"][0] != OAUTH_CLIENT_ID
|
|
or params["code_challenge_method"][0] != "S256"
|
|
or set(params["scope"][0].split()) != {"openid", "email", "profile"}
|
|
):
|
|
self.send_json(400, {"error": "invalid_google_authorization_request"})
|
|
return
|
|
|
|
redirect_uri = params["redirect_uri"][0]
|
|
state = params["state"][0]
|
|
|
|
with STATE.lock:
|
|
STATE.google_authorize_requests += 1
|
|
STATE.google_code_counter += 1
|
|
code = f"google-boundary-code-{STATE.google_code_counter}"
|
|
STATE.google_codes[code] = {
|
|
"challenge": params["code_challenge"][0],
|
|
"redirect_uri": redirect_uri,
|
|
"nonce": params["nonce"][0],
|
|
"used": False,
|
|
}
|
|
|
|
location = (
|
|
f"{redirect_uri}?{urllib.parse.urlencode({'code': code, 'state': state})}"
|
|
)
|
|
self.send_response(302)
|
|
self.send_header("location", location)
|
|
self.send_header("content-length", "0")
|
|
self.end_headers()
|
|
|
|
def google_token(self, body: bytes) -> None:
|
|
params = {
|
|
key: values[0]
|
|
for key, values in urllib.parse.parse_qs(body.decode()).items()
|
|
}
|
|
|
|
with STATE.lock:
|
|
STATE.google_token_requests += 1
|
|
mode = STATE.google_mode
|
|
code = STATE.google_codes.get(params.get("code", ""))
|
|
|
|
verifier = params.get("code_verifier", "")
|
|
challenge = base64url(hashlib.sha256(verifier.encode()).digest())
|
|
checks = {
|
|
"client_id": params.get("client_id") == OAUTH_CLIENT_ID,
|
|
"client_secret": params.get("client_secret") == OAUTH_CLIENT_SECRET,
|
|
"grant_type": params.get("grant_type") == "authorization_code",
|
|
"known_code": code is not None,
|
|
"unused_code": code is not None and not code["used"],
|
|
"redirect_uri": (
|
|
code is not None and code["redirect_uri"] == params.get("redirect_uri")
|
|
),
|
|
"pkce": code is not None and code["challenge"] == challenge,
|
|
}
|
|
valid = all(checks.values())
|
|
|
|
if not valid:
|
|
failed_checks = sorted(name for name, passed in checks.items() if not passed)
|
|
if not checks["pkce"] and code is not None:
|
|
failed_checks.append(
|
|
"pkce_lengths_"
|
|
f"{len(verifier)}_{len(code['challenge'])}_{len(challenge)}"
|
|
)
|
|
self.send_json(
|
|
400,
|
|
{
|
|
"error": "invalid_grant",
|
|
"error_description": ",".join(failed_checks),
|
|
},
|
|
)
|
|
return
|
|
|
|
with STATE.lock:
|
|
code["used"] = True
|
|
|
|
self.send_json(
|
|
200,
|
|
{
|
|
"access_token": "google-boundary-access-token",
|
|
"id_token": google_id_token(code, mode),
|
|
"scope": "openid email profile",
|
|
"token_type": "Bearer",
|
|
},
|
|
)
|
|
|
|
def google_jwks(self) -> None:
|
|
with STATE.lock:
|
|
STATE.google_jwks_requests += 1
|
|
|
|
self.send_json(
|
|
200,
|
|
{
|
|
"keys": [
|
|
{
|
|
"alg": "RS256",
|
|
"e": base64url(unsigned_bytes(GOOGLE_RSA_EXPONENT)),
|
|
"kid": GOOGLE_KEY_ID,
|
|
"kty": "RSA",
|
|
"n": base64url(unsigned_bytes(GOOGLE_RSA_MODULUS)),
|
|
"use": "sig",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
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()
|