From 5871defc4ed85376c72f7fc83eac9a9161f76aed Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Mon, 13 Oct 2025 15:26:36 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E7=99=BB=E5=BD=95?= =?UTF-8?q?=E8=AF=B7=E6=B1=82=E9=80=9F=E7=8E=87=E9=99=90=E5=88=B6=E4=B8=AD?= =?UTF-8?q?=E9=97=B4=E4=BB=B6=EF=BC=8C=E9=98=B2=E6=AD=A2=E6=9A=B4=E5=8A=9B?= =?UTF-8?q?=E7=A0=B4=E8=A7=A3=E5=AF=86=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- server/main.py | 61 ++++++++++++++++- test/bruteforce_simulation.py | 123 ++++++++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+), 1 deletion(-) create mode 100644 test/bruteforce_simulation.py diff --git a/server/main.py b/server/main.py index 701fd773..18d41b99 100644 --- a/server/main.py +++ b/server/main.py @@ -1,6 +1,11 @@ +import asyncio +import time +from collections import defaultdict, deque + import uvicorn -from fastapi import FastAPI, Request +from fastapi import FastAPI, Request, status from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware from server.routers import router @@ -11,6 +16,14 @@ from server.utils.common_utils import setup_logging # 设置日志配置 setup_logging() +RATE_LIMIT_MAX_ATTEMPTS = 10 +RATE_LIMIT_WINDOW_SECONDS = 60 +RATE_LIMIT_ENDPOINTS = {("/api/auth/token", "POST")} + +# In-memory login attempt tracker to reduce brute-force exposure per worker +_login_attempts: defaultdict[str, deque[float]] = defaultdict(deque) +_attempt_lock = asyncio.Lock() + app = FastAPI() app.include_router(router, prefix="/api") @@ -24,6 +37,51 @@ app.add_middleware( ) +def _extract_client_ip(request: Request) -> str: + forwarded_for = request.headers.get("x-forwarded-for") + if forwarded_for: + return forwarded_for.split(",")[0].strip() + if request.client: + return request.client.host + return "unknown" + + +class LoginRateLimitMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next): + normalized_path = request.url.path.rstrip("/") or "/" + request_signature = (normalized_path, request.method.upper()) + + if request_signature in RATE_LIMIT_ENDPOINTS: + client_ip = _extract_client_ip(request) + now = time.monotonic() + + async with _attempt_lock: + attempt_history = _login_attempts[client_ip] + + while attempt_history and now - attempt_history[0] > RATE_LIMIT_WINDOW_SECONDS: + attempt_history.popleft() + + if len(attempt_history) >= RATE_LIMIT_MAX_ATTEMPTS: + retry_after = int(max(1, RATE_LIMIT_WINDOW_SECONDS - (now - attempt_history[0]))) + return JSONResponse( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + content={"detail": "登录尝试过于频繁，请稍后再试"}, + headers={"Retry-After": str(retry_after)}, + ) + + attempt_history.append(now) + + response = await call_next(request) + + if response.status_code < 400: + async with _attempt_lock: + _login_attempts.pop(client_ip, None) + + return response + + return await call_next(request) + + # 鉴权中间件 class AuthMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): @@ -58,6 +116,7 @@ class AuthMiddleware(BaseHTTPMiddleware): # 添加鉴权中间件 +app.add_middleware(LoginRateLimitMiddleware) app.add_middleware(AuthMiddleware) diff --git a/test/bruteforce_simulation.py b/test/bruteforce_simulation.py new file mode 100644 index 00000000..74ea74de --- /dev/null +++ b/test/bruteforce_simulation.py @@ -0,0 +1,123 @@ +""" +Quick script to hammer the login endpoint with invalid credentials and observe rate limiting. + +Usage: + uv run python test/bruteforce_simulation.py --username demo --attempts 20 +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import random +import string +import sys +import time +from collections import Counter + +import httpx + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Simulate brute-force login attempts.") + parser.add_argument("--base-url", default=os.getenv("TEST_BASE_URL", "http://localhost:5050"), help="API base URL") + parser.add_argument("--username", default=os.getenv("TEST_USERNAME", "admin"), help="Login identifier to attack") + parser.add_argument( + "--attempts", type=int, default=20, help="Total number of attempts to issue (default: 20)" + ) + parser.add_argument( + "--concurrency", + type=int, + default=4, + help="Concurrent request limit (default: 4)", + ) + parser.add_argument( + "--delay", + type=float, + default=0.05, + help="Delay in seconds between scheduling attempts (default: 0.05)", + ) + parser.add_argument( + "--password", + default=None, + help="Explicit password to reuse for each request; random values are used when omitted", + ) + return parser.parse_args() + + +def build_payload(username: str, password: str) -> dict[str, str]: + return {"username": username, "password": password} + + +def random_password() -> str: + base = string.ascii_letters + string.digits + "!@#$%^&*" + return "".join(random.choices(base, k=12)) + + +async def attempt_login( + client: httpx.AsyncClient, + semaphore: asyncio.Semaphore, + attempt_no: int, + username: str, + static_password: str | None, +) -> tuple[int, float]: + async with semaphore: + password = static_password or random_password() + payload = build_payload(username, password) + started = time.perf_counter() + response = await client.post("/api/auth/token", data=payload) + elapsed = time.perf_counter() - started + detail = response.json().get("detail") if response.headers.get("content-type", "").startswith("application/json") else response.text + print( + f"[{attempt_no:02d}] {response.status_code} in {elapsed*1000:.1f} ms " + f"(pwd={password!r}) detail={detail!r}" + ) + return response.status_code, elapsed + + +async def run_simulation(args: argparse.Namespace) -> int: + timeout = httpx.Timeout(10.0, connect=3.0) + limits = httpx.Limits(max_connections=args.concurrency, max_keepalive_connections=args.concurrency) + semaphore = asyncio.Semaphore(args.concurrency) + status_counts: Counter[int] = Counter() + + async with httpx.AsyncClient(base_url=args.base_url.rstrip("/"), timeout=timeout, limits=limits) as client: + tasks = [] + for attempt_no in range(1, args.attempts + 1): + tasks.append( + asyncio.create_task( + attempt_login(client, semaphore, attempt_no, args.username, args.password) + ) + ) + if args.delay: + await asyncio.sleep(args.delay) + + for task in asyncio.as_completed(tasks): + status_code, _ = await task + status_counts[status_code] += 1 + + print("\nSummary:") + for code, total in sorted(status_counts.items()): + print(f" HTTP {code}: {total} hits") + + if 429 in status_counts: + print("Rate limiting engaged (HTTP 429 encountered).") + return 0 + + print("No rate limiting observed; consider increasing attempt count or reducing delay.") + return 1 + + +def main() -> None: + args = parse_args() + try: + exit_code = asyncio.run(run_simulation(args)) + except KeyboardInterrupt: + print("\nInterrupted.") + exit_code = 130 + sys.exit(exit_code) + + +if __name__ == "__main__": + main()