diff --git a/app/main.py b/app/main.py index e1bb96f..c4a4c5f 100644 --- a/app/main.py +++ b/app/main.py @@ -56,9 +56,11 @@ config = SecurityConfig( ) guard_deco = SecurityDecorator(config) +_security_middleware = SecurityMiddleware(app.router, config=config) app.add_middleware(SecurityMiddleware, config=config) app.add_middleware(RequestLoggingMiddleware) app.state.guard_decorator = guard_deco +app.state._security_middleware = _security_middleware mongo_worker = MongoWorker() _rabbit_worker: Optional[RabbitWorker] = None diff --git a/app/tests/conftest.py b/app/tests/conftest.py index 3c22d0b..af3770a 100644 --- a/app/tests/conftest.py +++ b/app/tests/conftest.py @@ -8,20 +8,6 @@ from guard import ip_ban_manager from guard_core.handlers.ratelimit_handler import RateLimitManager -def _clear_middleware_suspicious_counts(): - """Search for SecurityMiddleware in middleware stack and clear suspicious_request_counts.""" - from guard.middleware import SecurityMiddleware - from main import app - current = getattr(app, 'middleware_stack', None) or app - visited = set() - while current is not None and id(current) not in visited: - visited.add(id(current)) - if isinstance(current, SecurityMiddleware): - current.suspicious_request_counts.clear() - break - current = getattr(current, 'app', None) - - def _reset_all(): """Full reset of rate-limiter, IP-ban, and suspicious counts.""" # Rate limit timestamps @@ -33,8 +19,14 @@ def _reset_all(): ip_ban_manager.banned_ips.clear() ip_ban_manager.banned_networks.clear() - # Suspicious request counts - _clear_middleware_suspicious_counts() + # Suspicious request counts via direct reference stored in app.state + try: + from main import app + sm = getattr(app.state, '_security_middleware', None) + if sm is not None: + sm.suspicious_request_counts.clear() + except Exception: + pass @pytest.fixture(autouse=True)