From a8b47d9ed2f2e66f9300a138e4e333e2a9d06dff Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:57:30 +0000 Subject: [PATCH 1/3] Initial plan From 8c5c284c8c65997a592baaf84139b1bdf9d1330b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:15:06 +0000 Subject: [PATCH 2/3] style: apply Black formatting to src and tests Co-authored-by: VIPHACKER100 <110089810+VIPHACKER100@users.noreply.github.com> --- src/api/app.py | 106 +++++++++++++++++++++------ src/api/auth.py | 66 +++++++++++++---- src/api/dependencies.py | 7 +- src/api/routers/analytics.py | 6 +- src/api/routers/auth.py | 64 ++++++++++++---- src/api/routers/predict.py | 34 +++++++-- src/api/schemas.py | 9 +++ src/core/database.py | 16 +++- src/core/worker.py | 6 +- src/features/store.py | 8 +- src/features/vectorizer.py | 6 +- src/integrations/gmail_client.py | 68 +++++++++++------ src/models/ab_testing.py | 8 +- src/models/adaptive_learning.py | 29 +++++--- src/models/deep_learning.py | 31 +++++--- src/models/drift_monitor.py | 75 ++++++++++++------- src/models/ensemble.py | 56 +++++++++----- src/models/evaluate.py | 31 ++++++-- src/models/explainability.py | 33 +++++---- src/models/feedback.py | 16 +++- src/models/hyperparameter_tuning.py | 5 +- src/models/optimized_pipeline.py | 12 ++- src/models/predict.py | 53 ++++++++++---- src/models/registry.py | 10 ++- src/models/train_models.py | 84 +++++++++++++++------ src/models/vector_search.py | 40 +++++----- src/preprocessing/text_cleaner.py | 12 ++- src/security/alerts.py | 12 +-- src/security/brand_spoof_detector.py | 44 ++++++----- src/security/domain_graph.py | 21 +++--- src/security/domain_intelligence.py | 27 ++++--- src/security/header_analyzer.py | 28 +++---- src/security/homograph_detector.py | 36 ++++----- src/security/obfuscation_detector.py | 20 ++--- src/security/risk_scoring.py | 60 ++++++++------- src/security/rules_engine.py | 13 ++-- src/security/script_detector.py | 15 ++-- src/security/text_behavior.py | 38 +++++++--- src/security/threat_intel.py | 46 +++++++----- src/security/url_extractor.py | 16 ++-- src/security/url_features.py | 55 ++++++++++---- src/security/url_normalizer.py | 40 ++++++---- src/stream/consumer.py | 20 ++++- src/utils/anonymizer.py | 5 +- src/utils/compliance.py | 14 +++- src/utils/config_loader.py | 10 ++- src/utils/data_loader.py | 15 ++-- src/utils/data_versioning.py | 7 +- src/utils/logger.py | 31 +++++--- src/utils/sanitizer.py | 21 +++++- src/utils/secrets.py | 7 +- tests/conftest.py | 3 +- tests/test_abuse_protection.py | 55 +++++++++++--- tests/test_api.py | 83 +++++++++++++++------ tests/test_auth_unit.py | 12 ++- tests/test_idor_security.py | 48 +++++++++--- tests/test_input_validation.py | 21 +++++- tests/test_model.py | 1 + tests/test_new_modules.py | 63 ++++++++++------ tests/test_preprocessing.py | 1 + tests/test_secrets_audit.py | 35 +++++++-- tests/test_security_features.py | 12 ++- tests/test_vectorizer.py | 1 + 63 files changed, 1277 insertions(+), 550 deletions(-) diff --git a/src/api/app.py b/src/api/app.py index 0f97528..918e0ef 100644 --- a/src/api/app.py +++ b/src/api/app.py @@ -18,7 +18,12 @@ from slowapi.middleware import SlowAPIMiddleware from prometheus_fastapi_instrumentator import Instrumentator -from src.utils.logger import logger, security_logger, log_security_event, generate_request_id +from src.utils.logger import ( + logger, + security_logger, + log_security_event, + generate_request_id, +) from src.api.routers import auth, predict, analytics from src.core.database import init_db @@ -27,15 +32,27 @@ # Blacklisted bot / scanner User-Agents _BOT_USER_AGENTS = { - "sqlmap", "nikto", "nmap", "masscan", "gocurl", "bytespider", - "zgrab", "libwww-perl", "python-urllib", "dirbuster", "w3af", - "acunetix", "nessus", "openvas" + "sqlmap", + "nikto", + "nmap", + "masscan", + "gocurl", + "bytespider", + "zgrab", + "libwww-perl", + "python-urllib", + "dirbuster", + "w3af", + "acunetix", + "nessus", + "openvas", } # --------------------------------------------------------------------------- # Lifespan — run setup/teardown around the app lifecycle # --------------------------------------------------------------------------- + @asynccontextmanager async def lifespan(app): """Initialize database tables on startup.""" @@ -44,6 +61,7 @@ async def lifespan(app): log_security_event("SYSTEM_STARTUP", "127.0.0.1", "system", f"Environment: {_ENV}") yield + # --------------------------------------------------------------------------- # App setup # --------------------------------------------------------------------------- @@ -53,8 +71,12 @@ async def lifespan(app): description="Classify emails as Spam or Ham using ML models and multi-layered security heuristics.", version="3.1.0", lifespan=lifespan, - docs_url="/docs" if _ENV != "prod" else None, # Disable Swagger UI in production unless configured - redoc_url="/redoc" if _ENV != "prod" else None, # Disable ReDoc in production unless configured + docs_url=( + "/docs" if _ENV != "prod" else None + ), # Disable Swagger UI in production unless configured + redoc_url=( + "/redoc" if _ENV != "prod" else None + ), # Disable ReDoc in production unless configured ) TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") @@ -68,13 +90,17 @@ async def lifespan(app): Instrumentator().instrument(app).expose(app) # HTTPS Enforcement Middleware in production or when explicitly enabled -if _ENFORCE_HTTPS or (_ENV == "prod" and os.environ.get("ENABLE_HTTPS_REDIRECT", "false").lower() == "true"): +if _ENFORCE_HTTPS or ( + _ENV == "prod" + and os.environ.get("ENABLE_HTTPS_REDIRECT", "false").lower() == "true" +): app.add_middleware(HTTPSRedirectMiddleware) # --------------------------------------------------------------------------- # Anti-Bot & Abuse Protection Middleware # --------------------------------------------------------------------------- + @app.middleware("http") async def anti_bot_middleware(request: Request, call_next): """ @@ -87,24 +113,50 @@ async def anti_bot_middleware(request: Request, call_next): # 1. Honeypot Trap Header Check if request.headers.get("X-Honeypot-Trap") or request.headers.get("X-Bot-Check"): - log_security_event("BOT_HONEYPOT_TRIGGERED", client_ip=client_ip, detail=f"Path={path} UA={user_agent}") - return JSONResponse(status_code=403, content={"detail": "Automated access detected and blocked."}) + log_security_event( + "BOT_HONEYPOT_TRIGGERED", + client_ip=client_ip, + detail=f"Path={path} UA={user_agent}", + ) + return JSONResponse( + status_code=403, + content={"detail": "Automated access detected and blocked."}, + ) # 2. API Routes: Enforce non-empty User-Agent - api_prefixes = ("/predict", "/auth/", "/analyze-security", "/export-report", "/feedback") + api_prefixes = ( + "/predict", + "/auth/", + "/analyze-security", + "/export-report", + "/feedback", + ) if any(path.startswith(prefix) for prefix in api_prefixes): if not user_agent: - log_security_event("BOT_BLOCKED_EMPTY_UA", client_ip=client_ip, detail=f"Path={path}") - return JSONResponse(status_code=403, content={"detail": "User-Agent header is required for API access."}) + log_security_event( + "BOT_BLOCKED_EMPTY_UA", client_ip=client_ip, detail=f"Path={path}" + ) + return JSONResponse( + status_code=403, + content={"detail": "User-Agent header is required for API access."}, + ) # 3. Malicious Scraper / Bot User-Agent Blacklist ua_lower = user_agent.lower() if any(bot in ua_lower for bot in _BOT_USER_AGENTS): - log_security_event("BOT_ATTACK_BLOCKED", client_ip=client_ip, detail=f"Path={path} UA={user_agent}") - return JSONResponse(status_code=403, content={"detail": "Automated request pattern blocked."}) + log_security_event( + "BOT_ATTACK_BLOCKED", + client_ip=client_ip, + detail=f"Path={path} UA={user_agent}", + ) + return JSONResponse( + status_code=403, + content={"detail": "Automated request pattern blocked."}, + ) return await call_next(request) + # --------------------------------------------------------------------------- # Rate limiting & RateLimitExceeded Security Logging # --------------------------------------------------------------------------- @@ -113,18 +165,21 @@ async def anti_bot_middleware(request: Request, call_next): app.state.limiter = limiter -async def custom_rate_limit_handler(request: Request, exc: RateLimitExceeded) -> Response: +async def custom_rate_limit_handler( + request: Request, exc: RateLimitExceeded +) -> Response: client_ip = get_remote_address(request) log_security_event( "RATE_LIMIT_EXCEEDED", client_ip=client_ip, - detail=f"Path={request.url.path} Limit={exc.detail}" + detail=f"Path={request.url.path} Limit={exc.detail}", ) return JSONResponse( status_code=429, - content={"detail": "Rate limit exceeded. Please slow down your requests."} + content={"detail": "Rate limit exceeded. Please slow down your requests."}, ) + app.add_exception_handler(RateLimitExceeded, custom_rate_limit_handler) app.add_middleware(SlowAPIMiddleware) @@ -132,6 +187,7 @@ async def custom_rate_limit_handler(request: Request, exc: RateLimitExceeded) -> # Middleware — Request ID, Security Headers, and Error Audit Logging # --------------------------------------------------------------------------- + @app.middleware("http") async def security_and_request_middleware(request: Request, call_next): req_id = generate_request_id() @@ -144,7 +200,11 @@ async def security_and_request_middleware(request: Request, call_next): logger.info( "[%s] %s %s → %s (%.3fs)", - req_id, request.method, request.url.path, response.status_code, duration, + req_id, + request.method, + request.url.path, + response.status_code, + duration, extra={"request_id": req_id}, ) @@ -154,7 +214,9 @@ async def security_and_request_middleware(request: Request, call_next): response.headers["X-Frame-Options"] = "DENY" response.headers["X-XSS-Protection"] = "1; mode=block" response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" - response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains; preload" + response.headers["Strict-Transport-Security"] = ( + "max-age=31536000; includeSubDomains; preload" + ) response.headers["Content-Security-Policy"] = ( "default-src 'self'; " "script-src 'self' 'unsafe-inline' https://fonts.googleapis.com; " @@ -168,13 +230,13 @@ async def security_and_request_middleware(request: Request, call_next): log_security_event( "SERVER_ERROR", client_ip=client_ip, - detail=f"Method={request.method} Path={request.url.path} Status={response.status_code}" + detail=f"Method={request.method} Path={request.url.path} Status={response.status_code}", ) elif response.status_code in (400, 401, 403, 404, 422): log_security_event( "CLIENT_ERROR", client_ip=client_ip, - detail=f"Method={request.method} Path={request.url.path} Status={response.status_code}" + detail=f"Method={request.method} Path={request.url.path} Status={response.status_code}", ) return response @@ -184,6 +246,7 @@ async def security_and_request_middleware(request: Request, call_next): # Pages # --------------------------------------------------------------------------- + @app.get("/", response_class=HTMLResponse) async def home(request: Request): return templates.TemplateResponse(request, "index.html") @@ -198,6 +261,7 @@ async def dashboard(request: Request): # Health & Readiness # --------------------------------------------------------------------------- + @app.get("/health") async def health(): return {"status": "healthy", "timestamp": datetime.now(timezone.utc).isoformat()} diff --git a/src/api/auth.py b/src/api/auth.py index 6ed6182..8e5fb4b 100644 --- a/src/api/auth.py +++ b/src/api/auth.py @@ -28,9 +28,9 @@ # Configuration — tune via environment variables # --------------------------------------------------------------------------- -JWT_SECRET: str = vault.get("JWT_SECRET") # Must be set; vault raises in prod if absent +JWT_SECRET: str = vault.get("JWT_SECRET") # Must be set; vault raises in prod if absent JWT_ALGORITHM: str = "HS256" -JWT_EXPIRY_HOURS: int = int(os.environ.get("JWT_EXPIRY_HOURS", "1")) # Default: 1 hour +JWT_EXPIRY_HOURS: int = int(os.environ.get("JWT_EXPIRY_HOURS", "1")) # Default: 1 hour LOGIN_MAX_ATTEMPTS: int = int(os.environ.get("LOGIN_MAX_ATTEMPTS", "5")) LOGIN_LOCKOUT_MINUTES: int = int(os.environ.get("LOGIN_LOCKOUT_MINUTES", "15")) RESET_TOKEN_EXPIRY_HOURS: int = int(os.environ.get("RESET_TOKEN_EXPIRY_HOURS", "1")) @@ -39,6 +39,7 @@ # Internal helpers # --------------------------------------------------------------------------- + def _hash_token(raw_token: str) -> str: """Return a SHA-256 hex digest of a raw token for safe DB storage.""" return hashlib.sha256(raw_token.encode()).hexdigest() @@ -53,20 +54,25 @@ def _is_locked(user: User) -> bool: from src.utils.logger import logger, log_security_event + def _record_failed_login(session, user: User, client_ip: str = "N/A") -> None: """Increment failure counter and lock the account when the threshold is exceeded.""" user.failed_login_attempts = (user.failed_login_attempts or 0) + 1 if user.failed_login_attempts >= LOGIN_MAX_ATTEMPTS: - user.locked_until = datetime.now(timezone.utc) + timedelta(minutes=LOGIN_LOCKOUT_MINUTES) + user.locked_until = datetime.now(timezone.utc) + timedelta( + minutes=LOGIN_LOCKOUT_MINUTES + ) logger.warning( "Account '%s' locked for %d minutes after %d failed attempts.", - user.username, LOGIN_LOCKOUT_MINUTES, user.failed_login_attempts, + user.username, + LOGIN_LOCKOUT_MINUTES, + user.failed_login_attempts, ) log_security_event( "ACCOUNT_LOCKED", client_ip=client_ip, username=user.username, - detail=f"Locked for {LOGIN_LOCKOUT_MINUTES} min after {user.failed_login_attempts} failed attempts" + detail=f"Locked for {LOGIN_LOCKOUT_MINUTES} min after {user.failed_login_attempts} failed attempts", ) session.commit() @@ -77,10 +83,12 @@ def _reset_login_counter(session, user: User) -> None: user.locked_until = None session.commit() + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- + def init_db(): """No-op. Database schema is initialized via src.core.database.init_db.""" pass @@ -136,7 +144,9 @@ def register_user(username: str, password: str, email: str | None = None) -> dic session.close() -def authenticate_user(username: str, password: str, client_ip: str = "N/A") -> dict | None: +def authenticate_user( + username: str, password: str, client_ip: str = "N/A" +) -> dict | None: """ Validate credentials and return a short-lived JWT token. @@ -156,7 +166,12 @@ def authenticate_user(username: str, password: str, client_ip: str = "N/A") -> d # Account lockout check if _is_locked(user): logger.warning("Login attempt on locked account: %s", username) - log_security_event("AUTH_FAILURE", client_ip=client_ip, username=username, detail="Locked account login attempt") + log_security_event( + "AUTH_FAILURE", + client_ip=client_ip, + username=username, + detail="Locked account login attempt", + ) return None # Verify password @@ -179,7 +194,11 @@ def authenticate_user(username: str, password: str, client_ip: str = "N/A") -> d JWT_SECRET, algorithm=JWT_ALGORITHM, ) - return {"token": token, "username": username, "expires_in": JWT_EXPIRY_HOURS * 3600} + return { + "token": token, + "username": username, + "expires_in": JWT_EXPIRY_HOURS * 3600, + } finally: session.close() @@ -234,7 +253,9 @@ def request_password_reset(username: str) -> str | None: raw_token = secrets.token_urlsafe(32) user.password_reset_token_hash = _hash_token(raw_token) - user.password_reset_expires = datetime.now(timezone.utc) + timedelta(hours=RESET_TOKEN_EXPIRY_HOURS) + user.password_reset_expires = datetime.now(timezone.utc) + timedelta( + hours=RESET_TOKEN_EXPIRY_HOURS + ) session.commit() logger.info("Password reset token issued for user: %s", username) @@ -281,7 +302,9 @@ def reset_password(username: str, raw_token: str, new_password: str) -> bool: return False # Update password and clear reset fields - user.password_hash = bcrypt.hashpw(new_password.encode(), bcrypt.gensalt()).decode() + user.password_hash = bcrypt.hashpw( + new_password.encode(), bcrypt.gensalt() + ).decode() user.password_reset_token_hash = None user.password_reset_expires = None user.failed_login_attempts = 0 @@ -297,7 +320,9 @@ def reset_password(username: str, raw_token: str, new_password: str) -> bool: session.close() -def log_usage(user_id: int | None, endpoint: str, request_body: str, response_body: str): +def log_usage( + user_id: int | None, endpoint: str, request_body: str, response_body: str +): """Log an API usage record.""" session = SessionLocal() try: @@ -329,11 +354,17 @@ def get_user_by_username(username: str) -> User | None: # Ownership-checked resource access (IDOR Prevention) # --------------------------------------------------------------------------- + def get_user_logs(user_id: int) -> list[dict]: """Retrieve usage logs strictly filtered by user_id.""" session = SessionLocal() try: - logs = session.query(UsageLog).filter_by(user_id=user_id).order_by(UsageLog.id.desc()).all() + logs = ( + session.query(UsageLog) + .filter_by(user_id=user_id) + .order_by(UsageLog.id.desc()) + .all() + ) return [ { "id": l.id, @@ -376,9 +407,12 @@ def get_user_feedback_by_id(feedback_id: int, user_id: int) -> dict | None: Strictly verifies ownership: feedback.user_id == user_id. Returns None if not owned. """ from src.core.database import Feedback as DBFeedback + session = SessionLocal() try: - fb = session.query(DBFeedback).filter_by(id=feedback_id, user_id=user_id).first() + fb = ( + session.query(DBFeedback).filter_by(id=feedback_id, user_id=user_id).first() + ) if fb is None: return None return { @@ -400,9 +434,12 @@ def delete_user_feedback(feedback_id: int, user_id: int) -> bool: Strictly verifies ownership: feedback.user_id == user_id. Returns False if not owned. """ from src.core.database import Feedback as DBFeedback + session = SessionLocal() try: - fb = session.query(DBFeedback).filter_by(id=feedback_id, user_id=user_id).first() + fb = ( + session.query(DBFeedback).filter_by(id=feedback_id, user_id=user_id).first() + ) if fb is None: return False session.delete(fb) @@ -414,4 +451,3 @@ def delete_user_feedback(feedback_id: int, user_id: int) -> bool: raise finally: session.close() - diff --git a/src/api/dependencies.py b/src/api/dependencies.py index c281f6d..645dca7 100644 --- a/src/api/dependencies.py +++ b/src/api/dependencies.py @@ -3,9 +3,11 @@ from src.api.auth import verify_token, verify_api_key, get_user_by_username from src.core.database import User from src.models.ab_testing import ABTest + # Global instances ab_test = ABTest("naive_bayes", "svm", split=0.5) + async def optional_auth(request: Request) -> Optional[User]: """ Attempt to authenticate via Bearer token or X-API-Key header. @@ -27,7 +29,10 @@ async def optional_auth(request: Request) -> Optional[User]: return None -async def get_current_user(request: Request, user: Optional[User] = Depends(optional_auth)) -> User: + +async def get_current_user( + request: Request, user: Optional[User] = Depends(optional_auth) +) -> User: """ Require authentication via Bearer token or X-API-Key header. Raises HTTPException(401) if unauthenticated or user invalid. diff --git a/src/api/routers/analytics.py b/src/api/routers/analytics.py index 4bbce61..e97d433 100644 --- a/src/api/routers/analytics.py +++ b/src/api/routers/analytics.py @@ -5,13 +5,17 @@ router = APIRouter(tags=["analytics"]) + @router.get("/ab/summary") async def get_ab_summary(): return ab_test.summary() + @router.get("/analytics") async def analytics(): - metrics_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "models", "metrics.json") + metrics_path = os.path.join( + os.path.dirname(__file__), "..", "..", "..", "models", "metrics.json" + ) if not os.path.exists(metrics_path): return {"error": "No metrics available. Train models first."} with open(metrics_path) as f: diff --git a/src/api/routers/auth.py b/src/api/routers/auth.py index 7831b93..b99bc21 100644 --- a/src/api/routers/auth.py +++ b/src/api/routers/auth.py @@ -49,10 +49,17 @@ async def register(request: Request, body: RegisterRequest): client_ip = get_remote_address(request) try: result = register_user(body.username, body.password, body.email) - log_security_event("USER_REGISTERED", client_ip=client_ip, username=body.username) + log_security_event( + "USER_REGISTERED", client_ip=client_ip, username=body.username + ) return result except ValueError as e: - log_security_event("REGISTER_FAILED", client_ip=client_ip, username=body.username, detail=str(e)) + log_security_event( + "REGISTER_FAILED", + client_ip=client_ip, + username=body.username, + detail=str(e), + ) raise HTTPException(status_code=409, detail=str(e)) @@ -68,7 +75,12 @@ async def login(request: Request, body: LoginRequest): client_ip = get_remote_address(request) result = authenticate_user(body.username, body.password, client_ip=client_ip) if result is None: - log_security_event("AUTH_FAILURE", client_ip=client_ip, username=body.username, detail="Invalid credentials or locked account") + log_security_event( + "AUTH_FAILURE", + client_ip=client_ip, + username=body.username, + detail="Invalid credentials or locked account", + ) raise HTTPException(status_code=401, detail="Invalid credentials") log_security_event("AUTH_SUCCESS", client_ip=client_ip, username=body.username) return result @@ -88,7 +100,9 @@ async def password_reset_request(request: Request, body: PasswordResetRequestSch """ client_ip = get_remote_address(request) raw_token = request_password_reset(body.username) - log_security_event("PASSWORD_RESET_REQUESTED", client_ip=client_ip, username=body.username) + log_security_event( + "PASSWORD_RESET_REQUESTED", client_ip=client_ip, username=body.username + ) # Return a uniform response body to prevent user enumeration. # In production: send raw_token via email and do NOT include it in the response. return { @@ -111,14 +125,26 @@ async def password_reset(request: Request, body: PasswordResetSchema): try: ok = reset_password(body.username, body.token, body.new_password) except ValueError as e: - log_security_event("PASSWORD_RESET_FAILED", client_ip=client_ip, username=body.username, detail=str(e)) + log_security_event( + "PASSWORD_RESET_FAILED", + client_ip=client_ip, + username=body.username, + detail=str(e), + ) raise HTTPException(status_code=422, detail=str(e)) if not ok: - log_security_event("PASSWORD_RESET_FAILED", client_ip=client_ip, username=body.username, detail="Invalid or expired token") + log_security_event( + "PASSWORD_RESET_FAILED", + client_ip=client_ip, + username=body.username, + detail="Invalid or expired token", + ) raise HTTPException(status_code=400, detail="Invalid or expired reset token.") - log_security_event("PASSWORD_RESET_SUCCESS", client_ip=client_ip, username=body.username) + log_security_event( + "PASSWORD_RESET_SUCCESS", client_ip=client_ip, username=body.username + ) return {"detail": "Password updated successfully."} @@ -135,7 +161,9 @@ async def get_my_profile(current_user: User = Depends(get_current_user)): "username": current_user.username, "email": current_user.email, "is_email_verified": current_user.is_email_verified, - "created_at": current_user.created_at.isoformat() if current_user.created_at else None, + "created_at": ( + current_user.created_at.isoformat() if current_user.created_at else None + ), } @@ -153,29 +181,39 @@ async def get_my_log_by_id(log_id: int, current_user: User = Depends(get_current """ log = get_user_log_by_id(log_id, current_user.id) if not log: - raise HTTPException(status_code=404, detail="Log entry not found or access denied") + raise HTTPException( + status_code=404, detail="Log entry not found or access denied" + ) return log @router.get("/feedback/{feedback_id}") -async def get_my_feedback_by_id(feedback_id: int, current_user: User = Depends(get_current_user)): +async def get_my_feedback_by_id( + feedback_id: int, current_user: User = Depends(get_current_user) +): """ Retrieve a specific feedback entry by feedback_id. Prevents IDOR by verifying ownership: feedback.user_id == current_user.id. """ fb = get_user_feedback_by_id(feedback_id, current_user.id) if not fb: - raise HTTPException(status_code=404, detail="Feedback entry not found or access denied") + raise HTTPException( + status_code=404, detail="Feedback entry not found or access denied" + ) return fb @router.delete("/feedback/{feedback_id}") -async def delete_my_feedback(feedback_id: int, current_user: User = Depends(get_current_user)): +async def delete_my_feedback( + feedback_id: int, current_user: User = Depends(get_current_user) +): """ Delete a specific feedback entry by feedback_id. Prevents IDOR by verifying ownership: feedback.user_id == current_user.id. """ deleted = delete_user_feedback(feedback_id, current_user.id) if not deleted: - raise HTTPException(status_code=404, detail="Feedback entry not found or access denied") + raise HTTPException( + status_code=404, detail="Feedback entry not found or access denied" + ) return {"detail": "Feedback entry deleted successfully."} diff --git a/src/api/routers/predict.py b/src/api/routers/predict.py index aa8e3d8..d9c65f0 100644 --- a/src/api/routers/predict.py +++ b/src/api/routers/predict.py @@ -25,6 +25,7 @@ _arq_pool = None + async def _get_arq_pool(): """Lazily create and cache the ARQ Redis connection pool.""" global _arq_pool @@ -33,7 +34,10 @@ async def _get_arq_pool(): try: _arq_pool = await create_pool(RedisSettings(host=redis_host, port=6379)) except Exception as e: - logger.warning("ARQ pool unavailable (Redis not running?): %s. Falling back to sync.", e) + logger.warning( + "ARQ pool unavailable (Redis not running?): %s. Falling back to sync.", + e, + ) return _arq_pool @@ -41,6 +45,7 @@ async def _get_arq_pool(): # Endpoints (with per-IP Rate Limiting & Anti-Abuse Controls) # --------------------------------------------------------------------------- + @router.post("/predict") @_limiter.limit("15/minute") async def predict(request: Request, body: PredictRequest): @@ -50,18 +55,21 @@ async def predict(request: Request, body: PredictRequest): model_name = ab_test.select() result = predict_email(body.text, model_name, body.headers or "") - logger.info("Email prediction: %s", result['prediction']) + logger.info("Email prediction: %s", result["prediction"]) # Enqueue background tasks via ARQ (falls back to sync if Redis unavailable) pool = await _get_arq_pool() if pool: - await pool.enqueue_job("check_drift", body.text, result['prediction']) + await pool.enqueue_job("check_drift", body.text, result["prediction"]) if result.get("security_risk_score", 0) >= 75: - await pool.enqueue_job("send_security_alert", "Phishing Detection", result) + await pool.enqueue_job( + "send_security_alert", "Phishing Detection", result + ) else: # Sync fallback when Redis is not available (e.g., local dev) from src.models.drift_monitor import DriftMonitor - DriftMonitor().check([body.text], [result['prediction']]) + + DriftMonitor().check([body.text], [result["prediction"]]) return result except Exception as e: @@ -73,6 +81,7 @@ async def predict(request: Request, body: PredictRequest): @_limiter.limit("15/minute") async def analyze_security_endpoint(request: Request, body: PredictRequest): from src.security.risk_scoring import calculate_security_risk + try: analysis = calculate_security_risk(body.text, body.headers or "") return analysis @@ -83,7 +92,11 @@ async def analyze_security_endpoint(request: Request, body: PredictRequest): @router.post("/predict/batch") @_limiter.limit("5/minute") -async def predict_batch_endpoint(request: Request, body: BatchPredictRequest, user: Optional[User] = Depends(optional_auth)): +async def predict_batch_endpoint( + request: Request, + body: BatchPredictRequest, + user: Optional[User] = Depends(optional_auth), +): try: results = predict_batch(body.emails, body.model_name) return {"results": results, "count": len(results)} @@ -99,13 +112,14 @@ async def predict_batch_endpoint(request: Request, body: BatchPredictRequest, us async def export_report_endpoint(request: Request, body: PredictRequest): from src.security.risk_scoring import calculate_security_risk from src.security.url_extractor import extract_urls + try: analysis = calculate_security_risk(body.text, body.headers or "") report = { "timestamp": datetime.now(timezone.utc).isoformat(), "original_email": body.text, "extracted_urls": extract_urls(body.text), - "forensic_analysis": analysis + "forensic_analysis": analysis, } return JSONResponse(content=report) except Exception as e: @@ -115,7 +129,11 @@ async def export_report_endpoint(request: Request, body: PredictRequest): @router.post("/feedback") @_limiter.limit("10/minute") -async def submit_feedback(request: Request, body: FeedbackRequest, user: Optional[User] = Depends(optional_auth)): +async def submit_feedback( + request: Request, + body: FeedbackRequest, + user: Optional[User] = Depends(optional_auth), +): clean_text = anonymize_text(body.email_text) user_id = user.id if user else None entry = save_feedback( diff --git a/src/api/schemas.py b/src/api/schemas.py index 8066d4d..e45765e 100644 --- a/src/api/schemas.py +++ b/src/api/schemas.py @@ -11,6 +11,7 @@ MAX_TEXT_LENGTH = 50_000 MAX_HEADER_LENGTH = 10_000 + class PredictRequest(BaseModel): text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH) model: Optional[str] = None @@ -42,6 +43,7 @@ def check_honeypot(cls, v: Optional[str]) -> Optional[str]: raise ValueError("Bot activity detected via honeypot input.") return v + class BatchPredictRequest(BaseModel): model_config = {"protected_namespaces": ()} @@ -59,6 +61,7 @@ def sanitise_batch_emails(cls, v: list[str]) -> list[str]: def check_batch_model_name(cls, v: Optional[str]) -> Optional[str]: return validate_model_name(v) + class RegisterRequest(BaseModel): username: str = Field(..., min_length=3, max_length=50) # Minimum 8 chars enforced; maximum 128 chars protects bcrypt from CPU DoS @@ -89,6 +92,7 @@ def check_honeypot(cls, v: Optional[str]) -> Optional[str]: raise ValueError("Bot registration attempt detected.") return v + class LoginRequest(BaseModel): username: str = Field(..., min_length=1, max_length=50) password: str = Field(..., min_length=1, max_length=128) @@ -98,8 +102,10 @@ class LoginRequest(BaseModel): def check_login_username(cls, v: str) -> str: return validate_username(v) + class PasswordResetRequestSchema(BaseModel): """Request a password-reset token to be sent out-of-band (e.g. email).""" + username: str = Field(..., min_length=1, max_length=50) @field_validator("username") @@ -107,8 +113,10 @@ class PasswordResetRequestSchema(BaseModel): def check_reset_username(cls, v: str) -> str: return validate_username(v) + class PasswordResetSchema(BaseModel): """Consume a password-reset token and set a new password.""" + username: str = Field(..., min_length=1, max_length=50) token: str = Field(..., min_length=1, max_length=128) new_password: str = Field(..., min_length=8, max_length=128) @@ -123,6 +131,7 @@ def check_reset_user(cls, v: str) -> str: def check_reset_token(cls, v: str) -> str: return validate_token(v) + class FeedbackRequest(BaseModel): model_config = {"protected_namespaces": ()} diff --git a/src/core/database.py b/src/core/database.py index c30426f..214ed0f 100644 --- a/src/core/database.py +++ b/src/core/database.py @@ -1,5 +1,14 @@ import os -from sqlalchemy import create_engine, Column, Integer, String, Text, Boolean, ForeignKey, DateTime +from sqlalchemy import ( + create_engine, + Column, + Integer, + String, + Text, + Boolean, + ForeignKey, + DateTime, +) from sqlalchemy.orm import declarative_base, sessionmaker, relationship from datetime import datetime, timezone @@ -16,6 +25,7 @@ Base = declarative_base() + class User(Base): __tablename__ = "users" @@ -41,6 +51,7 @@ class User(Base): logs = relationship("UsageLog", back_populates="user") + class UsageLog(Base): __tablename__ = "usage_logs" @@ -53,6 +64,7 @@ class UsageLog(Base): user = relationship("User", back_populates="logs") + class Feedback(Base): __tablename__ = "feedback" @@ -66,9 +78,11 @@ class Feedback(Base): user = relationship("User") + def init_db(): Base.metadata.create_all(bind=engine) + def get_db(): db = SessionLocal() try: diff --git a/src/core/worker.py b/src/core/worker.py index 4a10cb5..8f66715 100644 --- a/src/core/worker.py +++ b/src/core/worker.py @@ -5,15 +5,19 @@ drift_monitor = DriftMonitor() + async def check_drift(ctx, text: str, prediction: str): """Background task to check for model drift.""" drift_monitor.check([text], [prediction]) + async def send_security_alert(ctx, alert_type: str, result: dict): """Background task to send a security alert.""" trigger_security_alert(alert_type, result) + class WorkerSettings: """ARQ Worker configuration""" + functions = [check_drift, send_security_alert] - redis_settings = RedisSettings(host=os.getenv('REDIS_HOST', 'localhost'), port=6379) + redis_settings = RedisSettings(host=os.getenv("REDIS_HOST", "localhost"), port=6379) diff --git a/src/features/store.py b/src/features/store.py index abc0f2f..8275876 100644 --- a/src/features/store.py +++ b/src/features/store.py @@ -26,7 +26,9 @@ def _save_meta(meta: dict): json.dump(meta, f, indent=2) -def save_features(X, name: str, dataset_version: str = "unknown", method: str = "tfidf") -> str: +def save_features( + X, name: str, dataset_version: str = "unknown", method: str = "tfidf" +) -> str: """ Persist a feature matrix and record metadata. @@ -66,7 +68,9 @@ def load_features(name: str): """Load a previously saved feature matrix.""" meta = _load_meta() if name not in meta: - raise KeyError(f"Feature set '{name}' not found in store. Available: {list(meta.keys())}") + raise KeyError( + f"Feature set '{name}' not found in store. Available: {list(meta.keys())}" + ) path = meta[name]["path"] X = joblib.load(path) logger.info("Loaded feature set '%s' from %s", name, path) diff --git a/src/features/vectorizer.py b/src/features/vectorizer.py index 6dbfe02..335daa2 100644 --- a/src/features/vectorizer.py +++ b/src/features/vectorizer.py @@ -26,7 +26,7 @@ VECTORIZER_TYPES: dict[str, type] = { "bow": CountVectorizer, "tfidf": TfidfVectorizer, - "tfidf_char": TfidfVectorizer, # character n-gram variant + "tfidf_char": TfidfVectorizer, # character n-gram variant } # Sensible defaults that work well for email phishing / spam detection @@ -100,7 +100,8 @@ def fit_and_save( joblib.dump(vec, path) logger.info( "Fitted %s vectorizer — vocab=%d features, matrix=%s → saved to %s", - method.upper(), X_transformed.shape[1], + method.upper(), + X_transformed.shape[1], "×".join(str(d) for d in X_transformed.shape), path, ) @@ -118,6 +119,7 @@ def load_vectorizer(path: str = "models/vectorizer.pkl"): # Security feature matrix (unchanged from v1) # --------------------------------------------------------------------------- + def extract_security_matrix(texts: list): """ Extract hand-crafted security features from a list of texts. diff --git a/src/integrations/gmail_client.py b/src/integrations/gmail_client.py index d3af597..33f2703 100644 --- a/src/integrations/gmail_client.py +++ b/src/integrations/gmail_client.py @@ -13,7 +13,8 @@ from src.utils.logger import logger # Required scopes for reading and modifying labels -SCOPES = ['https://www.googleapis.com/auth/gmail.modify'] +SCOPES = ["https://www.googleapis.com/auth/gmail.modify"] + class GmailScanner: def __init__(self, credentials_path: str = "config/credentials.json"): @@ -24,23 +25,27 @@ def __init__(self, credentials_path: str = "config/credentials.json"): def _authenticate(self): """Standard Google OAuth2 flow.""" - if os.path.exists('token.json'): - self.creds = Credentials.from_authorized_user_file('token.json', SCOPES) - + if os.path.exists("token.json"): + self.creds = Credentials.from_authorized_user_file("token.json", SCOPES) + if not self.creds or not self.creds.valid: if self.creds and self.creds.expired and self.creds.refresh_token: self.creds.refresh(Request()) else: if not os.path.exists(self.creds_path): - logger.warning("Gmail credentials.json missing. Running in MOCK MODE.") + logger.warning( + "Gmail credentials.json missing. Running in MOCK MODE." + ) return - flow = InstalledAppFlow.from_client_secrets_file(self.creds_path, SCOPES) + flow = InstalledAppFlow.from_client_secrets_file( + self.creds_path, SCOPES + ) self.creds = flow.run_local_server(port=0) - - with open('token.json', 'w') as token: + + with open("token.json", "w") as token: token.write(self.creds.to_json()) - self.service = build('gmail', 'v1', credentials=self.creds) + self.service = build("gmail", "v1", credentials=self.creds) def scan_inbox(self, max_results: int = 10): """Fetch latest emails, classify them, and label threats.""" @@ -49,28 +54,44 @@ def scan_inbox(self, max_results: int = 10): return try: - results = self.service.users().messages().list(userId='me', labelIds=['INBOX'], maxResults=max_results).execute() - messages = results.get('messages', []) + results = ( + self.service.users() + .messages() + .list(userId="me", labelIds=["INBOX"], maxResults=max_results) + .execute() + ) + messages = results.get("messages", []) for msg_meta in messages: - msg = self.service.users().messages().get(userId='me', id=msg_meta['id'], format='full').execute() - + msg = ( + self.service.users() + .messages() + .get(userId="me", id=msg_meta["id"], format="full") + .execute() + ) + # Extract text parts - payload = msg.get('payload', {}) + payload = msg.get("payload", {}) body = "" - if 'parts' in payload: - for part in payload['parts']: - if part['mimeType'] == 'text/plain': - body = base64.urlsafe_b64decode(part['body']['data']).decode() + if "parts" in payload: + for part in payload["parts"]: + if part["mimeType"] == "text/plain": + body = base64.urlsafe_b64decode( + part["body"]["data"] + ).decode() else: - body = base64.urlsafe_b64decode(payload['body']['data']).decode() + body = base64.urlsafe_b64decode(payload["body"]["data"]).decode() # Predict res = predict_email(body) - - if res['prediction'] == 'spam' or res['security_risk_score'] > 70: - logger.warning("Threat Found in Gmail: %s (Risk: %d)", msg_meta['id'], res['security_risk_score']) - self._mark_as_threat(msg_meta['id']) + + if res["prediction"] == "spam" or res["security_risk_score"] > 70: + logger.warning( + "Threat Found in Gmail: %s (Risk: %d)", + msg_meta["id"], + res["security_risk_score"], + ) + self._mark_as_threat(msg_meta["id"]) except Exception as e: logger.error("Gmail Scan Error: %s", e) @@ -84,6 +105,7 @@ def _mark_as_threat(self, msg_id: str): # body={'ids': [msg_id], 'addLabelIds': ['SPAM'], 'removeLabelIds': ['INBOX']} # ).execute() + if __name__ == "__main__": scanner = GmailScanner() scanner.scan_inbox() diff --git a/src/models/ab_testing.py b/src/models/ab_testing.py index fa6448f..f66771d 100644 --- a/src/models/ab_testing.py +++ b/src/models/ab_testing.py @@ -8,7 +8,9 @@ from datetime import datetime, timezone from src.utils.logger import logger -AB_LOG_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "logs", "ab_test_results.json") +AB_LOG_PATH = os.path.join( + os.path.dirname(__file__), "..", "..", "logs", "ab_test_results.json" +) class ABTest: @@ -32,7 +34,9 @@ def select(self) -> str: """Randomly select a model based on the traffic split.""" return self.model_a if random.random() < self.split else self.model_b - def record(self, model_used: str, prediction: str, correct_label: str | None = None): + def record( + self, model_used: str, prediction: str, correct_label: str | None = None + ): """Record an A/B test outcome.""" entry = { "model": model_used, diff --git a/src/models/adaptive_learning.py b/src/models/adaptive_learning.py index f6cc349..7ec2ef0 100644 --- a/src/models/adaptive_learning.py +++ b/src/models/adaptive_learning.py @@ -6,14 +6,15 @@ from src.utils.logger import logger from src.core.database import SessionLocal, Feedback + class AdaptiveLearner: """ Monitors feedback and triggers incremental retraining. """ - + def __init__(self, threshold: int = 50): - self.threshold = threshold # retrain every 50 new feedback entries - + self.threshold = threshold # retrain every 50 new feedback entries + def check_and_retrain(self): """Check if enough new data exists to justify retraining.""" session = SessionLocal() @@ -28,15 +29,22 @@ def check_and_retrain(self): "email_text": f.email_text, "predicted_label": f.predicted_label, "correct_label": f.correct_label, - "model_used": f.model_used + "model_used": f.model_used, } for f in feedbacks ] df = pd.DataFrame(data) - logger.info("Adaptive Learning: Threshold reached (%d entries). Starting incremental retrain...", len(df)) + logger.info( + "Adaptive Learning: Threshold reached (%d entries). Starting incremental retrain...", + len(df), + ) self._run_retraining_job(df) else: - logger.info("Adaptive Learning: Not enough new data yet (%d/%d)", len(feedbacks), self.threshold) + logger.info( + "Adaptive Learning: Not enough new data yet (%d/%d)", + len(feedbacks), + self.threshold, + ) except Exception as e: logger.error("Adaptive Learning Error: %s", e) finally: @@ -44,20 +52,21 @@ def check_and_retrain(self): def _run_retraining_job(self, new_data: pd.DataFrame): """ - Simulate retraining logic. + Simulate retraining logic. In production, this would call scripts/train_pipeline.py with the augmented dataset. """ logger.info("Loading baseline dataset...") # ... load v2_emails.csv ... - + logger.info("Merging with %d live feedback samples...", len(new_data)) # ... logic to merge and create updated features ... - + logger.info("Retraining Naive Bayes and SVM models...") # joblib.dump(new_model, "models/adaptive_model.pkl") - + logger.info("Retraining COMPLETE. New model registered in registry.") + if __name__ == "__main__": learner = AdaptiveLearner(threshold=5) learner.check_and_retrain() diff --git a/src/models/deep_learning.py b/src/models/deep_learning.py index 77e145c..f5c5650 100644 --- a/src/models/deep_learning.py +++ b/src/models/deep_learning.py @@ -2,20 +2,27 @@ Advanced Deep Learning Module Uses HuggingFace Transformers for semantic phishing and spam detection. """ + import logging from transformers import pipeline + class DeepLearningModel: def __init__(self, model_name: str = "ealvaradob/bert-finetuned-phishing"): self.model_name = model_name self.classifier = None self.is_loaded = False - + def load(self): """Load the transformer model into memory.""" logging.info("Loading Deep Learning transformer: %s", self.model_name) try: - self.classifier = pipeline("text-classification", model=self.model_name, truncation=True, max_length=512) + self.classifier = pipeline( + "text-classification", + model=self.model_name, + truncation=True, + max_length=512, + ) self.is_loaded = True except Exception as e: logging.error(f"Failed to load transformer model: {e}") @@ -25,15 +32,19 @@ def predict(self, text: str) -> dict: """Process text sequences using transformer.""" if not self.is_loaded: self.load() - + result = self.classifier(text)[0] - label = result['label'].lower() - + label = result["label"].lower() + # Typically LABEL_1 is Phishing/Spam in binary classifications - prediction = "spam" if "phishing" in label or "spam" in label or label == "label_1" else "ham" - + prediction = ( + "spam" + if "phishing" in label or "spam" in label or label == "label_1" + else "ham" + ) + return { - "prediction": prediction, - "confidence": round(result['score'], 4), - "model_used": self.model_name + "prediction": prediction, + "confidence": round(result["score"], 4), + "model_used": self.model_name, } diff --git a/src/models/drift_monitor.py b/src/models/drift_monitor.py index c0b616e..4fe3343 100644 --- a/src/models/drift_monitor.py +++ b/src/models/drift_monitor.py @@ -8,7 +8,9 @@ from collections import Counter from src.utils.logger import logger -DRIFT_LOG = os.path.join(os.path.dirname(__file__), "..", "..", "logs", "drift_log.json") +DRIFT_LOG = os.path.join( + os.path.dirname(__file__), "..", "..", "logs", "drift_log.json" +) class DriftMonitor: @@ -32,9 +34,9 @@ def __init__(self, training_stats: dict | None = None): - vocab: set or list of frequent training words """ self.training_stats = training_stats or {} - self.threshold_spam_ratio = 0.15 # alert if ratio shifts > 15% - self.threshold_length_pct = 0.30 # alert if avg length shifts > 30% - self.threshold_vocab_drift = 0.40 # alert if >40% new words + self.threshold_spam_ratio = 0.15 # alert if ratio shifts > 15% + self.threshold_length_pct = 0.30 # alert if avg length shifts > 30% + self.threshold_vocab_drift = 0.40 # alert if >40% new words def compute_training_stats(self, labels: list[str], texts: list[str]) -> dict: """Compute and store baseline statistics from training data.""" @@ -45,10 +47,12 @@ def compute_training_stats(self, labels: list[str], texts: list[str]) -> dict: "vocab": list(set(" ".join(texts).split()))[:5000], # top 5k unique words "total_samples": len(labels), } - logger.info("Computed training stats: spam_ratio=%.3f, avg_len=%.1f, vocab_size=%d", - self.training_stats["spam_ratio"], - self.training_stats["avg_text_length"], - len(self.training_stats["vocab"])) + logger.info( + "Computed training stats: spam_ratio=%.3f, avg_len=%.1f, vocab_size=%d", + self.training_stats["spam_ratio"], + self.training_stats["avg_text_length"], + len(self.training_stats["vocab"]), + ) return self.training_stats def check(self, predictions: list[str], texts: list[str]) -> dict: @@ -60,31 +64,43 @@ def check(self, predictions: list[str], texts: list[str]) -> dict: dict with alert flags and details. """ if not self.training_stats: - return {"error": "No training stats available. Call compute_training_stats first."} + return { + "error": "No training stats available. Call compute_training_stats first." + } alerts = [] # 1. Prediction distribution drift - live_spam_ratio = sum(1 for p in predictions if p == "spam") / len(predictions) if predictions else 0 + live_spam_ratio = ( + sum(1 for p in predictions if p == "spam") / len(predictions) + if predictions + else 0 + ) train_spam_ratio = self.training_stats.get("spam_ratio", 0.5) ratio_diff = abs(live_spam_ratio - train_spam_ratio) if ratio_diff > self.threshold_spam_ratio: - alerts.append({ - "type": "prediction_distribution", - "message": f"Spam ratio shifted: train={train_spam_ratio:.3f} → live={live_spam_ratio:.3f} (Δ={ratio_diff:.3f})", - "severity": "warning" if ratio_diff < 0.3 else "critical", - }) + alerts.append( + { + "type": "prediction_distribution", + "message": f"Spam ratio shifted: train={train_spam_ratio:.3f} → live={live_spam_ratio:.3f} (Δ={ratio_diff:.3f})", + "severity": "warning" if ratio_diff < 0.3 else "critical", + } + ) # 2. Input length drift live_avg_len = np.mean([len(t) for t in texts]) if texts else 0 train_avg_len = self.training_stats.get("avg_text_length", 1) - len_shift = abs(live_avg_len - train_avg_len) / train_avg_len if train_avg_len else 0 + len_shift = ( + abs(live_avg_len - train_avg_len) / train_avg_len if train_avg_len else 0 + ) if len_shift > self.threshold_length_pct: - alerts.append({ - "type": "input_length", - "message": f"Avg text length shifted: train={train_avg_len:.0f} → live={live_avg_len:.0f} ({len_shift:.0%})", - "severity": "warning", - }) + alerts.append( + { + "type": "input_length", + "message": f"Avg text length shifted: train={train_avg_len:.0f} → live={live_avg_len:.0f} ({len_shift:.0%})", + "severity": "warning", + } + ) # 3. Vocabulary drift train_vocab = set(self.training_stats.get("vocab", [])) @@ -93,11 +109,13 @@ def check(self, predictions: list[str], texts: list[str]) -> dict: new_words = live_words - train_vocab drift_ratio = len(new_words) / len(live_words) if live_words else 0 if drift_ratio > self.threshold_vocab_drift: - alerts.append({ - "type": "vocabulary_drift", - "message": f"{drift_ratio:.0%} of live words not in training vocabulary ({len(new_words)} new words)", - "severity": "warning", - }) + alerts.append( + { + "type": "vocabulary_drift", + "message": f"{drift_ratio:.0%} of live words not in training vocabulary ({len(new_words)} new words)", + "severity": "warning", + } + ) result = { "has_drift": len(alerts) > 0, @@ -113,7 +131,9 @@ def check(self, predictions: list[str], texts: list[str]) -> dict: self._log_check(result) if alerts: for a in alerts: - logger.warning("DRIFT ALERT [%s]: %s", a["severity"].upper(), a["message"]) + logger.warning( + "DRIFT ALERT [%s]: %s", a["severity"].upper(), a["message"] + ) else: logger.info("No drift detected (%d predictions checked)", len(predictions)) @@ -122,6 +142,7 @@ def check(self, predictions: list[str], texts: list[str]) -> dict: def _log_check(self, result: dict): os.makedirs(os.path.dirname(DRIFT_LOG), exist_ok=True) from datetime import datetime, timezone + logs = [] if os.path.exists(DRIFT_LOG): with open(DRIFT_LOG) as f: diff --git a/src/models/ensemble.py b/src/models/ensemble.py index 88af8ed..01bc74d 100644 --- a/src/models/ensemble.py +++ b/src/models/ensemble.py @@ -31,6 +31,7 @@ # --------------------------------------------------------------------------- try: from lightgbm import LGBMClassifier + _LGBM_AVAILABLE = True except ImportError: _LGBM_AVAILABLE = False @@ -45,26 +46,45 @@ def _base_estimators() -> list[tuple[str, object]]: # HistGB works on dense arrays; we wrap it so sparse input is handled ( "hgb", - Pipeline([ - ("scaler", MaxAbsScaler()), - ("dense", DenseTransformer()), - ("clf", HistGradientBoostingClassifier( - max_iter=200, learning_rate=0.1, max_leaf_nodes=63, random_state=42 - )), - ]), + Pipeline( + [ + ("scaler", MaxAbsScaler()), + ("dense", DenseTransformer()), + ( + "clf", + HistGradientBoostingClassifier( + max_iter=200, + learning_rate=0.1, + max_leaf_nodes=63, + random_state=42, + ), + ), + ] + ), ), ] if _LGBM_AVAILABLE: - estimators.append(( - "lgbm", - Pipeline([ - ("scaler", MaxAbsScaler()), - ("clf", LGBMClassifier( - n_estimators=300, learning_rate=0.05, - num_leaves=63, n_jobs=-1, random_state=42, verbosity=-1, - )), - ]) - )) + estimators.append( + ( + "lgbm", + Pipeline( + [ + ("scaler", MaxAbsScaler()), + ( + "clf", + LGBMClassifier( + n_estimators=300, + learning_rate=0.05, + num_leaves=63, + n_jobs=-1, + random_state=42, + verbosity=-1, + ), + ), + ] + ), + ) + ) return estimators @@ -99,7 +119,7 @@ def train_ensemble( X_train, y_train, model_dir: str = "models", - kind: str = "voting", # "voting" | "stacking" + kind: str = "voting", # "voting" | "stacking" ) -> object: """ Train and save an ensemble model. diff --git a/src/models/evaluate.py b/src/models/evaluate.py index d01383a..2b7f241 100644 --- a/src/models/evaluate.py +++ b/src/models/evaluate.py @@ -24,20 +24,31 @@ def evaluate_model(model, X_test, y_test, model_name: str) -> dict: metrics = { "model": model_name, "accuracy": round(accuracy_score(y_test, y_pred), 4), - "precision": round(precision_score(y_test, y_pred, pos_label="spam", zero_division=0), 4), - "recall": round(recall_score(y_test, y_pred, pos_label="spam", zero_division=0), 4), - "f1_score": round(f1_score(y_test, y_pred, pos_label="spam", zero_division=0), 4), + "precision": round( + precision_score(y_test, y_pred, pos_label="spam", zero_division=0), 4 + ), + "recall": round( + recall_score(y_test, y_pred, pos_label="spam", zero_division=0), 4 + ), + "f1_score": round( + f1_score(y_test, y_pred, pos_label="spam", zero_division=0), 4 + ), "confusion_matrix": cm, } logger.info( "%s → acc=%.4f prec=%.4f rec=%.4f f1=%.4f", - model_name, metrics["accuracy"], metrics["precision"], - metrics["recall"], metrics["f1_score"], + model_name, + metrics["accuracy"], + metrics["precision"], + metrics["recall"], + metrics["f1_score"], ) return metrics -def evaluate_all(models: dict, X_test, y_test, output_path: str = "models/metrics.json") -> dict: +def evaluate_all( + models: dict, X_test, y_test, output_path: str = "models/metrics.json" +) -> dict: """ Evaluate all models and save results. @@ -60,6 +71,10 @@ def evaluate_all(models: dict, X_test, y_test, output_path: str = "models/metric os.makedirs(os.path.dirname(output_path), exist_ok=True) with open(output_path, "w") as f: json.dump(report, f, indent=2) - logger.info("Best model: %s (F1=%.4f). Metrics saved → %s", - best["model"], best["f1_score"], output_path) + logger.info( + "Best model: %s (F1=%.4f). Metrics saved → %s", + best["model"], + best["f1_score"], + output_path, + ) return report diff --git a/src/models/explainability.py b/src/models/explainability.py index 0ddad34..8000401 100644 --- a/src/models/explainability.py +++ b/src/models/explainability.py @@ -1,22 +1,26 @@ """ Explainability Layer — Provide reasoning and feature importance for predictions using SHAP. """ + import shap import numpy as np from src.preprocessing.text_cleaner import preprocess_text -def get_prediction_explanation(text: str, prediction: str, security_results: dict, model=None, vectorizer=None) -> dict: + +def get_prediction_explanation( + text: str, prediction: str, security_results: dict, model=None, vectorizer=None +) -> dict: """ Generate a human-readable explanation and SHAP feature importances. """ clean_text = preprocess_text(text) - + # 1. Extract security logic (Deterministic rules) sec_reasons = security_results.get("threat_reasons", []) risk_level = security_results.get("risk_level", "LOW") - + shap_keywords = [] - + # 2. Extract ML Feature Importance (SHAP) # This requires the trained scikit-learn pipeline (vectorizer + model) if model and vectorizer: @@ -26,25 +30,24 @@ def get_prediction_explanation(text: str, prediction: str, security_results: dic explainer = shap.LinearExplainer(model, vectorizer.transform([""])) vec_text = vectorizer.transform([clean_text]) shap_values = explainer.shap_values(vec_text) - + # Get feature names from vectorizer feature_names = np.array(vectorizer.get_feature_names_out()) - + # Get top contributing words if isinstance(shap_values, list): # For multi-class (or some binary configs), grab the positive class vals = shap_values[1][0] if len(shap_values) > 1 else shap_values[0][0] else: vals = shap_values[0] - + # Sort by absolute impact top_indices = np.argsort(-np.abs(vals))[:5] for idx in top_indices: if vals[idx] != 0: - shap_keywords.append({ - "word": feature_names[idx], - "impact": round(vals[idx], 4) - }) + shap_keywords.append( + {"word": feature_names[idx], "impact": round(vals[idx], 4)} + ) except Exception as e: # Fallback if SHAP fails or model is incompatible pass @@ -58,8 +61,10 @@ def get_prediction_explanation(text: str, prediction: str, security_results: dic if shap_keywords: top_words = [kw["word"] for kw in shap_keywords if kw["impact"] > 0][:3] if top_words: - factors.append(f"Suspicious semantic patterns identified around: {', '.join(top_words)}.") - + factors.append( + f"Suspicious semantic patterns identified around: {', '.join(top_words)}." + ) + if not factors: summary += "Based on deep model confidence patterns." else: @@ -71,5 +76,5 @@ def get_prediction_explanation(text: str, prediction: str, security_results: dic "summary": summary, "risk_level": risk_level, "shap_analysis": shap_keywords, - "security_threats": sec_reasons + "security_threats": sec_reasons, } diff --git a/src/models/feedback.py b/src/models/feedback.py index cdef6c9..6f0da74 100644 --- a/src/models/feedback.py +++ b/src/models/feedback.py @@ -75,7 +75,12 @@ def save_feedback( session.add(db_feedback) session.commit() entry["id"] = db_feedback.id - logger.info("Feedback recorded: id=%s, predicted=%s, correct=%s", db_feedback.id, predicted_label, correct_label) + logger.info( + "Feedback recorded: id=%s, predicted=%s, correct=%s", + db_feedback.id, + predicted_label, + correct_label, + ) except Exception as e: session.rollback() logger.error("Failed to save feedback to database: %s", e) @@ -88,6 +93,7 @@ def load_feedback(): """Load all feedback entries as a list of dicts.""" _ensure_file() import pandas as pd + df = pd.read_csv(FEEDBACK_PATH) logger.info("Loaded %d feedback entries", len(df)) return df @@ -106,6 +112,7 @@ def merge_feedback_into_dataset(original_csv: str, output_csv: str) -> str: Uses feedback's correct_label as the ground truth. """ import pandas as pd + _ensure_file() original = pd.read_csv(original_csv) @@ -119,5 +126,10 @@ def merge_feedback_into_dataset(original_csv: str, output_csv: str) -> str: ) merged = pd.concat([original, new_rows], ignore_index=True) merged.to_csv(output_csv, index=False) - logger.info("Merged %d feedback rows into dataset → %s (%d total)", len(new_rows), output_csv, len(merged)) + logger.info( + "Merged %d feedback rows into dataset → %s (%d total)", + len(new_rows), + output_csv, + len(merged), + ) return output_csv diff --git a/src/models/hyperparameter_tuning.py b/src/models/hyperparameter_tuning.py index 14e4ec1..76ce45b 100644 --- a/src/models/hyperparameter_tuning.py +++ b/src/models/hyperparameter_tuning.py @@ -77,6 +77,7 @@ def _to_dense_if_needed(X, model_name: str): if model_name in _NEEDS_DENSE: import scipy.sparse as sp + if sp.issparse(X): return X.toarray() return X @@ -117,7 +118,9 @@ def tune_model( rs.fit(X, y_train) logger.info( "Tuned %s → best_score=%.4f params=%s", - model_name, rs.best_score_, rs.best_params_, + model_name, + rs.best_score_, + rs.best_params_, ) return { "best_params": rs.best_params_, diff --git a/src/models/optimized_pipeline.py b/src/models/optimized_pipeline.py index 1b83c3a..a05852d 100644 --- a/src/models/optimized_pipeline.py +++ b/src/models/optimized_pipeline.py @@ -15,11 +15,13 @@ # Pre-load stopwords once try: from nltk.corpus import stopwords as _sw + _STOPWORDS = frozenset(_sw.words("english")) except Exception: _STOPWORDS = frozenset() from nltk.stem import PorterStemmer + _STEMMER = PorterStemmer() @@ -33,7 +35,11 @@ def __init__(self, vectorizer_path: str, model_path: str): self.vectorizer = joblib.load(vectorizer_path) self.model = joblib.load(model_path) self._has_proba = hasattr(self.model, "predict_proba") - logger.info("OptimizedPipeline loaded: vectorizer=%s model=%s", vectorizer_path, model_path) + logger.info( + "OptimizedPipeline loaded: vectorizer=%s model=%s", + vectorizer_path, + model_path, + ) def _preprocess(self, text: str) -> str: """Optimized preprocessing with compiled regex and cached stopwords.""" @@ -63,7 +69,9 @@ def predict_batch(self, texts: list[str]) -> list[dict]: if self._has_proba: probas = self.model.predict_proba(features) for pred, proba in zip(predictions, probas): - results.append({"prediction": pred, "confidence": round(float(max(proba)), 4)}) + results.append( + {"prediction": pred, "confidence": round(float(max(proba)), 4)} + ) else: for pred in predictions: results.append({"prediction": pred, "confidence": None}) diff --git a/src/models/predict.py b/src/models/predict.py index 015b5c2..12a8b93 100644 --- a/src/models/predict.py +++ b/src/models/predict.py @@ -40,14 +40,17 @@ def _get_vector_db() -> VectorSearchDB: _vector_db.initialize() return _vector_db + @functools.lru_cache(maxsize=1) def _load_vectorizer(path: str): return joblib.load(path) + @functools.lru_cache(maxsize=5) def _load_model(path: str): return joblib.load(path) + _BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "models") @@ -69,7 +72,9 @@ def _sanitize_model_name(model_name: str | None) -> str: candidate = _default_model_name() if model_name is None else model_name allowed = _allowed_model_names() if candidate not in allowed: - raise ValueError(f"Unsupported model '{candidate}'. Allowed models: {sorted(allowed)}") + raise ValueError( + f"Unsupported model '{candidate}'. Allowed models: {sorted(allowed)}" + ) return next(name for name in allowed if name == candidate) @@ -93,7 +98,9 @@ def _default_model_name() -> str: return "naive_bayes" -def predict_email(text: str, model_name: str | None = None, raw_headers: str = "") -> dict: +def predict_email( + text: str, model_name: str | None = None, raw_headers: str = "" +) -> dict: """ Classify a single email text with security scanning. @@ -135,6 +142,7 @@ def predict_email(text: str, model_name: str | None = None, raw_headers: str = " ml_confidence = float(max(proba)) elif hasattr(model, "decision_function"): import numpy as np + df = model.decision_function(features)[0] ml_confidence = float(1.0 / (1.0 + np.exp(-abs(df)))) @@ -144,7 +152,9 @@ def predict_email(text: str, model_name: str | None = None, raw_headers: str = " dl_model = _get_dl_model() dl_result = dl_model.predict(text) except Exception as e: - logger.warning("Deep learning model unavailable, falling back to ML-only: %s", e) + logger.warning( + "Deep learning model unavailable, falling back to ML-only: %s", e + ) # 4. Run Vector Similarity Search similar_threats = [] @@ -171,24 +181,32 @@ def predict_email(text: str, model_name: str | None = None, raw_headers: str = " # Vector similarity boost: if similar known phishing emails found if similar_threats: - spam_matches = sum(1 for m in similar_threats - if m.get("metadata", {}).get("label") == "spam" - and m.get("distance", 999) < 0.5) + spam_matches = sum( + 1 + for m in similar_threats + if m.get("metadata", {}).get("label") == "spam" + and m.get("distance", 999) < 0.5 + ) if spam_matches >= 2: final_prediction = "spam" - logger.info("Vector similarity override: %d close phishing matches found", spam_matches) + logger.info( + "Vector similarity override: %d close phishing matches found", + spam_matches, + ) # Security rule override (highest priority) if security_analysis["risk_score"] > 80: final_prediction = "spam" - logger.info("Security Rule Override: Flagging as SPAM due to high risk score (%d)", - security_analysis["risk_score"]) + logger.info( + "Security Rule Override: Flagging as SPAM due to high risk score (%d)", + security_analysis["risk_score"], + ) # 6. Explainability (SHAP with model + vectorizer passed in) from src.models.explainability import get_prediction_explanation + explanation = get_prediction_explanation( - text, final_prediction, security_analysis, - model=model, vectorizer=vectorizer + text, final_prediction, security_analysis, model=model, vectorizer=vectorizer ) result = { @@ -210,12 +228,17 @@ def predict_email(text: str, model_name: str | None = None, raw_headers: str = " "threat_reasons": security_analysis["threat_reasons"], "threat_explanation": explanation["summary"], "shap_analysis": explanation.get("shap_analysis", []), - "security_details": security_analysis["details"] + "security_details": security_analysis["details"], } - logger.info("Prediction: %s (risk=%d, ml=%s, dl=%s, model=%s)", - final_prediction, security_analysis["risk_score"], - ml_prediction, dl_result.get("prediction"), model_name) + logger.info( + "Prediction: %s (risk=%d, ml=%s, dl=%s, model=%s)", + final_prediction, + security_analysis["risk_score"], + ml_prediction, + dl_result.get("prediction"), + model_name, + ) return result diff --git a/src/models/registry.py b/src/models/registry.py index 6fada11..0b26b6e 100644 --- a/src/models/registry.py +++ b/src/models/registry.py @@ -91,7 +91,9 @@ def list_models(status: str | None = None) -> list[dict]: def get_active_model(model_name: str) -> dict | None: """Get the latest active version of a model.""" entries = _load_registry() - candidates = [e for e in entries if e["model_name"] == model_name and e["status"] == "active"] + candidates = [ + e for e in entries if e["model_name"] == model_name and e["status"] == "active" + ] return max(candidates, key=lambda e: e["version"]) if candidates else None @@ -107,5 +109,9 @@ def archive_model(model_name: str, version: int): def get_best_model() -> dict | None: """Return the active model with the highest F1 score.""" - entries = [e for e in _load_registry() if e["status"] == "active" and "f1_score" in e.get("metrics", {})] + entries = [ + e + for e in _load_registry() + if e["status"] == "active" and "f1_score" in e.get("metrics", {}) + ] return max(entries, key=lambda e: e["metrics"]["f1_score"]) if entries else None diff --git a/src/models/train_models.py b/src/models/train_models.py index 951723c..a9e4756 100644 --- a/src/models/train_models.py +++ b/src/models/train_models.py @@ -33,21 +33,27 @@ from sklearn.base import BaseEstimator, TransformerMixin from src.utils.logger import logger + class DenseTransformer(BaseEstimator, TransformerMixin): """Simple transformer that converts sparse matrices to dense arrays.""" + def fit(self, X, y=None): return self + def transform(self, X): import scipy.sparse as sp + if sp.issparse(X): return X.toarray() return X + # --------------------------------------------------------------------------- # Optional LightGBM # --------------------------------------------------------------------------- try: from lightgbm import LGBMClassifier + _LGBM_AVAILABLE = True except ImportError: _LGBM_AVAILABLE = False @@ -60,35 +66,56 @@ def transform(self, X): # Factory functions — build estimators (some wrapped in sparse-safe Pipelines) # --------------------------------------------------------------------------- + def _make_naive_bayes(): return MultinomialNB(alpha=0.1) + def _make_logistic_regression(): return LogisticRegression(C=5.0, solver="saga", max_iter=1000, random_state=42) + def _make_svm(): # Wrap in CalibratedClassifierCV so predict_proba is available return CalibratedClassifierCV(LinearSVC(C=1.0, max_iter=10_000), cv=3) + def _make_random_forest(): # MaxAbsScaler preserves sparsity → no dense conversion needed - return Pipeline([ - ("scaler", MaxAbsScaler()), - ("clf", RandomForestClassifier( - n_estimators=200, max_depth=30, min_samples_leaf=2, - n_jobs=-1, random_state=42, - )), - ]) + return Pipeline( + [ + ("scaler", MaxAbsScaler()), + ( + "clf", + RandomForestClassifier( + n_estimators=200, + max_depth=30, + min_samples_leaf=2, + n_jobs=-1, + random_state=42, + ), + ), + ] + ) + def _make_gradient_boosting(): # HistGB requires dense arrays. We densify it inside the pipeline. - return Pipeline([ - ("scaler", MaxAbsScaler()), - ("dense", DenseTransformer()), - ("clf", HistGradientBoostingClassifier( - max_iter=200, learning_rate=0.1, max_leaf_nodes=63, random_state=42, - )), - ]) + return Pipeline( + [ + ("scaler", MaxAbsScaler()), + ("dense", DenseTransformer()), + ( + "clf", + HistGradientBoostingClassifier( + max_iter=200, + learning_rate=0.1, + max_leaf_nodes=63, + random_state=42, + ), + ), + ] + ) MODEL_FACTORIES: dict[str, callable] = { @@ -103,14 +130,25 @@ def _make_gradient_boosting(): MODEL_REGISTRY = MODEL_FACTORIES if _LGBM_AVAILABLE: + def _make_lgbm(): - return Pipeline([ - ("scaler", MaxAbsScaler()), - ("clf", LGBMClassifier( - n_estimators=300, learning_rate=0.05, num_leaves=63, - n_jobs=-1, random_state=42, verbosity=-1, - )), - ]) + return Pipeline( + [ + ("scaler", MaxAbsScaler()), + ( + "clf", + LGBMClassifier( + n_estimators=300, + learning_rate=0.05, + num_leaves=63, + n_jobs=-1, + random_state=42, + verbosity=-1, + ), + ), + ] + ) + MODEL_FACTORIES["lgbm"] = _make_lgbm @@ -124,7 +162,9 @@ def split_data(X, y, test_size: float = 0.2, random_state: int = 42): next(iter(set(y))), ) return train_test_split(X, y, test_size=test_size, random_state=random_state) - return train_test_split(X, y, test_size=test_size, random_state=random_state, stratify=y) + return train_test_split( + X, y, test_size=test_size, random_state=random_state, stratify=y + ) def train_model(X_train, y_train, model_name: str, params: dict | None = None): diff --git a/src/models/vector_search.py b/src/models/vector_search.py index 3bb1c7d..48fb645 100644 --- a/src/models/vector_search.py +++ b/src/models/vector_search.py @@ -2,13 +2,17 @@ Vector Similarity Search Module Uses ChromaDB and Sentence-Transformers to find semantically similar known phishing emails. """ + import os import chromadb from sentence_transformers import SentenceTransformer import logging + class VectorSearchDB: - def __init__(self, db_path: str = "./data/chroma_db", model_name: str = "all-MiniLM-L6-v2"): + def __init__( + self, db_path: str = "./data/chroma_db", model_name: str = "all-MiniLM-L6-v2" + ): self.db_path = db_path self.model_name = model_name self.client = None @@ -27,41 +31,39 @@ def add_email(self, doc_id: str, text: str, label: str, metadata: dict = None): """Add a known email (spam or ham) to the vector database.""" if not self.collection: self.initialize() - + embedding = self.encoder.encode(text).tolist() - + meta = {"label": label} if metadata: meta.update(metadata) - + self.collection.add( - ids=[doc_id], - embeddings=[embedding], - documents=[text], - metadatas=[meta] + ids=[doc_id], embeddings=[embedding], documents=[text], metadatas=[meta] ) def search_similar(self, text: str, n_results: int = 3) -> list: """Search for the most semantically similar emails.""" if not self.collection: self.initialize() - + if self.collection.count() == 0: return [] - + embedding = self.encoder.encode(text).tolist() results = self.collection.query( - query_embeddings=[embedding], - n_results=n_results + query_embeddings=[embedding], n_results=n_results ) - + matches = [] if results and "documents" in results and len(results["documents"]) > 0: for idx in range(len(results["documents"][0])): - matches.append({ - "id": results["ids"][0][idx], - "document": results["documents"][0][idx], - "metadata": results["metadatas"][0][idx], - "distance": results["distances"][0][idx] - }) + matches.append( + { + "id": results["ids"][0][idx], + "document": results["documents"][0][idx], + "metadata": results["metadatas"][0][idx], + "distance": results["distances"][0][idx], + } + ) return matches diff --git a/src/preprocessing/text_cleaner.py b/src/preprocessing/text_cleaner.py index 37f5091..04548f7 100644 --- a/src/preprocessing/text_cleaner.py +++ b/src/preprocessing/text_cleaner.py @@ -22,6 +22,7 @@ # --------------------------------------------------------------------------- try: import nltk + for _pkg in ("punkt", "punkt_tab", "stopwords"): nltk.download(_pkg, quiet=True) from nltk.corpus import stopwords as _nltk_stopwords @@ -47,6 +48,7 @@ # Public API — vectorised (operate on a whole Series at once) # ───────────────────────────────────────────────────────────────────────────── + def vectorized_clean(series: pd.Series) -> pd.Series: """ Apply the full cleaning pipeline to a pandas Series of strings. @@ -94,6 +96,7 @@ def _stem_row(text: str) -> str: return series.apply(_stem_row) from multiprocessing import Pool + with Pool(processes=n_jobs) as pool: stemmed = pool.map(_stem_row, series.tolist()) return pd.Series(stemmed, index=series.index) @@ -103,11 +106,12 @@ def _stem_row(text: str) -> str: # DataFrame-level entry point (called by the pipeline) # ───────────────────────────────────────────────────────────────────────────── + def preprocess_dataframe( df: pd.DataFrame, text_column: str = "text", remove_stops: bool = True, - stem: bool = False, # off by default — adds time without much gain + stem: bool = False, # off by default — adds time without much gain n_jobs: int = 1, ) -> pd.DataFrame: """ @@ -139,6 +143,7 @@ def preprocess_dataframe( # Legacy single-string API (kept for backwards compatibility with tests/CLI) # ───────────────────────────────────────────────────────────────────────────── + def preprocess_text(text: str) -> str: """Process a single string through the cleaning pipeline.""" s = pd.Series([text]) @@ -151,19 +156,24 @@ def preprocess_text(text: str) -> str: def to_lowercase(text: str) -> str: return text.lower() + def remove_punctuation(text: str) -> str: s = _RE_PUNCT.sub(" ", text) return _RE_WHITESPACE.sub(" ", s).strip() + def remove_numbers(text: str) -> str: return _RE_NUMBER.sub("", text) + def tokenize(text: str) -> list[str]: return text.split() + def remove_stopwords(tokens: list[str]) -> list[str]: return [t for t in tokens if t not in _STOP_WORDS] + def stem_tokens(tokens: list[str]) -> list[str]: if not _NLTK_AVAILABLE: return tokens diff --git a/src/security/alerts.py b/src/security/alerts.py index 7b8e7c9..af3295c 100644 --- a/src/security/alerts.py +++ b/src/security/alerts.py @@ -5,33 +5,35 @@ from src.utils.logger import logger from datetime import datetime + def trigger_security_alert(event_type: str, details: dict): """ Simulate an enterprise alerting system (Email, Slack, Webhook). """ risk_score = details.get("security_risk_score", 0) - + if risk_score >= 90: alert_level = "CRITICAL" elif risk_score >= 75: alert_level = "WARNING" else: - return # No alert for low risk + return # No alert for low risk alert_payload = { "timestamp": datetime.now().isoformat(), "level": alert_level, "type": event_type, "reasons": details.get("threat_reasons", []), - "email_preview": details.get("email_text", "")[:100] + "..." + "email_preview": details.get("email_text", "")[:100] + "...", } - + # In a real system, you'd send this to PagerDuty or an ELK stack. logger.warning("!!! SECURITY ALERT [%s] !!! %s", alert_level, alert_payload) - + # Store in a local alert log for forensic dashboard (Phase 63) _log_alert_to_file(alert_payload) + def _log_alert_to_file(payload: dict): alert_file = "data/security_alerts.log" with open(alert_file, "a") as f: diff --git a/src/security/brand_spoof_detector.py b/src/security/brand_spoof_detector.py index 7fd742d..b0788d4 100644 --- a/src/security/brand_spoof_detector.py +++ b/src/security/brand_spoof_detector.py @@ -16,13 +16,14 @@ "metamask": ["metamask-support", "mextamask"], } + def levenshtein_distance(s1: str, s2: str) -> int: """Native implementation of Levenshtein distance.""" if len(s1) < len(s2): return levenshtein_distance(s2, s1) if not s2: return len(s1) - + previous_row = range(len(s2) + 1) for i, c1 in enumerate(s1): current_row = [i + 1] @@ -34,49 +35,54 @@ def levenshtein_distance(s1: str, s2: str) -> int: previous_row = current_row return previous_row[-1] + def detect_brand_impersonation(text: str, urls: list[str]) -> dict: """ Check if any URLs in the text are mimicking a protected brand. """ domains = get_normalized_domains(urls) - + threats = [] max_similarity = 0.0 - + for domain in domains: # Check against protected brand names - main_part = domain.split('.')[0] - + main_part = domain.split(".")[0] + for brand in PROTECTED_BRANDS: # Skip if it IS the actual brand domain (e.g. apple.com) if main_part == brand: continue - + dist = levenshtein_distance(main_part, brand) similarity = 1 - (dist / max(len(main_part), len(brand))) - + # Threshold for suspicion if similarity > 0.75: - threats.append({ - "brand": brand, - "spoofed_domain": domain, - "similarity": round(similarity, 2) - }) + threats.append( + { + "brand": brand, + "spoofed_domain": domain, + "similarity": round(similarity, 2), + } + ) max_similarity = max(max_similarity, similarity) - + # Also check manual variation list for variant in PROTECTED_BRANDS[brand]: if variant in main_part: - threats.append({ - "brand": brand, - "spoofed_domain": domain, - "reason": "Suspicious brand keyword" - }) + threats.append( + { + "brand": brand, + "spoofed_domain": domain, + "reason": "Suspicious brand keyword", + } + ) max_similarity = max(max_similarity, 0.9) return { "brand_impersonation_detected": len(threats) > 0, "impersonated_brands": list(set([t["brand"] for t in threats])), "threat_details": threats, - "max_brand_similarity": round(max_similarity, 2) + "max_brand_similarity": round(max_similarity, 2), } diff --git a/src/security/domain_graph.py b/src/security/domain_graph.py index 400df3a..31b9e64 100644 --- a/src/security/domain_graph.py +++ b/src/security/domain_graph.py @@ -3,20 +3,21 @@ Maps malicious host networks via shared registrars and infrastructure. """ + class DomainGraphAnalyzer: def __init__(self): # Local mapping table placeholders for nodes and edges self.nodes = {} self.edges = [] - + def add_domain_node(self, domain: str, ip_address: str, registrar: str): """Append node relationships to track clustered IPs.""" self.nodes[domain] = { "type": "domain", - "ip": ip_address, - "registrar": registrar + "ip": ip_address, + "registrar": registrar, } - + def cluster_analysis(self, origin_domain: str) -> dict: """ Traverse network tree to identify if origin_domain resolves to the same @@ -24,13 +25,15 @@ def cluster_analysis(self, origin_domain: str) -> dict: """ if origin_domain not in self.nodes: return {"network_threat": False, "cluster_id": None} - + # Example graph logic target_ip = self.nodes[origin_domain]["ip"] - shared_ip_count = sum(1 for data in self.nodes.values() if data.get("ip") == target_ip) - + shared_ip_count = sum( + 1 for data in self.nodes.values() if data.get("ip") == target_ip + ) + return { - "network_threat": shared_ip_count > 5, + "network_threat": shared_ip_count > 5, "cluster_id": f"IP_Cluster_{target_ip}", - "related_malicious_domains": [] + "related_malicious_domains": [], } diff --git a/src/security/domain_intelligence.py b/src/security/domain_intelligence.py index b18afcd..86bf257 100644 --- a/src/security/domain_intelligence.py +++ b/src/security/domain_intelligence.py @@ -5,6 +5,7 @@ from src.security.threat_intel import check_domain_reputation from src.security.url_normalizer import normalize_url + def analyze_domain_intelligence(urls: list[str]) -> dict: """ Perform deep analysis on domains found in URLs. @@ -13,43 +14,41 @@ def analyze_domain_intelligence(urls: list[str]) -> dict: "max_risk_score": 0, "threat_flags": [], "known_malicious_count": 0, - "details": [] + "details": [], } - + unique_domains = set() for url in urls: norm = normalize_url(url) domain = norm.get("normalized_host") if domain: unique_domains.add(domain) - + for domain in unique_domains: rep = check_domain_reputation(domain) - + domain_threats = [] domain_score = 0 - + if rep["known_threat"]: domain_threats.append(f"Known threat: {rep['reason']}") domain_score += 80 results["known_malicious_count"] += 1 - + # Basic heuristic: Domain length if len(domain) > 30: domain_threats.append("Excessively long domain name") domain_score += 10 - + # Check for numeric domains (very suspicious) # e.g. 192.168.1.1 (already handled in url_features, but check for 123-app.top) - if any(c.isdigit() for c in domain.split('.')[0]): + if any(c.isdigit() for c in domain.split(".")[0]): domain_score += 5 - + if domain_score > 0: - results["details"].append({ - "domain": domain, - "score": domain_score, - "threats": domain_threats - }) + results["details"].append( + {"domain": domain, "score": domain_score, "threats": domain_threats} + ) results["max_risk_score"] = max(results["max_risk_score"], domain_score) results["threat_flags"].extend(domain_threats) diff --git a/src/security/header_analyzer.py b/src/security/header_analyzer.py index 9834006..be878fa 100644 --- a/src/security/header_analyzer.py +++ b/src/security/header_analyzer.py @@ -5,6 +5,7 @@ import re from src.utils.logger import logger + def parse_headers(raw_headers: str) -> dict: """ Very basic parser for email header strings. @@ -13,11 +14,12 @@ def parse_headers(raw_headers: str) -> dict: headers = {} lines = raw_headers.splitlines() for line in lines: - if ':' in line: - key, val = line.split(':', 1) + if ":" in line: + key, val = line.split(":", 1) headers[key.strip().lower()] = val.strip() return headers + def analyze_headers(headers: dict) -> dict: """ Check for common spoofing signals: @@ -30,18 +32,18 @@ def analyze_headers(headers: dict) -> dict: "mismatched_domains": False, "suspicious_relay": False, "missing_security": False, - "threats": [] + "threats": [], } - from_field = headers.get('from', '') - return_path = headers.get('return-path', '') - reply_to = headers.get('reply-to', '') - + from_field = headers.get("from", "") + return_path = headers.get("return-path", "") + reply_to = headers.get("reply-to", "") + # 1. From vs Return-Path check if from_field and return_path: - from_domain = re.search(r'@([\w.-]+)', from_field) - rp_domain = re.search(r'@([\w.-]+)', return_path) - + from_domain = re.search(r"@([\w.-]+)", from_field) + rp_domain = re.search(r"@([\w.-]+)", return_path) + if from_domain and rp_domain and from_domain.group(1) != rp_domain.group(1): results["is_spoofed"] = True results["mismatched_domains"] = True @@ -53,12 +55,12 @@ def analyze_headers(headers: dict) -> dict: results["threats"].append("Reply-To differs from From address") # 3. Check for auth results (simulated) - auth_res = headers.get('authentication-results', '') + auth_res = headers.get("authentication-results", "") if auth_res: - if 'spf=fail' in auth_res.lower() or 'spf=softfail' in auth_res.lower(): + if "spf=fail" in auth_res.lower() or "spf=softfail" in auth_res.lower(): results["is_spoofed"] = True results["threats"].append("SPF authentication failed") - if 'dkim=fail' in auth_res.lower(): + if "dkim=fail" in auth_res.lower(): results["is_spoofed"] = True results["threats"].append("DKIM authentication failed") diff --git a/src/security/homograph_detector.py b/src/security/homograph_detector.py index 99b1671..3493419 100644 --- a/src/security/homograph_detector.py +++ b/src/security/homograph_detector.py @@ -8,17 +8,18 @@ # Common confusable character mappings (Latin <-> Greek/Cyrillic/Etc) # This is a reduced set of the most common substitutions CONFUSABLES = { - 'a': ['α', 'а'], - 'o': ['ο', 'о'], - 'e': ['е', 'ε'], - 'p': ['р'], - 'c': ['с'], - 'y': ['у'], - 'i': ['і', 'ɩ'], - 'k': ['κ'], - 'x': ['х'], + "a": ["α", "а"], + "o": ["ο", "о"], + "e": ["е", "ε"], + "p": ["р"], + "c": ["с"], + "y": ["у"], + "i": ["і", "ɩ"], + "k": ["κ"], + "x": ["х"], } + def detect_homograph(text: str) -> dict: """ Scan text for homograph attacks in domains. @@ -28,32 +29,33 @@ def detect_homograph(text: str) -> dict: """ urls = extract_urls(text) domains = [get_domain(u) for u in urls] - + homograph_detected = False mixed_script_domains = False spoofed_parts = [] - + for domain in domains: # Check for mixed scripts in the domain scripts = set() for char in domain: - if char in ('.', '-'): continue + if char in (".", "-"): + continue try: # Get the script name (e.g., LATIN, GREEK, CYRILLIC) script = unicodedata.name(char).split()[0] scripts.add(script) except: continue - + if len(scripts) > 1: mixed_script_domains = True homograph_detected = True spoofed_parts.append(domain) - + # Punycode check (IDNA) try: - punycode = domain.encode('idna').decode('ascii') - if punycode.startswith('xn--'): + punycode = domain.encode("idna").decode("ascii") + if punycode.startswith("xn--"): homograph_detected = True spoofed_parts.append(f"{domain} (punycode: {punycode})") except: @@ -63,5 +65,5 @@ def detect_homograph(text: str) -> dict: "homograph_detected": homograph_detected, "mixed_script_domain": mixed_script_domains, "spoofed_domains": spoofed_parts, - "homograph_count": len(spoofed_parts) + "homograph_count": len(spoofed_parts), } diff --git a/src/security/obfuscation_detector.py b/src/security/obfuscation_detector.py index 7a3f3f7..4911861 100644 --- a/src/security/obfuscation_detector.py +++ b/src/security/obfuscation_detector.py @@ -7,24 +7,25 @@ from src.utils.logger import logger # Zero-width spaces, joiners, and other invisible markers -INVISIBLE_CHARS_PATTERN = r'[\u200B-\u200D\uFEFF\u00AD\u2060]' +INVISIBLE_CHARS_PATTERN = r"[\u200B-\u200D\uFEFF\u00AD\u2060]" + def detect_obfuscation(text: str) -> dict: """Detect hidden/invisible characters used to bypass spam filters.""" matches = re.findall(INVISIBLE_CHARS_PATTERN, text) count = len(matches) - + # Check for excessive character repetition (another evasion trick) repetition = False - if re.search(r'(.)\1{10,}', text): + if re.search(r"(.)\1{10,}", text): repetition = True - + # Check for excessive punctuation / symbol ratio - symbols = len(re.findall(r'[^\w\s]', text)) + symbols = len(re.findall(r"[^\w\s]", text)) symbol_ratio = symbols / len(text) if len(text) > 0 else 0 - + is_obfuscated = count > 0 or repetition or symbol_ratio > 0.4 - + threats = [] if count > 0: threats.append(f"Detected {count} zero-width / invisible characters") @@ -37,9 +38,10 @@ def detect_obfuscation(text: str) -> dict: "is_obfuscated": is_obfuscated, "invisible_count": count, "symbol_ratio": symbol_ratio, - "threats": threats + "threats": threats, } + if __name__ == "__main__": - test_text = "Free G\u200Bift Ca\u200Brd" + test_text = "Free G\u200bift Ca\u200brd" print(detect_obfuscation(test_text)) diff --git a/src/security/risk_scoring.py b/src/security/risk_scoring.py index d220a9d..c52d187 100644 --- a/src/security/risk_scoring.py +++ b/src/security/risk_scoring.py @@ -16,6 +16,7 @@ from src.utils.config_loader import settings + def calculate_security_risk(text: str, raw_headers: str = "") -> dict: """ Apply multi-layer rule detection and calculate risk score (0-100). @@ -29,25 +30,25 @@ def calculate_security_risk(text: str, raw_headers: str = "") -> dict: behavior_feats = analyze_text_behavior(text) domain_feats = analyze_domain_intelligence(urls) obfuscation_feats = detect_obfuscation(text) - + # Load weights from config w = settings.get("security.weights", {}) - + # Header analysis header_feats = {"is_spoofed": False, "threats": []} if raw_headers: parsed = parse_headers(raw_headers) header_feats = analyze_headers(parsed) - + score = 0 reasons = [] - + # 1. Critical Threats (Known Data) d_threats = domain_feats.get("threat_flags", []) if d_threats and isinstance(d_threats, list): score += w.get("domain_blacklist", 60) reasons.extend(d_threats) - + if header_feats.get("is_spoofed") and isinstance(header_feats.get("threats"), list): score += w.get("header_spoof", 40) reasons.extend(header_feats["threats"]) @@ -60,53 +61,61 @@ def calculate_security_risk(text: str, raw_headers: str = "") -> dict: if homograph_feats["homograph_detected"]: score += w.get("homograph", 40) reasons.append("Homograph (IDN) attack detected") - + if script_feats["mixed_script"]: score += w.get("mixed_script", 25) reasons.append("Mixed Unicode scripts detected") if brand_feats["brand_impersonation_detected"]: score += w.get("brand_spoof", 45) - reasons.append(f"Brand Impersonation detected: {', '.join(brand_feats['impersonated_brands'])}") - + reasons.append( + f"Brand Impersonation detected: {', '.join(brand_feats['impersonated_brands'])}" + ) + if behavior_feats["behavior_score"] > w.get("behavioral_threshold", 40): score += w.get("behavioral_low", 15) reasons.append("Suspicious writing behavior (Urgency/Reward Bait)") - + if url_feats.get("has_cyrillic"): score += w.get("cyrillic_url", 50) - reasons.append("URL contains suspicious Cyrillic characters (potential homograph attack)") + reasons.append( + "URL contains suspicious Cyrillic characters (potential homograph attack)" + ) if url_feats["has_ip_url"]: score += w.get("ip_url", 45) reasons.append("Url contains direct IP address") - + if url_feats["suspicious_tld_count"] > 0: score += w.get("suspicious_tld", 20) * url_feats["suspicious_tld_count"] - reasons.append(f"Suspicious TLD(s) detected: {url_feats['suspicious_tld_count']}") - + reasons.append( + f"Suspicious TLD(s) detected: {url_feats['suspicious_tld_count']}" + ) + if url_feats["has_shortener"]: score += w.get("url_shortener", 15) reasons.append("URL shortener used") - + if url_feats["url_count"] > 3: score += w.get("multiple_urls", 10) reasons.append("Multiple URLs detected (> 3)") # Cap score at 100 risk_score = min(score, 100) - + # Derive risk level from config thresholds - thresholds = settings.get("security.thresholds", {"high_risk": 75, "suspicious": 30}) + thresholds = settings.get( + "security.thresholds", {"high_risk": 75, "suspicious": 30} + ) if risk_score > thresholds["high_risk"]: risk_level = "HIGH RISK" elif risk_score > thresholds["suspicious"]: risk_level = "SUSPICIOUS" else: risk_level = "LOW" - + logger.info("Security Risk Analysis: score=%d, reasons=%s", risk_score, reasons) - + return { "risk_score": risk_score, "risk_level": risk_level, @@ -121,7 +130,7 @@ def calculate_security_risk(text: str, raw_headers: str = "") -> dict: "behavioral_threat": behavior_feats["behavior_score"] > 50, "domain_risk": domain_feats["known_malicious_count"] > 0, "header_spoof": header_feats["is_spoofed"], - "obfuscation": obfuscation_feats["is_obfuscated"] + "obfuscation": obfuscation_feats["is_obfuscated"], }, "details": { "urls": url_feats, @@ -131,19 +140,16 @@ def calculate_security_risk(text: str, raw_headers: str = "") -> dict: "behavior": behavior_feats, "domain_intel": domain_feats, "headers": header_feats, - "obfuscation": obfuscation_feats - } + "obfuscation": obfuscation_feats, + }, } def run_security_rules(text: str, raw_headers: str = "") -> dict: """Fastheuristic filter for rule-based overrides.""" analysis = calculate_security_risk(text, raw_headers) - + # Decision: Auto-flag as spam if risk is very high is_rule_spam = analysis["risk_score"] >= 75 - - return { - "is_rule_spam": is_rule_spam, - "risk_analysis": analysis - } + + return {"is_rule_spam": is_rule_spam, "risk_analysis": analysis} diff --git a/src/security/rules_engine.py b/src/security/rules_engine.py index ea99436..bfd018f 100644 --- a/src/security/rules_engine.py +++ b/src/security/rules_engine.py @@ -3,6 +3,7 @@ Overrides ML logic when specific fatal heuristics are triggered. """ + def run_security_rules(security_flags: dict, risk_score: int) -> dict: """ Decides if a deterministic rule should immediately trigger a Spam classification, @@ -14,25 +15,21 @@ def run_security_rules(security_flags: dict, risk_score: int) -> dict: if security_flags.get("ip_url"): reasons.append("Url contains direct IP address") override_spam = True - + if security_flags.get("homograph"): reasons.append("Homograph visual spoiler detected") override_spam = True - + if security_flags.get("cyrillic_url"): reasons.append("Suspicious Cyrillic characters in URL detected") override_spam = True - if risk_score > 75: reasons.append("Aggregate forensic risk critically high") override_spam = True - + if security_flags.get("mixed_script"): reasons.append("Highly suspicious Unicode mixed-scripts detected") override_spam = True - return { - "rule_spam_flag": override_spam, - "reason": reasons - } + return {"rule_spam_flag": override_spam, "reason": reasons} diff --git a/src/security/script_detector.py b/src/security/script_detector.py index 4b20fd7..f1eb113 100644 --- a/src/security/script_detector.py +++ b/src/security/script_detector.py @@ -5,6 +5,7 @@ import unicodedata from collections import Counter + def detect_script_types(text: str) -> dict: """Analyze the distribution of scripts in the text.""" if not text: @@ -26,11 +27,11 @@ def detect_script_types(text: str) -> dict: counts = Counter(scripts) total = sum(counts.values()) - + latin_count = counts.get("LATIN", 0) greek_count = counts.get("GREEK", 0) cyrillic_count = counts.get("CYRILLIC", 0) - + # Calculate ratios ratios = { "latin_ratio": round(latin_count / total, 3), @@ -39,10 +40,6 @@ def detect_script_types(text: str) -> dict: } # Signal mixed scripts if non-Latin scripts are present alongside Latin - mixed_script = (latin_count > 0 and (greek_count > 0 or cyrillic_count > 0)) - - return { - **ratios, - "mixed_script": mixed_script, - "script_counts": dict(counts) - } + mixed_script = latin_count > 0 and (greek_count > 0 or cyrillic_count > 0) + + return {**ratios, "mixed_script": mixed_script, "script_counts": dict(counts)} diff --git a/src/security/text_behavior.py b/src/security/text_behavior.py index 2578628..c92c204 100644 --- a/src/security/text_behavior.py +++ b/src/security/text_behavior.py @@ -5,15 +5,32 @@ import re URGENCY_WORDS = { - "urgent", "action required", "suspended", "immediately", "act now", - "locked", "security alert", "unauthorized", "expired", "final notice" + "urgent", + "action required", + "suspended", + "immediately", + "act now", + "locked", + "security alert", + "unauthorized", + "expired", + "final notice", } REWARD_WORDS = { - "winner", "prize", "gift card", "reward", "lottery", "cash", - "congratulations", "won", "free", "discount" + "winner", + "prize", + "gift card", + "reward", + "lottery", + "cash", + "congratulations", + "won", + "free", + "discount", } + def analyze_text_behavior(text: str) -> dict: """ Extract behavioral and psychological signals from email text. @@ -22,21 +39,22 @@ def analyze_text_behavior(text: str) -> dict: return {"behavior_score": 0} low_text = text.lower() - + # 1. Count triggers urgency_matches = [w for w in URGENCY_WORDS if w in low_text] reward_matches = [w for w in REWARD_WORDS if w in low_text] - + # 2. Structural signals - exclamation_count = text.count('!') + exclamation_count = text.count("!") caps_ratio = sum(1 for c in text if c.isupper()) / len(text) if len(text) > 0 else 0 - + # 3. Decision scores behavior_score = 0 behavior_score += len(urgency_matches) * 15 behavior_score += len(reward_matches) * 10 behavior_score += min(exclamation_count * 5, 20) - if caps_ratio > 0.3: behavior_score += 20 + if caps_ratio > 0.3: + behavior_score += 20 return { "behavior_score": min(behavior_score, 100), @@ -45,5 +63,5 @@ def analyze_text_behavior(text: str) -> dict: "is_urgent": len(urgency_matches) > 0, "is_reward_bait": len(reward_matches) > 0, "high_caps_flag": caps_ratio > 0.3, - "exclamation_excess": exclamation_count > 3 + "exclamation_excess": exclamation_count > 3, } diff --git a/src/security/threat_intel.py b/src/security/threat_intel.py index 386630c..8b2c0fe 100644 --- a/src/security/threat_intel.py +++ b/src/security/threat_intel.py @@ -9,37 +9,45 @@ DB_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "data", "threat_intel.db") + def init_threat_db(): """Initialise the threat intelligence database.""" os.makedirs(os.path.dirname(DB_PATH), exist_ok=True) conn = sqlite3.connect(DB_PATH) c = conn.cursor() - + # Table for malicious domains - c.execute('''CREATE TABLE IF NOT EXISTS bad_domains - (domain TEXT PRIMARY KEY, reason TEXT, added_at DATETIME)''') - + c.execute( + """CREATE TABLE IF NOT EXISTS bad_domains + (domain TEXT PRIMARY KEY, reason TEXT, added_at DATETIME)""" + ) + # Table for known phishing patterns (regex/keywords) - c.execute('''CREATE TABLE IF NOT EXISTS bad_patterns - (pattern TEXT PRIMARY KEY, type TEXT, added_at DATETIME)''') - + c.execute( + """CREATE TABLE IF NOT EXISTS bad_patterns + (pattern TEXT PRIMARY KEY, type TEXT, added_at DATETIME)""" + ) + # Seed with some defaults defaults = [ - ('earn-money.xyz', 'Phishing TLD variation'), - ('login-verification.top', 'Credential harvesting'), - ('gift-card-claim.click', 'Reward bait domain'), + ("earn-money.xyz", "Phishing TLD variation"), + ("login-verification.top", "Credential harvesting"), + ("gift-card-claim.click", "Reward bait domain"), ] - c.executemany("INSERT OR IGNORE INTO bad_domains VALUES (?, ?, ?)", - [(d, r, datetime.now(timezone.utc).isoformat()) for d, r in defaults]) - + c.executemany( + "INSERT OR IGNORE INTO bad_domains VALUES (?, ?, ?)", + [(d, r, datetime.now(timezone.utc).isoformat()) for d, r in defaults], + ) + conn.commit() conn.close() + def check_domain_reputation(domain: str) -> dict: """Check if a domain exists in our knowledge base.""" if not os.path.exists(DB_PATH): init_threat_db() - + try: clean_domain = validate_domain(domain) except ValueError: @@ -50,11 +58,12 @@ def check_domain_reputation(domain: str) -> dict: c.execute("SELECT reason FROM bad_domains WHERE domain = ?", (clean_domain,)) res = c.fetchone() conn.close() - + if res: return {"known_threat": True, "reason": res[0]} return {"known_threat": False} + def report_malicious_domain(domain: str, reason: str): """Add a domain to the local blocklist after strict validation.""" clean_domain = validate_domain(domain) @@ -62,11 +71,14 @@ def report_malicious_domain(domain: str, reason: str): conn = sqlite3.connect(DB_PATH) c = conn.cursor() - c.execute("INSERT OR REPLACE INTO bad_domains VALUES (?, ?, ?)", - (clean_domain, clean_reason, datetime.now(timezone.utc).isoformat())) + c.execute( + "INSERT OR REPLACE INTO bad_domains VALUES (?, ?, ?)", + (clean_domain, clean_reason, datetime.now(timezone.utc).isoformat()), + ) conn.commit() conn.close() + if __name__ == "__main__": init_threat_db() print("Threat Intel DB initialised.") diff --git a/src/security/url_extractor.py b/src/security/url_extractor.py index f3ae0ed..65959fd 100644 --- a/src/security/url_extractor.py +++ b/src/security/url_extractor.py @@ -7,20 +7,22 @@ # Improved URL detection regex URL_PATTERN = re.compile( - r'(?:http|https)://[^\s/$.?#].[^\s]*' # http/https - r'|www\.[^\s/$.?#].[^\s]*' # www. - r'|[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|io|biz|info|xyz|top|click|gq|tk|ml|ga|cf)\b' # domain-like + r"(?:http|https)://[^\s/$.?#].[^\s]*" # http/https + r"|www\.[^\s/$.?#].[^\s]*" # www. + r"|[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|io|biz|info|xyz|top|click|gq|tk|ml|ga|cf)\b" # domain-like ) + def extract_urls(text: str) -> List[str]: """Extract all URLs and domain-like strings from text.""" return URL_PATTERN.findall(text) + def get_domain(url: str) -> str: """Extract domain part from a URL string.""" domain = url.lower() - if domain.startswith(('http://', 'https://')): - domain = domain.split('//')[1] - if domain.startswith('www.'): + if domain.startswith(("http://", "https://")): + domain = domain.split("//")[1] + if domain.startswith("www."): domain = domain[4:] - return domain.split('/')[0].split('?')[0] + return domain.split("/")[0].split("?")[0] diff --git a/src/security/url_features.py b/src/security/url_features.py index df860bc..e3aa03b 100644 --- a/src/security/url_features.py +++ b/src/security/url_features.py @@ -5,38 +5,61 @@ import re from src.security.url_extractor import extract_urls, get_domain -SUSPICIOUS_TLDS = {".xyz", ".top", ".click", ".gq", ".tk", ".ml", ".ga", ".cf", ".bid", ".men"} -SHORTENERS = {"bit.ly", "tinyurl.com", "t.co", "goo.gl", "ow.ly", "is.gd", "buff.ly", "rebrand.ly"} +SUSPICIOUS_TLDS = { + ".xyz", + ".top", + ".click", + ".gq", + ".tk", + ".ml", + ".ga", + ".cf", + ".bid", + ".men", +} +SHORTENERS = { + "bit.ly", + "tinyurl.com", + "t.co", + "goo.gl", + "ow.ly", + "is.gd", + "buff.ly", + "rebrand.ly", +} # Regex for IPv4 in URLs -IP_PATTERN = re.compile(r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b') +IP_PATTERN = re.compile(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b") # Regex for Cyrillic characters (basic range U+0400 to U+04FF) -CYRILLIC_PATTERN = re.compile(r'[\u0400-\u04FF]') +CYRILLIC_PATTERN = re.compile(r"[\u0400-\u04FF]") + def get_url_security_features(text: str) -> dict: """Analyze text for URL-based security features.""" urls = extract_urls(text) domains = [get_domain(u) for u in urls] - + features = { "url_count": len(urls), "has_ip_url": any(IP_PATTERN.search(u) for u in urls), "has_shortener": any(any(s in d for s in SHORTENERS) for d in domains), - "suspicious_tld_count": sum(1 for d in domains if any(d.endswith(t) for t in SUSPICIOUS_TLDS)), - "max_subdomains": max([d.count('.') for d in domains]) if domains else 0, - "max_hyphens_in_domain": max([d.count('-') for d in domains]) if domains else 0, + "suspicious_tld_count": sum( + 1 for d in domains if any(d.endswith(t) for t in SUSPICIOUS_TLDS) + ), + "max_subdomains": max([d.count(".") for d in domains]) if domains else 0, + "max_hyphens_in_domain": max([d.count("-") for d in domains]) if domains else 0, "long_domain_detected": any(len(d) > 25 for d in domains), - "has_cyrillic": any(CYRILLIC_PATTERN.search(u) for u in urls) + "has_cyrillic": any(CYRILLIC_PATTERN.search(u) for u in urls), } - + # Suspicious flag for rules engine features["is_url_suspicious"] = ( - features["has_ip_url"] or - features["has_cyrillic"] or - features["suspicious_tld_count"] > 0 or - features["url_count"] > 3 or - features["max_subdomains"] > 2 + features["has_ip_url"] + or features["has_cyrillic"] + or features["suspicious_tld_count"] > 0 + or features["url_count"] > 3 + or features["max_subdomains"] > 2 ) - + return features diff --git a/src/security/url_normalizer.py b/src/security/url_normalizer.py index eec3a48..fdab7d6 100644 --- a/src/security/url_normalizer.py +++ b/src/security/url_normalizer.py @@ -5,6 +5,7 @@ import re from urllib.parse import urlparse, unquote + def normalize_url(url: str) -> dict: """ Standardize a URL for reliable threat analysis. @@ -15,38 +16,45 @@ def normalize_url(url: str) -> dict: try: # Decode URL-encoded characters (e.g., %20 -> space) url = unquote(url).strip() - + # Add scheme if missing for parsing - if not url.startswith(('http://', 'https://')): - url = 'http://' + url - + if not url.startswith(("http://", "https://")): + url = "http://" + url + parsed = urlparse(url) - + # Basic normalization host = parsed.netloc.lower() - if host.startswith('www.'): + if host.startswith("www."): host = host[4:] - + # Remove default ports - if ':' in host: - host, port = host.split(':', 1) - if (parsed.scheme == 'http' and port == '80') or (parsed.scheme == 'https' and port == '443'): - pass # keep host as is + if ":" in host: + host, port = host.split(":", 1) + if (parsed.scheme == "http" and port == "80") or ( + parsed.scheme == "https" and port == "443" + ): + pass # keep host as is else: host = f"{host}:{port}" - - path = parsed.path.rstrip('/') - + + path = parsed.path.rstrip("/") + return { "original": url, "normalized_host": host, "path": path, "scheme": parsed.scheme, "query": parsed.query, - "full_normalized": f"{parsed.scheme}://{host}{path}" + "full_normalized": f"{parsed.scheme}://{host}{path}", } except Exception: - return {"original": url, "normalized_host": url, "error": "Normalization failed"} + return { + "original": url, + "normalized_host": url, + "error": "Normalization failed", + } + def get_normalized_domains(urls: list[str]) -> list[str]: """Helper to get list of clean domains.""" diff --git a/src/stream/consumer.py b/src/stream/consumer.py index 5cfc285..eccaab5 100644 --- a/src/stream/consumer.py +++ b/src/stream/consumer.py @@ -6,12 +6,15 @@ KAFKA_BROKER_URL — Kafka bootstrap server address (default: localhost:9092 for dev only) KAFKA_TOPIC — Topic to consume from (default: email_ingest) """ + import os from src.utils.logger import logger from src.models.predict import predict_email # Read broker configuration from environment — never hardcode in source -_DEFAULT_BROKER = "localhost:9092" # Dev fallback only; override via KAFKA_BROKER_URL in prod +_DEFAULT_BROKER = ( + "localhost:9092" # Dev fallback only; override via KAFKA_BROKER_URL in prod +) _DEFAULT_TOPIC = "email_ingest" @@ -21,7 +24,9 @@ def __init__( broker_url: str = None, topic: str = None, ): - self.broker_url = broker_url or os.environ.get("KAFKA_BROKER_URL", _DEFAULT_BROKER) + self.broker_url = broker_url or os.environ.get( + "KAFKA_BROKER_URL", _DEFAULT_BROKER + ) self.topic = topic or os.environ.get("KAFKA_TOPIC", _DEFAULT_TOPIC) if self.broker_url == _DEFAULT_BROKER: @@ -31,7 +36,11 @@ def __init__( _DEFAULT_BROKER, ) - logger.info("Initialized stream consumer for topic %s at %s", self.topic, self.broker_url) + logger.info( + "Initialized stream consumer for topic %s at %s", + self.topic, + self.broker_url, + ) def listen(self): """Simulate consuming from a high-throughput message queue.""" @@ -46,7 +55,10 @@ def process(self, email_text: str): """Invoke classification inference on stream ingestion.""" res = predict_email(email_text) if res["prediction"] == "spam": - logger.warning("[STREAM] Intercepted SPAM email! Risk: %d", res.get("security_risk_score", 0)) + logger.warning( + "[STREAM] Intercepted SPAM email! Risk: %d", + res.get("security_risk_score", 0), + ) if __name__ == "__main__": diff --git a/src/utils/anonymizer.py b/src/utils/anonymizer.py index 42dbe67..61f4f6b 100644 --- a/src/utils/anonymizer.py +++ b/src/utils/anonymizer.py @@ -7,7 +7,9 @@ # Regex patterns for common PII _EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b") -_PHONE_RE = re.compile(r"\b(?:\+?\d{1,3}[-.\s]?)?\(?\d{2,4}\)?[-.\s]?\d{3,4}[-.\s]?\d{3,4}\b") +_PHONE_RE = re.compile( + r"\b(?:\+?\d{1,3}[-.\s]?)?\(?\d{2,4}\)?[-.\s]?\d{3,4}[-.\s]?\d{3,4}\b" +) _URL_RE = re.compile(r"https?://\S+|www\.\S+") # Simple name pattern — capitalized words that look like names @@ -33,6 +35,7 @@ def anonymize_text(text: str) -> str: def anonymize_dataframe(df, text_column: str = "text"): """Apply anonymization to a DataFrame text column.""" import pandas as pd + df = df.copy() df[text_column] = df[text_column].astype(str).apply(anonymize_text) logger.info("Anonymized %d rows in column '%s'", len(df), text_column) diff --git a/src/utils/compliance.py b/src/utils/compliance.py index 30bcab5..4b9c1e8 100644 --- a/src/utils/compliance.py +++ b/src/utils/compliance.py @@ -11,6 +11,7 @@ COMPLIANCE_LOG = settings.get("compliance.log_path", "logs/compliance.log") RETENTION_DAYS = settings.get("compliance.retention_days", 30) + def log_compliance_event(event_type: str, details: str): """Secure audit log for regulatory compliance (GDPR/ISO).""" os.makedirs(os.path.dirname(COMPLIANCE_LOG), exist_ok=True) @@ -18,27 +19,32 @@ def log_compliance_event(event_type: str, details: str): with open(COMPLIANCE_LOG, "a") as f: f.write(f"{timestamp} | {event_type} | {details}\n") + def enforce_data_retention(): """Delete sensitive email entries older than the retention threshold.""" session = SessionLocal() try: cutoff = datetime.now(timezone.utc) - timedelta(days=RETENTION_DAYS) - + # Count items to be deleted for audit to_delete = session.query(Feedback).filter(Feedback.timestamp < cutoff) count = to_delete.count() - + if count > 0: to_delete.delete(synchronize_session=False) session.commit() - log_compliance_event("DATA_RETENTION", f"Deleted {count} expired records (older than {RETENTION_DAYS} days)") + log_compliance_event( + "DATA_RETENTION", + f"Deleted {count} expired records (older than {RETENTION_DAYS} days)", + ) logger.info("Compliance: Cleaned up %d expired feedback records.", count) - + except Exception as e: session.rollback() logger.error("Data Retention Error: %s", e) finally: session.close() + if __name__ == "__main__": enforce_data_retention() diff --git a/src/utils/config_loader.py b/src/utils/config_loader.py index f946e1e..0255802 100644 --- a/src/utils/config_loader.py +++ b/src/utils/config_loader.py @@ -6,7 +6,10 @@ import yaml from typing import Any -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "config", "config.yaml") +CONFIG_PATH = os.path.join( + os.path.dirname(__file__), "..", "..", "config", "config.yaml" +) + class Config: _instance = None @@ -26,17 +29,18 @@ def _load(self): # Fallback defaults if file missing self._values = { "security": {"thresholds": {"high_risk": 75, "suspicious": 30}}, - "api": {"port": 8000} + "api": {"port": 8000}, } def get(self, key_path: str, default: Any = None) -> Any: try: val = self._values - for part in key_path.split('.'): + for part in key_path.split("."): val = val[part] return val except (KeyError, TypeError): return default + # Global settings instance settings = Config() diff --git a/src/utils/data_loader.py b/src/utils/data_loader.py index 40dc937..95a6306 100644 --- a/src/utils/data_loader.py +++ b/src/utils/data_loader.py @@ -64,8 +64,7 @@ def _fill_template(template: str) -> str: """Fill a template string with random placeholder values.""" return ( - template - .replace("{name}", random.choice(NAMES)) + template.replace("{name}", random.choice(NAMES)) .replace("{item}", random.choice(ITEMS)) .replace("{amount}", str(random.randint(50, 50000))) .replace("{percent}", str(random.randint(5, 80))) @@ -80,8 +79,12 @@ def generate_sample_dataset(path: str, n: int = 1000) -> pd.DataFrame: """ rows = [] for _ in range(n // 2): - rows.append({"text": _fill_template(random.choice(SPAM_TEMPLATES)), "label": "spam"}) - rows.append({"text": _fill_template(random.choice(HAM_TEMPLATES)), "label": "ham"}) + rows.append( + {"text": _fill_template(random.choice(SPAM_TEMPLATES)), "label": "spam"} + ) + rows.append( + {"text": _fill_template(random.choice(HAM_TEMPLATES)), "label": "ham"} + ) random.shuffle(rows) df = pd.DataFrame(rows) os.makedirs(os.path.dirname(path), exist_ok=True) @@ -118,9 +121,7 @@ def load_dataset(path: str) -> pd.DataFrame: # ── text column ────────────────────────────────────────────────────────── if "text" not in df.columns: - matched = next( - (a for a in _TEXT_ALIASES if a in df.columns or a in cols), None - ) + matched = next((a for a in _TEXT_ALIASES if a in df.columns or a in cols), None) if matched is None: # last resort: pick the first string-typed column str_cols = [c for c in df.columns if df[c].dtype == object] diff --git a/src/utils/data_versioning.py b/src/utils/data_versioning.py index 1a029ae..72c9cfc 100644 --- a/src/utils/data_versioning.py +++ b/src/utils/data_versioning.py @@ -79,7 +79,12 @@ def register_version( } entries.append(entry) _save_metadata(entries) - logger.info("Registered dataset version %s (%d samples) → %s", version_tag, len(df), dest_path) + logger.info( + "Registered dataset version %s (%d samples) → %s", + version_tag, + len(df), + dest_path, + ) return entry diff --git a/src/utils/logger.py b/src/utils/logger.py index e90883a..d06ea4c 100644 --- a/src/utils/logger.py +++ b/src/utils/logger.py @@ -15,7 +15,9 @@ SECURITY_LOG_FILE = os.path.join(LOG_DIR, "security_audit.log") -def setup_logger(name: str = "email_classifier", level: int = logging.INFO) -> logging.Logger: +def setup_logger( + name: str = "email_classifier", level: int = logging.INFO +) -> logging.Logger: """Configure and return a logger with rotating file handler and console output.""" os.makedirs(LOG_DIR, exist_ok=True) @@ -32,12 +34,16 @@ def setup_logger(name: str = "email_classifier", level: int = logging.INFO) -> l ) # Rotating file handler (5 MB, keep 5 backups) - file_handler = RotatingFileHandler(LOG_FILE, maxBytes=5_000_000, backupCount=5, encoding="utf-8") + file_handler = RotatingFileHandler( + LOG_FILE, maxBytes=5_000_000, backupCount=5, encoding="utf-8" + ) file_handler.setLevel(level) file_handler.setFormatter(formatter) # Console handler — use UTF-8 to avoid Windows cp1252 encoding errors - stream = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace", line_buffering=True) + stream = io.TextIOWrapper( + sys.stdout.buffer, encoding="utf-8", errors="replace", line_buffering=True + ) console_handler = logging.StreamHandler(stream) console_handler.setLevel(level) console_handler.setFormatter(formatter) @@ -48,7 +54,9 @@ def setup_logger(name: str = "email_classifier", level: int = logging.INFO) -> l return logger_inst -def setup_security_logger(name: str = "security_audit", level: int = logging.INFO) -> logging.Logger: +def setup_security_logger( + name: str = "security_audit", level: int = logging.INFO +) -> logging.Logger: """Configure and return dedicated security audit logger writing to logs/security_audit.log.""" os.makedirs(LOG_DIR, exist_ok=True) @@ -64,12 +72,16 @@ def setup_security_logger(name: str = "security_audit", level: int = logging.INF ) # Rotating file handler (10 MB, keep 10 backups) - file_handler = RotatingFileHandler(SECURITY_LOG_FILE, maxBytes=10_000_000, backupCount=10, encoding="utf-8") + file_handler = RotatingFileHandler( + SECURITY_LOG_FILE, maxBytes=10_000_000, backupCount=10, encoding="utf-8" + ) file_handler.setLevel(level) file_handler.setFormatter(formatter) # Console handler - stream = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace", line_buffering=True) + stream = io.TextIOWrapper( + sys.stdout.buffer, encoding="utf-8", errors="replace", line_buffering=True + ) console_handler = logging.StreamHandler(stream) console_handler.setLevel(level) console_handler.setFormatter(formatter) @@ -85,11 +97,12 @@ def generate_request_id() -> str: return str(uuid.uuid4())[:8] -def log_security_event(event_type: str, client_ip: str, username: str = "anonymous", detail: str = ""): +def log_security_event( + event_type: str, client_ip: str, username: str = "anonymous", detail: str = "" +): """Convenience helper to record structured security audit events.""" security_logger.info( - "EVENT=%s IP=%s USER=%s DETAIL=%s", - event_type, client_ip, username, detail + "EVENT=%s IP=%s USER=%s DETAIL=%s", event_type, client_ip, username, detail ) diff --git a/src/utils/sanitizer.py b/src/utils/sanitizer.py index cbae47b..193861a 100644 --- a/src/utils/sanitizer.py +++ b/src/utils/sanitizer.py @@ -10,9 +10,18 @@ # Regex Patterns _USERNAME_RE = re.compile(r"^[a-zA-Z0-9_-]{3,50}$") -_DOMAIN_RE = re.compile(r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,63}$") +_DOMAIN_RE = re.compile( + r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,63}$" +) _TOKEN_RE = re.compile(r"^[a-zA-Z0-9._-]{1,128}$") -_ALLOWED_MODELS = {"naive_bayes", "svm", "ensemble", "transformers", "deep_learning", "unknown"} +_ALLOWED_MODELS = { + "naive_bayes", + "svm", + "ensemble", + "transformers", + "deep_learning", + "unknown", +} _ALLOWED_LABELS = {"spam", "ham"} # Control Character Pattern (preserves \n, \r, \t, removes ASCII control chars 0-31 except 9,10,13 and 127) @@ -62,7 +71,9 @@ def validate_domain(domain: str) -> str: raise ValueError("Domain is required.") clean = domain.strip().lower() if not _DOMAIN_RE.match(clean): - raise ValueError("Invalid domain name format. Must be a valid FQDN (e.g. malicious-site.com).") + raise ValueError( + "Invalid domain name format. Must be a valid FQDN (e.g. malicious-site.com)." + ) return clean @@ -87,7 +98,9 @@ def validate_model_name(name: Optional[str]) -> Optional[str]: return None clean = name.strip().lower() if clean not in _ALLOWED_MODELS: - raise ValueError(f"Invalid model name '{name}'. Allowed models: {', '.join(sorted(_ALLOWED_MODELS))}") + raise ValueError( + f"Invalid model name '{name}'. Allowed models: {', '.join(sorted(_ALLOWED_MODELS))}" + ) return clean diff --git a/src/utils/secrets.py b/src/utils/secrets.py index 5a1ebec..02a4139 100644 --- a/src/utils/secrets.py +++ b/src/utils/secrets.py @@ -12,6 +12,7 @@ from src.utils.secrets import vault jwt_secret = vault.get("JWT_SECRET") """ + import os import json from src.utils.logger import logger @@ -34,7 +35,9 @@ def _load(self): self.keys.update(json.load(f)) logger.debug("Loaded secrets from '%s'.", self.file_path) except Exception as exc: - logger.warning("Could not read secrets file '%s': %s", self.file_path, exc) + logger.warning( + "Could not read secrets file '%s': %s", self.file_path, exc + ) # 2. Environment variables always win — overlay ALL keys that appear in the JSON, # plus the core required keys. This ensures any new key added to secrets.json @@ -71,7 +74,7 @@ def _validate(self): raise RuntimeError( "FATAL: JWT_SECRET is not set or is using the development placeholder. " "Generate a strong secret with: " - "python -c \"import secrets; print(secrets.token_hex(32))\"" + 'python -c "import secrets; print(secrets.token_hex(32))"' ) elif not secret: # Dev mode: warn loudly and fall back to the placeholder diff --git a/tests/conftest.py b/tests/conftest.py index f1a852d..35bb569 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,9 @@ """Pytest configuration and fixtures.""" + import warnings # Suppress httpx deprecation warning for TestClient warnings.filterwarnings("ignore", category=DeprecationWarning, module="httpx") # Suppress sklearn feature names warning for LGBM -warnings.filterwarnings("ignore", message="X does not have valid feature names") \ No newline at end of file +warnings.filterwarnings("ignore", message="X does not have valid feature names") diff --git a/tests/test_abuse_protection.py b/tests/test_abuse_protection.py index a65178a..9514553 100644 --- a/tests/test_abuse_protection.py +++ b/tests/test_abuse_protection.py @@ -20,41 +20,49 @@ # Honeypot Field Validation (schema-level) # --------------------------------------------------------------------------- + class TestHoneypotValidation: """PredictRequest and RegisterRequest bot_trap field must reject non-empty values.""" def test_predict_request_empty_bot_trap_passes(self): from src.api.schemas import PredictRequest + req = PredictRequest(text="Hello test email", bot_trap="") assert req.bot_trap == "" def test_predict_request_none_bot_trap_passes(self): from src.api.schemas import PredictRequest + req = PredictRequest(text="Hello test email") assert req.bot_trap == "" def test_predict_request_populated_bot_trap_raises(self): from pydantic import ValidationError from src.api.schemas import PredictRequest + with pytest.raises(ValidationError) as exc_info: - PredictRequest(text="Hello test email", bot_trap="I am a bot filling all fields") + PredictRequest( + text="Hello test email", bot_trap="I am a bot filling all fields" + ) errors = exc_info.value.errors() assert any("Bot activity" in str(e) for e in errors) def test_register_request_populated_bot_trap_raises(self): from pydantic import ValidationError from src.api.schemas import RegisterRequest + with pytest.raises(ValidationError) as exc_info: RegisterRequest( username="botuser", password="Secure@123", - bot_trap="auto-filled by script" + bot_trap="auto-filled by script", ) errors = exc_info.value.errors() assert any("Bot" in str(e) for e in errors) def test_register_request_empty_bot_trap_passes(self): from src.api.schemas import RegisterRequest + req = RegisterRequest(username="realuser", password="Secure@123", bot_trap="") assert req.bot_trap == "" @@ -63,28 +71,33 @@ def test_register_request_empty_bot_trap_passes(self): # Batch Size Cap (schema-level) # --------------------------------------------------------------------------- + class TestBatchSizeCap: """BatchPredictRequest must reject lists exceeding 50 emails.""" def test_batch_of_50_emails_passes(self): from src.api.schemas import BatchPredictRequest + req = BatchPredictRequest(emails=["email body"] * 50) assert len(req.emails) == 50 def test_batch_of_51_emails_raises(self): from pydantic import ValidationError from src.api.schemas import BatchPredictRequest + with pytest.raises(ValidationError): BatchPredictRequest(emails=["email body"] * 51) def test_batch_of_1_email_passes(self): from src.api.schemas import BatchPredictRequest + req = BatchPredictRequest(emails=["single email"]) assert len(req.emails) == 1 def test_empty_batch_raises(self): from pydantic import ValidationError from src.api.schemas import BatchPredictRequest + with pytest.raises(ValidationError): BatchPredictRequest(emails=[]) @@ -94,11 +107,23 @@ def test_empty_batch_raises(self): # --------------------------------------------------------------------------- _BOT_USER_AGENTS = { - "sqlmap", "nikto", "nmap", "masscan", "gocurl", "bytespider", - "zgrab", "libwww-perl", "python-urllib", "dirbuster", "w3af", - "acunetix", "nessus", "openvas" + "sqlmap", + "nikto", + "nmap", + "masscan", + "gocurl", + "bytespider", + "zgrab", + "libwww-perl", + "python-urllib", + "dirbuster", + "w3af", + "acunetix", + "nessus", + "openvas", } + def _is_blacklisted_ua(user_agent: str) -> bool: """Mirror of app.py bot detection logic.""" ua_lower = user_agent.lower() @@ -133,10 +158,13 @@ def test_dirbuster_detected(self): assert _is_blacklisted_ua("DirBuster-1.0") is True def test_legitimate_chrome_allowed(self): - assert _is_blacklisted_ua( - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " - "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" - ) is False + assert ( + _is_blacklisted_ua( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" + ) + is False + ) def test_legitimate_api_client_allowed(self): assert _is_blacklisted_ua("PhishShieldClient/1.0") is False @@ -156,10 +184,17 @@ def test_empty_ua_not_blacklisted_separately(self): # Middleware Path Matching Logic # --------------------------------------------------------------------------- + class TestApiPathMatching: """Anti-bot middleware applies to specific API path prefixes only.""" - API_PREFIXES = ("/predict", "/auth/", "/analyze-security", "/export-report", "/feedback") + API_PREFIXES = ( + "/predict", + "/auth/", + "/analyze-security", + "/export-report", + "/feedback", + ) def _is_api_path(self, path: str) -> bool: return any(path.lower().startswith(prefix) for prefix in self.API_PREFIXES) diff --git a/tests/test_api.py b/tests/test_api.py index ddd4640..1e7574a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -3,6 +3,7 @@ import sys import os import warnings + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import pytest @@ -15,6 +16,7 @@ @pytest.fixture def client(): from src.api.app import app + with TestClient(app) as c: yield c @@ -50,19 +52,24 @@ def test_analytics(client): def test_register_and_login(client): import uuid + username = f"testuser_{uuid.uuid4().hex[:8]}" # Password meets policy: 8+ chars, digit, special character good_password = "Secure@99!" # Register - res = client.post("/auth/register", json={"username": username, "password": good_password}) + res = client.post( + "/auth/register", json={"username": username, "password": good_password} + ) assert res.status_code == 200 data = res.json() assert "api_key" in data assert data["api_key"].startswith("pse_"), "API key should use the pse_ prefix" # Login succeeds with correct credentials - res = client.post("/auth/login", json={"username": username, "password": good_password}) + res = client.post( + "/auth/login", json={"username": username, "password": good_password} + ) assert res.status_code == 200 data = res.json() assert "token" in data @@ -76,6 +83,7 @@ def test_register_and_login(client): def test_register_rejects_weak_password(client): """Passwords shorter than 8 chars or lacking digit/special chars must be rejected.""" import uuid + username = f"weakuser_{uuid.uuid4().hex[:8]}" # Too short @@ -83,37 +91,50 @@ def test_register_rejects_weak_password(client): assert res.status_code == 422 # No digit - res = client.post("/auth/register", json={"username": username, "password": "NoDigit!!"}) + res = client.post( + "/auth/register", json={"username": username, "password": "NoDigit!!"} + ) assert res.status_code in (409, 422) # No special character - res = client.post("/auth/register", json={"username": username, "password": "NoSpecial1"}) + res = client.post( + "/auth/register", json={"username": username, "password": "NoSpecial1"} + ) assert res.status_code in (409, 422) def test_account_lockout(client): """After LOGIN_MAX_ATTEMPTS consecutive failures the account is locked.""" import uuid + username = f"lockuser_{uuid.uuid4().hex[:8]}" good_password = "Locked@99!" # Register - res = client.post("/auth/register", json={"username": username, "password": good_password}) + res = client.post( + "/auth/register", json={"username": username, "password": good_password} + ) assert res.status_code == 200 # Exhaust the default 5 attempts for _ in range(5): - res = client.post("/auth/login", json={"username": username, "password": "WrongPass!1"}) + res = client.post( + "/auth/login", json={"username": username, "password": "WrongPass!1"} + ) assert res.status_code == 401 # Next attempt should still return 401 (account locked, same generic error) - res = client.post("/auth/login", json={"username": username, "password": good_password}) + res = client.post( + "/auth/login", json={"username": username, "password": good_password} + ) assert res.status_code == 401 def test_password_reset_no_enumeration(client): """password-reset-request always returns 200, even for unknown usernames.""" - res = client.post("/auth/password-reset-request", json={"username": "nonexistent_user_xyz"}) + res = client.post( + "/auth/password-reset-request", json={"username": "nonexistent_user_xyz"} + ) assert res.status_code == 200 data = res.json() assert "detail" in data @@ -124,12 +145,15 @@ def test_password_reset_no_enumeration(client): def test_password_reset_full_flow(client): """Request a reset token, then use it to set a new password.""" import uuid + username = f"resetuser_{uuid.uuid4().hex[:8]}" original_password = "Original@1!" new_password = "NewSecure@2!" # Register - client.post("/auth/register", json={"username": username, "password": original_password}) + client.post( + "/auth/register", json={"username": username, "password": original_password} + ) # Request reset token res = client.post("/auth/password-reset-request", json={"username": username}) @@ -138,19 +162,26 @@ def test_password_reset_full_flow(client): assert token is not None # Reset password - res = client.post("/auth/password-reset", json={ - "username": username, - "token": token, - "new_password": new_password, - }) + res = client.post( + "/auth/password-reset", + json={ + "username": username, + "token": token, + "new_password": new_password, + }, + ) assert res.status_code == 200 # Old password should no longer work - res = client.post("/auth/login", json={"username": username, "password": original_password}) + res = client.post( + "/auth/login", json={"username": username, "password": original_password} + ) assert res.status_code == 401 # New password should work - res = client.post("/auth/login", json={"username": username, "password": new_password}) + res = client.post( + "/auth/login", json={"username": username, "password": new_password} + ) assert res.status_code == 200 assert "token" in res.json() @@ -158,12 +189,18 @@ def test_password_reset_full_flow(client): def test_password_reset_invalid_token(client): """An invalid reset token must be rejected.""" import uuid - username = f"badtoken_{uuid.uuid4().hex[:8]}" - client.post("/auth/register", json={"username": username, "password": "ValidPwd@1!"}) - res = client.post("/auth/password-reset", json={ - "username": username, - "token": "not-a-real-token", - "new_password": "AnotherPwd@2!", - }) + username = f"badtoken_{uuid.uuid4().hex[:8]}" + client.post( + "/auth/register", json={"username": username, "password": "ValidPwd@1!"} + ) + + res = client.post( + "/auth/password-reset", + json={ + "username": username, + "token": "not-a-real-token", + "new_password": "AnotherPwd@2!", + }, + ) assert res.status_code == 400 diff --git a/tests/test_auth_unit.py b/tests/test_auth_unit.py index d48dbcd..3932a81 100644 --- a/tests/test_auth_unit.py +++ b/tests/test_auth_unit.py @@ -3,6 +3,7 @@ These import auth functions directly, bypassing the full FastAPI app stack (no prometheus, arq, or Redis dependencies required). """ + import os import sys import uuid @@ -20,6 +21,7 @@ def setup_module(module): """Create DB tables before any test in this module runs.""" from src.core.database import init_db + init_db() @@ -38,6 +40,7 @@ def unique_user(): # Registration # --------------------------------------------------------------------------- + class TestRegistration: def test_register_success(self): u = unique_user() @@ -73,6 +76,7 @@ def test_register_with_email(self): # Authentication & lockout # --------------------------------------------------------------------------- + class TestAuthentication: def test_login_success(self): u = unique_user() @@ -124,6 +128,7 @@ def test_verify_api_key_wrong(self): # Password reset # --------------------------------------------------------------------------- + class TestPasswordReset: def test_reset_request_unknown_user_returns_none(self): # Returns None (caller surfaces generic response; no exception) @@ -172,9 +177,11 @@ def test_weak_new_password_rejected(self): # IDOR & Ownership Protection Tests # --------------------------------------------------------------------------- + class TestIDORProtection: def test_log_usage_and_user_log_ownership(self): from src.core.database import SessionLocal, User as DBUser + session = SessionLocal() try: uA_name = unique_user() @@ -186,7 +193,9 @@ def test_log_usage_and_user_log_ownership(self): uB = session.query(DBUser).filter_by(username=uB_name).first() # Log usage for User A - auth_module.log_usage(uA.id, "/predict", '{"text":"hello"}', '{"prediction":"ham"}') + auth_module.log_usage( + uA.id, "/predict", '{"text":"hello"}', '{"prediction":"ham"}' + ) logsA = auth_module.get_user_logs(uA.id) assert len(logsA) > 0 log_id_A = logsA[0]["id"] @@ -205,6 +214,7 @@ def test_log_usage_and_user_log_ownership(self): def test_feedback_ownership_and_deletion(self): from src.core.database import SessionLocal, User as DBUser from src.models.feedback import save_feedback + session = SessionLocal() try: uA_name = unique_user() diff --git a/tests/test_idor_security.py b/tests/test_idor_security.py index f6b33ce..31bc50c 100644 --- a/tests/test_idor_security.py +++ b/tests/test_idor_security.py @@ -1,6 +1,7 @@ """ IDOR Security Tests — Verify strict resource ownership checks on user logs and feedback endpoints. """ + import os import sys import uuid @@ -14,6 +15,7 @@ os.environ.setdefault("ENV", "dev") from src.core.database import init_db + init_db() from src.api.app import app @@ -46,22 +48,34 @@ def test_idor_log_access(): userB, tokenB, keyB = create_user_and_login("userB") # Generate a log for User A by calling predict with User A's token - res = client.post("/predict/batch", json={"emails": ["Phishing alert!"]}, headers={"Authorization": f"Bearer {tokenA}"}) + res = client.post( + "/predict/batch", + json={"emails": ["Phishing alert!"]}, + headers={"Authorization": f"Bearer {tokenA}"}, + ) assert res.status_code == 200 # Retrieve logs for User A - logs_A = client.get("/auth/logs", headers={"Authorization": f"Bearer {tokenA}"}).json() + logs_A = client.get( + "/auth/logs", headers={"Authorization": f"Bearer {tokenA}"} + ).json() assert len(logs_A) > 0 log_id_A = logs_A[0]["id"] # User A can view their own log - res_A = client.get(f"/auth/logs/{log_id_A}", headers={"Authorization": f"Bearer {tokenA}"}) + res_A = client.get( + f"/auth/logs/{log_id_A}", headers={"Authorization": f"Bearer {tokenA}"} + ) assert res_A.status_code == 200 assert res_A.json()["id"] == log_id_A # User B attempts to access User A's log by ID (IDOR Attempt) - res_B = client.get(f"/auth/logs/{log_id_A}", headers={"Authorization": f"Bearer {tokenB}"}) - assert res_B.status_code == 404, "User B should receive 404 Not Found when requesting User A's log" + res_B = client.get( + f"/auth/logs/{log_id_A}", headers={"Authorization": f"Bearer {tokenB}"} + ) + assert ( + res_B.status_code == 404 + ), "User B should receive 404 Not Found when requesting User A's log" def test_idor_feedback_access_and_deletion(): @@ -84,19 +98,31 @@ def test_idor_feedback_access_and_deletion(): feedback_id_A = fb_res.json()["entry"]["id"] # User A can view their feedback - res_get_A = client.get(f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenA}"}) + res_get_A = client.get( + f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenA}"} + ) assert res_get_A.status_code == 200 assert res_get_A.json()["id"] == feedback_id_A # User B attempts to view User A's feedback (IDOR Read Attempt) - res_get_B = client.get(f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenB}"}) - assert res_get_B.status_code == 404, "User B should receive 404 Not Found when reading User A's feedback" + res_get_B = client.get( + f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenB}"} + ) + assert ( + res_get_B.status_code == 404 + ), "User B should receive 404 Not Found when reading User A's feedback" # User B attempts to delete User A's feedback (IDOR Delete Attempt) - res_del_B = client.delete(f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenB}"}) - assert res_del_B.status_code == 404, "User B should receive 404 Not Found when deleting User A's feedback" + res_del_B = client.delete( + f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenB}"} + ) + assert ( + res_del_B.status_code == 404 + ), "User B should receive 404 Not Found when deleting User A's feedback" # User A can delete their own feedback - res_del_A = client.delete(f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenA}"}) + res_del_A = client.delete( + f"/auth/feedback/{feedback_id_A}", headers={"Authorization": f"Bearer {tokenA}"} + ) assert res_del_A.status_code == 200 assert res_del_A.json()["detail"] == "Feedback entry deleted successfully." diff --git a/tests/test_input_validation.py b/tests/test_input_validation.py index 2e20637..c0b713b 100644 --- a/tests/test_input_validation.py +++ b/tests/test_input_validation.py @@ -31,6 +31,7 @@ # 1. Text Sanitization (XSS, Null Bytes, Control Chars) # --------------------------------------------------------------------------- + class TestTextSanitization: def test_html_script_tags_escaped(self): raw = "Hello world!" @@ -72,6 +73,7 @@ def test_max_length_truncation(self): # 2. Username Regex Validation # --------------------------------------------------------------------------- + class TestUsernameValidation: def test_valid_usernames(self): assert validate_username("alice") == "alice" @@ -107,6 +109,7 @@ def test_too_long_username_rejected(self): # 3. Domain Validation # --------------------------------------------------------------------------- + class TestDomainValidation: def test_valid_domains(self): assert validate_domain("example.com") == "example.com" @@ -130,6 +133,7 @@ def test_sql_injection_in_domain_rejected(self): # 4. Token & Whitelist Validation # --------------------------------------------------------------------------- + class TestTokenAndWhitelistValidation: def test_valid_tokens(self): assert validate_token("abc123XYZ_.-") == "abc123XYZ_.-" @@ -160,6 +164,7 @@ def test_invalid_label_rejected(self): # 5. Schema Validation Integration # --------------------------------------------------------------------------- + class TestSchemaValidationIntegration: def test_predict_request_escapes_html(self): req = PredictRequest(text="Check this ") @@ -172,7 +177,11 @@ def test_register_request_rejects_sql_username(self): def test_register_request_rejects_email_newline(self): with pytest.raises(ValidationError): - RegisterRequest(username="validuser", password="Secure@Password1!", email="test@example.com\r\nBcc: evil@hacker.io") + RegisterRequest( + username="validuser", + password="Secure@Password1!", + email="test@example.com\r\nBcc: evil@hacker.io", + ) def test_login_request_rejects_invalid_username(self): with pytest.raises(ValidationError): @@ -180,8 +189,14 @@ def test_login_request_rejects_invalid_username(self): def test_password_reset_rejects_malformed_token(self): with pytest.raises(ValidationError): - PasswordResetSchema(username="validuser", token="token with spaces!", new_password="NewSecure@Pass1!") + PasswordResetSchema( + username="validuser", + token="token with spaces!", + new_password="NewSecure@Pass1!", + ) def test_feedback_request_validates_labels(self): with pytest.raises(ValidationError): - FeedbackRequest(email_text="test", predicted_label="unknown_label", correct_label="ham") + FeedbackRequest( + email_text="test", predicted_label="unknown_label", correct_label="ham" + ) diff --git a/tests/test_model.py b/tests/test_model.py index 9a77e53..c70274e 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -2,6 +2,7 @@ import sys import os + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import pytest diff --git a/tests/test_new_modules.py b/tests/test_new_modules.py index 2c1e42e..221ccb5 100644 --- a/tests/test_new_modules.py +++ b/tests/test_new_modules.py @@ -13,9 +13,11 @@ from src.models.explainability import get_prediction_explanation from src.core.worker import check_drift, send_security_alert + def test_database_orm(): """Test database tables and SQLAlchemy operations.""" import uuid + unique = uuid.uuid4().hex[:8] uname = f"test_db_{unique}" akey = f"key_{unique}" @@ -34,14 +36,21 @@ def test_database_orm(): assert db_user.api_key == akey # 2. Test UsageLog - log = UsageLog(user_id=user.id, endpoint="/predict", request_body="{}", response_body="{}") + log = UsageLog( + user_id=user.id, endpoint="/predict", request_body="{}", response_body="{}" + ) session.add(log) session.commit() assert log.id is not None assert log.user.username == uname # 3. Test Feedback - feedback = Feedback(email_text="test text", predicted_label="ham", correct_label="spam", model_used="naive_bayes") + feedback = Feedback( + email_text="test text", + predicted_label="ham", + correct_label="spam", + model_used="naive_bayes", + ) session.add(feedback) session.commit() assert feedback.id is not None @@ -52,70 +61,80 @@ def test_database_orm(): finally: # Cleanup test data session.rollback() - session.query(UsageLog).filter_by(user_id=user.id if 'user' in dir() else -1).delete() - session.query(Feedback).filter_by(id=feedback.id if 'feedback' in dir() else -1).delete() + session.query(UsageLog).filter_by( + user_id=user.id if "user" in dir() else -1 + ).delete() + session.query(Feedback).filter_by( + id=feedback.id if "feedback" in dir() else -1 + ).delete() session.query(User).filter_by(username=uname).delete() session.commit() session.close() + def test_deep_learning_mocked(): """Test DeepLearningModel class with a mocked transformer pipeline.""" with patch("src.models.deep_learning.pipeline") as mock_pipeline: mock_classifier = MagicMock() mock_classifier.return_value = [{"label": "phishing", "score": 0.95}] mock_pipeline.return_value = mock_classifier - + model = DeepLearningModel() model.load() assert model.is_loaded - + res = model.predict("click this phishing link") assert res["prediction"] == "spam" assert res["confidence"] == 0.95 assert res["model_used"] == "ealvaradob/bert-finetuned-phishing" + def test_vector_search_mocked(): """Test VectorSearchDB with mocked ChromaDB and SentenceTransformer.""" - with patch("src.models.vector_search.chromadb.PersistentClient") as mock_client_cls, \ - patch("src.models.vector_search.SentenceTransformer") as mock_transformer_cls: - + with patch( + "src.models.vector_search.chromadb.PersistentClient" + ) as mock_client_cls, patch( + "src.models.vector_search.SentenceTransformer" + ) as mock_transformer_cls: + mock_client = MagicMock() mock_collection = MagicMock() mock_client.get_or_create_collection.return_value = mock_collection mock_client_cls.return_value = mock_client - + mock_encoder = MagicMock() mock_encoder.encode.return_value = MagicMock(tolist=lambda: [0.1, 0.2, 0.3]) mock_transformer_cls.return_value = mock_encoder - + vdb = VectorSearchDB(db_path="/mock/path") vdb.initialize() - + # Test add email vdb.add_email("doc1", "phishing text", "spam") mock_collection.add.assert_called_once() - + # Test search similar mock_collection.count.return_value = 1 mock_collection.query.return_value = { "ids": [["doc1"]], "documents": [["phishing text"]], "metadatas": [[{"label": "spam"}]], - "distances": [[0.123]] + "distances": [[0.123]], } - + results = vdb.search_similar("similar text", n_results=1) assert len(results) == 1 assert results[0]["id"] == "doc1" assert results[0]["distance"] == 0.123 + def test_explainability_with_shap(): """Test explainability layer with and without SHAP/model inputs.""" # Without model/vectorizer explanation = get_prediction_explanation( text="normal message", prediction="ham", - security_results={"threat_reasons": [], "risk_level": "LOW"} + security_results={"threat_reasons": [], "risk_level": "LOW"}, ) assert explanation["risk_level"] == "LOW" assert "Classified as HAM" in explanation["summary"] @@ -125,30 +144,32 @@ def test_explainability_with_shap(): mock_vectorizer = MagicMock() mock_vectorizer.transform.return_value = "vectorized_text" mock_vectorizer.get_feature_names_out.return_value = ["click", "free", "money"] - + with patch("src.models.explainability.shap.LinearExplainer") as mock_explainer_cls: mock_explainer = MagicMock() mock_explainer.shap_values.return_value = [[[0.1, 0.5, 0.0]]] mock_explainer_cls.return_value = mock_explainer - + explanation = get_prediction_explanation( text="free money", prediction="spam", security_results={"threat_reasons": ["high risk"], "risk_level": "HIGH"}, model=mock_model, - vectorizer=mock_vectorizer + vectorizer=mock_vectorizer, ) assert explanation["risk_level"] == "HIGH" assert len(explanation["shap_analysis"]) > 0 + def test_worker_tasks(): """Test arq background worker tasks with mocks.""" mock_drift = MagicMock() mock_alert = MagicMock() async def _run(): - with patch("src.core.worker.drift_monitor.check", mock_drift), \ - patch("src.core.worker.trigger_security_alert", mock_alert): + with patch("src.core.worker.drift_monitor.check", mock_drift), patch( + "src.core.worker.trigger_security_alert", mock_alert + ): await check_drift(None, "email text", "ham") mock_drift.assert_called_once_with(["email text"], ["ham"]) diff --git a/tests/test_preprocessing.py b/tests/test_preprocessing.py index 07d1e54..24e96e7 100644 --- a/tests/test_preprocessing.py +++ b/tests/test_preprocessing.py @@ -2,6 +2,7 @@ import sys import os + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from src.preprocessing.text_cleaner import ( diff --git a/tests/test_secrets_audit.py b/tests/test_secrets_audit.py index 76beff0..869d2a9 100644 --- a/tests/test_secrets_audit.py +++ b/tests/test_secrets_audit.py @@ -21,6 +21,7 @@ # SecretsVault — env var override behaviour # --------------------------------------------------------------------------- + class TestSecretsVaultEnvPriority: """Environment variables must always win over values from secrets.json.""" @@ -32,6 +33,7 @@ def test_env_var_overrides_json_value(self, tmp_path): os.environ["JWT_SECRET"] = "env-value" try: from src.utils.secrets import SecretsVault + vault = SecretsVault(secrets_file=str(secrets_file)) assert vault.get("JWT_SECRET") == "env-value" finally: @@ -45,6 +47,7 @@ def test_json_value_used_when_env_absent(self, tmp_path): os.environ.pop("JWT_SECRET", None) try: from src.utils.secrets import SecretsVault + vault = SecretsVault(secrets_file=str(secrets_file)) assert vault.get("JWT_SECRET") == "from-json-secret" finally: @@ -53,15 +56,20 @@ def test_json_value_used_when_env_absent(self, tmp_path): def test_new_key_in_json_overridden_by_env(self, tmp_path): """Any key found in JSON (not just the 4 hardcoded ones) must be env-overrideable.""" secrets_file = tmp_path / "secrets.json" - secrets_file.write_text(json.dumps({ - "JWT_SECRET": "json-jwt", - "SMTP_PASSWORD": "json-smtp-pass", - })) + secrets_file.write_text( + json.dumps( + { + "JWT_SECRET": "json-jwt", + "SMTP_PASSWORD": "json-smtp-pass", + } + ) + ) os.environ["JWT_SECRET"] = "env-jwt" os.environ["SMTP_PASSWORD"] = "env-smtp-pass" try: from src.utils.secrets import SecretsVault + vault = SecretsVault(secrets_file=str(secrets_file)) assert vault.get("JWT_SECRET") == "env-jwt" assert vault.get("SMTP_PASSWORD") == "env-smtp-pass" @@ -74,6 +82,7 @@ def test_missing_secrets_json_does_not_raise(self, tmp_path): os.environ["JWT_SECRET"] = "env-only-value" try: from src.utils.secrets import SecretsVault + vault = SecretsVault(secrets_file=str(tmp_path / "nonexistent.json")) assert vault.get("JWT_SECRET") == "env-only-value" finally: @@ -88,6 +97,7 @@ def test_app_secret_key_alias_for_jwt_secret(self, tmp_path): os.environ["APP_SECRET_KEY"] = "alias-secret" try: from src.utils.secrets import SecretsVault + vault = SecretsVault(secrets_file=str(secrets_file)) assert vault.get("JWT_SECRET") == "alias-secret" finally: @@ -99,6 +109,7 @@ def test_app_secret_key_alias_for_jwt_secret(self, tmp_path): # SecretsVault — production startup guard # --------------------------------------------------------------------------- + class TestSecretsVaultProductionGuard: """In ENV=prod, startup must be blocked for missing or placeholder JWT_SECRET.""" @@ -107,11 +118,13 @@ class TestSecretsVaultProductionGuard: def _make_vault(self, jwt_secret, tmp_path, env="prod"): """Helper: create a SecretsVault directly with controlled env state.""" from src.utils.secrets import SecretsVault + secrets_file = tmp_path / "secrets.json" secrets_file.write_text(json.dumps({})) # Patch the module-level _ENV used inside _validate import src.utils.secrets as sm + original_env = sm._ENV sm._ENV = env if jwt_secret: @@ -145,6 +158,7 @@ def test_prod_succeeds_with_strong_secret(self, tmp_path): # .gitignore coverage # --------------------------------------------------------------------------- + class TestGitignoreCoverage: """Critical secret file patterns must be listed in .gitignore.""" @@ -175,6 +189,7 @@ def test_env_production_is_gitignored(self): # Kafka consumer reads from env # --------------------------------------------------------------------------- + class TestKafkaConsumerEnvConfig: """EmailStreamConsumer must read broker URL and topic from environment.""" @@ -182,6 +197,7 @@ def _build_consumer(self, broker_url=None, topic=None, env_vars=None): """Instantiate a minimal consumer-like object using only the env logic, without importing the full consumer module (avoids transformers dep).""" import os + _DEFAULT_BROKER = "localhost:9092" _DEFAULT_TOPIC = "email_ingest" @@ -193,7 +209,9 @@ def _build_consumer(self, broker_url=None, topic=None, env_vars=None): else: os.environ[key] = val try: - resolved_broker = broker_url or os.environ.get("KAFKA_BROKER_URL", _DEFAULT_BROKER) + resolved_broker = broker_url or os.environ.get( + "KAFKA_BROKER_URL", _DEFAULT_BROKER + ) resolved_topic = topic or os.environ.get("KAFKA_TOPIC", _DEFAULT_TOPIC) finally: for key, original in saved.items(): @@ -205,7 +223,10 @@ def _build_consumer(self, broker_url=None, topic=None, env_vars=None): def test_broker_url_from_env(self): broker, topic = self._build_consumer( - env_vars={"KAFKA_BROKER_URL": "kafka.prod.example.com:9093", "KAFKA_TOPIC": "prod_email_ingest"} + env_vars={ + "KAFKA_BROKER_URL": "kafka.prod.example.com:9093", + "KAFKA_TOPIC": "prod_email_ingest", + } ) assert broker == "kafka.prod.example.com:9093" assert topic == "prod_email_ingest" @@ -221,7 +242,7 @@ def test_explicit_args_override_env(self): broker, topic = self._build_consumer( broker_url="explicit:9093", topic="explicit_topic", - env_vars={"KAFKA_BROKER_URL": "from-env:9092"} + env_vars={"KAFKA_BROKER_URL": "from-env:9092"}, ) assert broker == "explicit:9093" assert topic == "explicit_topic" diff --git a/tests/test_security_features.py b/tests/test_security_features.py index 956e17a..94b8a2f 100644 --- a/tests/test_security_features.py +++ b/tests/test_security_features.py @@ -1,35 +1,45 @@ import pytest + def test_url_extraction(): from src.security.url_extractor import extract_urls + text = "Check this out: http://evil.com and https://192.168.1.1/login" urls = extract_urls(text) assert len(urls) == 2 assert set(urls) == {"http://evil.com", "https://192.168.1.1/login"} + def test_homograph_detection(): from src.security.homograph_detector import detect_homograph + # "apple" using Cyrillic 'а' text = "Visit https://аpple.com" result = detect_homograph(text) assert result.get("homograph_detected", False) is True + def test_script_detection(): from src.security.script_detector import detect_script_types - text = "Paypаl login" # mixed Latin and Cyrillic + + text = "Paypаl login" # mixed Latin and Cyrillic result = detect_script_types(text) assert result.get("mixed_script", False) is True + def test_risk_scoring(): from src.security.risk_scoring import calculate_security_risk + text = "Urgent: Verify your account at http://192.168.1.55" result = calculate_security_risk(text) assert result.get("risk_score", 0) > 30 flags = result.get("security_flags", {}) assert flags.get("ip_url", False) is True + def test_cyrillic_url_detection(): from src.security.risk_scoring import calculate_security_risk + text = "Visit https://вася.com" result = calculate_security_risk(text) assert result.get("risk_score", 0) >= 50 diff --git a/tests/test_vectorizer.py b/tests/test_vectorizer.py index 3c26388..f2093bd 100644 --- a/tests/test_vectorizer.py +++ b/tests/test_vectorizer.py @@ -2,6 +2,7 @@ import sys import os + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import pytest From dc38110ae41ff598a3ffa3db440cc62faeb2fb1e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:17:59 +0000 Subject: [PATCH 3/3] fix: harden URL extraction regex usage against ReDoS Co-authored-by: VIPHACKER100 <110089810+VIPHACKER100@users.noreply.github.com> --- src/security/url_extractor.py | 56 ++++++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/src/security/url_extractor.py b/src/security/url_extractor.py index 65959fd..1e78959 100644 --- a/src/security/url_extractor.py +++ b/src/security/url_extractor.py @@ -5,17 +5,59 @@ import re from typing import List -# Improved URL detection regex -URL_PATTERN = re.compile( - r"(?:http|https)://[^\s/$.?#].[^\s]*" # http/https - r"|www\.[^\s/$.?#].[^\s]*" # www. - r"|[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|io|biz|info|xyz|top|click|gq|tk|ml|ga|cf)\b" # domain-like -) +# Improved URL detection regex (limited to protocol/www URLs to avoid ReDoS risk) +URL_PATTERN = re.compile(r"(?:https?://|www\.)[^\s<>'\"()]+") + +_MAX_REGEX_INPUT_LENGTH = 8192 +_DOMAIN_TLDS = { + "com", + "net", + "org", + "edu", + "gov", + "io", + "biz", + "info", + "xyz", + "top", + "click", + "gq", + "tk", + "ml", + "ga", + "cf", +} + + +def _extract_domain_like_tokens(text: str) -> List[str]: + domain_like: List[str] = [] + for token in text.split(): + clean = token.strip(".,;:!?()[]{}<>\"'") + lower = clean.lower() + if not lower or lower.startswith(("http://", "https://", "www.")): + continue + if "." not in lower: + continue + parts = lower.split(".") + if len(parts) < 2 or any(not part for part in parts): + continue + if any( + not all(ch.isalnum() or ch == "-" for ch in part) for part in parts[:-1] + ): + continue + if parts[-1] in _DOMAIN_TLDS: + domain_like.append(clean) + return domain_like def extract_urls(text: str) -> List[str]: """Extract all URLs and domain-like strings from text.""" - return URL_PATTERN.findall(text) + if not text: + return [] + bounded_text = text[:_MAX_REGEX_INPUT_LENGTH] + matches = URL_PATTERN.findall(bounded_text) + matches.extend(_extract_domain_like_tokens(bounded_text)) + return list(dict.fromkeys(matches)) def get_domain(url: str) -> str: