diff --git a/src/database.py b/src/database.py index b3217c3..c59245d 100644 --- a/src/database.py +++ b/src/database.py @@ -261,7 +261,7 @@ class DatabaseManager: session.add(detection) # Update IP stats - self._update_ip_stats(session, ip) + self._update_ip_stats(session, ip, is_suspicious) session.commit() return access_log.id @@ -313,13 +313,16 @@ class DatabaseManager: finally: self.close_session() - def _update_ip_stats(self, session: Session, ip: str) -> None: + def _update_ip_stats( + self, session: Session, ip: str, is_suspicious: bool = False + ) -> None: """ Update IP statistics (upsert pattern). Args: session: Active database session ip: IP address to update + is_suspicious: Whether the request was flagged as suspicious """ sanitized_ip = sanitize_ip(ip) now = datetime.now() @@ -329,9 +332,15 @@ class DatabaseManager: if ip_stats: ip_stats.total_requests += 1 ip_stats.last_seen = now + if is_suspicious: + ip_stats.need_reevaluation = True else: ip_stats = IpStats( - ip=sanitized_ip, total_requests=1, first_seen=now, last_seen=now + ip=sanitized_ip, + total_requests=1, + first_seen=now, + last_seen=now, + need_reevaluation=is_suspicious, ) session.add(ip_stats) @@ -385,6 +394,7 @@ class DatabaseManager: ip_stats.category = category ip_stats.category_scores = category_scores ip_stats.last_analysis = last_analysis + ip_stats.need_reevaluation = False try: session.commit() @@ -637,6 +647,24 @@ class DatabaseManager: finally: self.close_session() + def get_ips_needing_reevaluation(self) -> List[str]: + """ + Get all IP addresses that have been flagged for reevaluation. + + Returns: + List of IP addresses where need_reevaluation is True + """ + session = self.session + try: + ips = ( + session.query(IpStats.ip) + .filter(IpStats.need_reevaluation == True) + .all() + ) + return [ip[0] for ip in ips] + finally: + self.close_session() + def get_access_logs( self, limit: int = 100, diff --git a/src/migrations/runner.py b/src/migrations/runner.py index a46a07a..396fbc9 100644 --- a/src/migrations/runner.py +++ b/src/migrations/runner.py @@ -38,6 +38,16 @@ def _migrate_raw_request_column(cursor) -> bool: return True +def _migrate_need_reevaluation_column(cursor) -> bool: + """Add need_reevaluation column to ip_stats if missing.""" + if _column_exists(cursor, "ip_stats", "need_reevaluation"): + return False + cursor.execute( + "ALTER TABLE ip_stats ADD COLUMN need_reevaluation BOOLEAN DEFAULT 0" + ) + return True + + def _migrate_performance_indexes(cursor) -> List[str]: """Add performance indexes to attack_detections if missing.""" added = [] @@ -77,6 +87,9 @@ def run_migrations(database_path: str) -> None: if _migrate_raw_request_column(cursor): applied.append("add raw_request column to access_logs") + if _migrate_need_reevaluation_column(cursor): + applied.append("add need_reevaluation column to ip_stats") + idx_added = _migrate_performance_indexes(cursor) for idx in idx_added: applied.append(f"add index {idx}") diff --git a/src/models.py b/src/models.py index a38b1f6..c9a190f 100644 --- a/src/models.py +++ b/src/models.py @@ -200,6 +200,9 @@ class IpStats(Base): category_scores: Mapped[Dict[str, int]] = mapped_column(JSON, nullable=True) manual_category: Mapped[bool] = mapped_column(Boolean, default=False, nullable=True) last_analysis: Mapped[datetime] = mapped_column(DateTime, nullable=True) + need_reevaluation: Mapped[bool] = mapped_column( + Boolean, default=False, nullable=True + ) def __repr__(self) -> str: return f"" diff --git a/src/tasks/analyze_ips.py b/src/tasks/analyze_ips.py index f5fea5b..81ad1d9 100644 --- a/src/tasks/analyze_ips.py +++ b/src/tasks/analyze_ips.py @@ -94,12 +94,11 @@ def main(): "attack_url": 0, }, } - # Get IPs with recent activity (last minute to match cron schedule) - recent_accesses = db_manager.get_access_logs(limit=999999999, since_minutes=1) - ips_to_analyze = {item["ip"] for item in recent_accesses} + # Get IPs flagged for reevaluation (set when a suspicious request arrives) + ips_to_analyze = set(db_manager.get_ips_needing_reevaluation()) if not ips_to_analyze: - app_logger.debug("[Background Task] analyze-ips: No recent activity, skipping") + app_logger.debug("[Background Task] analyze-ips: No IPs need reevaluation, skipping") return for ip in ips_to_analyze: