diff --git a/requirements.txt b/requirements.txt index 5849375..30732cf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,4 +22,5 @@ jinja2>=3.1.5 # Scalable mode backends (PostgreSQL + Redis) psycopg2-binary>=2.9.0 -redis>=5.0.0 \ No newline at end of file +redis>=5.0.0 +prometheus-client>=0.25.0 diff --git a/src/app.py b/src/app.py index 604784a..baf156b 100644 --- a/src/app.py +++ b/src/app.py @@ -6,10 +6,11 @@ """ import gc -import sys import os from contextlib import asynccontextmanager +from prometheus_client import make_asgi_app + from fastapi import FastAPI, Request, Response from fastapi.responses import JSONResponse, FileResponse from fastapi.staticfiles import StaticFiles @@ -247,6 +248,15 @@ async def access_log_middleware(request: Request, call_next): config = get_config() secret = config.dashboard_secret_path.lstrip("/") static_dir = os.path.join(os.path.dirname(__file__), "templates", "static") + + # Add prometheus ASGI application endpoint + metrics_app = make_asgi_app() + application.mount( + f"/{secret}/metrics", + metrics_app, + name="metrics", + ) + application.mount( f"/{secret}/static", StaticFiles(directory=static_dir), diff --git a/src/database.py b/src/database.py index 95b2bc6..e510c13 100644 --- a/src/database.py +++ b/src/database.py @@ -9,13 +9,10 @@ import stat from datetime import datetime, timedelta from typing import Optional, List, Dict, Any -from zoneinfo import ZoneInfo from sqlalchemy import create_engine, func, distinct, event, or_, and_ from sqlalchemy.orm import sessionmaker, scoped_session, Session, joinedload -from sqlalchemy.engine import Engine -from ip_utils import is_local_or_private_ip, is_valid_public_ip from models import ( Base, @@ -25,7 +22,6 @@ IpStats, CategoryHistory, TrackedIp, - GeneratedPage, ) from sanitizer import ( sanitize_ip, @@ -2461,7 +2457,6 @@ def get_attack_types_stats( """ session = self.session try: - from sqlalchemy import func # Aggregate attack types with count query = session.query( @@ -3110,6 +3105,18 @@ def get_generated_pages_paginated( finally: self.close_session() + def count_category(self, category: str) -> int: + """Count the total number of ips in a given category.""" + session = self.session + try: + count = session.query(IpStats).filter(IpStats.category == category).count() + return count or 0 + except Exception as e: + applogger.error(f"Error counting {category}: {e}") + return 0 + finally: + self.close_session() + def count_generated_pages_created_today(self) -> int: """Count how many generated pages were created today. diff --git a/src/routes/api.py b/src/routes/api.py index d2512a5..5c30906 100644 --- a/src/routes/api.py +++ b/src/routes/api.py @@ -11,7 +11,7 @@ import secrets import time -from fastapi import APIRouter, Request, Response, Query, Cookie +from fastapi import APIRouter, Request, Response, Query from fastapi.responses import JSONResponse from pydantic import BaseModel diff --git a/src/tasks/analyze_ips.py b/src/tasks/analyze_ips.py index 4444244..fb9f297 100644 --- a/src/tasks/analyze_ips.py +++ b/src/tasks/analyze_ips.py @@ -7,6 +7,7 @@ from wordlists import get_wordlists from config import get_config from logger import get_app_logger +from prometheus_client import Gauge, REGISTRY # ---------------------- # TASK CONFIG @@ -19,12 +20,25 @@ "run_when_loaded": True, } +CATEGORIES = ("attacker", "good_crawler", "bad_crawler", "regular_user") + +num_clients = Gauge( + "clients_total", + "Total number of IPs per classification category", + labelnames=["category"], + namespace="krawl", + registry=REGISTRY, +) + def main(): config = get_config() db_manager = get_database() app_logger = get_app_logger() + for category in CATEGORIES: + num_clients.labels(category).set(db_manager.count_category(category)) + http_risky_methods_threshold = config.http_risky_methods_threshold violated_robots_threshold = config.violated_robots_threshold uneven_request_timing_threshold = config.uneven_request_timing_threshold @@ -382,9 +396,14 @@ def main(): "bad_crawler": bad_crawler_score, "regular_user": regular_user_score, } + category = max(category_scores, key=category_scores.get) + last_analysis = datetime.now() db_manager.update_ip_stats_analysis( ip, analyzed_metrics, category, category_scores, last_analysis ) + + for category in CATEGORIES: + num_clients.labels(category).set(db_manager.count_category(category)) return