123 lines
5.1 KiB
Python
123 lines
5.1 KiB
Python
"""Per-batch classification pipeline (Chapter 07): ties bot_identifier and
|
|
threat_scanner into the bot_hits/suspicious_events/ip_registry write path.
|
|
Called once per bulk-insert chunk from app/cli.py — DNS-based verification
|
|
belongs here (background batch job), never in a dashboard request.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import func, insert
|
|
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
|
|
|
from app.extensions import db
|
|
from app.models.bot_hit import BotHit
|
|
from app.models.ip_registry import IPRegistry
|
|
from app.models.suspicious_event import SuspiciousEvent
|
|
from app.services import bot_identifier, threat_scanner
|
|
from app.services.log_parser import ParsedEntry
|
|
|
|
|
|
@dataclass
|
|
class EntryClassification:
|
|
"""Cheap, no-I/O flags for the log_entries row itself."""
|
|
is_bot: bool
|
|
flagged: bool
|
|
|
|
|
|
def classify_entry(entry: ParsedEntry) -> EntryClassification:
|
|
"""No DNS I/O here — bot *verification* is batched separately below,
|
|
since it's stateful (cached per IP) and only worth doing once per IP
|
|
per batch, not once per line.
|
|
"""
|
|
return EntryClassification(
|
|
is_bot=bot_identifier.classify_bot(entry.user_agent) is not None,
|
|
flagged=threat_scanner.scan(entry) is not None,
|
|
)
|
|
|
|
|
|
def write_batch_side_effects(log_file_id: int, entries: list[ParsedEntry]) -> None:
|
|
"""Derive and bulk-write bot_hits, suspicious_events, ip_registry upserts."""
|
|
if not entries:
|
|
return
|
|
|
|
bot_hit_rows: list[dict] = []
|
|
suspicious_rows: list[dict] = []
|
|
ip_agg: dict[str, dict] = {}
|
|
verified_this_batch: dict[str, bool] = {} # avoid repeat DNS for the same IP in one batch
|
|
|
|
for entry in entries:
|
|
agg = ip_agg.setdefault(entry.ip, {"first": entry.timestamp, "last": entry.timestamp, "count": 0})
|
|
agg["count"] += 1
|
|
agg["first"] = min(agg["first"], entry.timestamp)
|
|
agg["last"] = max(agg["last"], entry.timestamp)
|
|
|
|
bot_name = bot_identifier.classify_bot(entry.user_agent)
|
|
if bot_name is not None:
|
|
verified = verified_this_batch.get(entry.ip)
|
|
if verified is None:
|
|
verified = _verify_with_cache(entry.ip, bot_name)
|
|
verified_this_batch[entry.ip] = verified
|
|
bot_hit_rows.append({
|
|
"log_file_id": log_file_id, "timestamp": entry.timestamp, "ip": entry.ip,
|
|
"bot_name": bot_name, "verified": verified, "path": entry.path,
|
|
"status_code": entry.status_code,
|
|
})
|
|
if not verified:
|
|
# Single detection, two dashboard consumers (Ch07): spoofed
|
|
# bot also surfaces as a suspicious_events row.
|
|
suspicious_rows.append({
|
|
"log_file_id": log_file_id, "ip": entry.ip, "timestamp": entry.timestamp,
|
|
"path": entry.path, "rule_matched": f"spoofed_bot:{bot_name}", "severity": "medium",
|
|
})
|
|
|
|
threat = threat_scanner.scan(entry)
|
|
if threat is not None:
|
|
suspicious_rows.append({
|
|
"log_file_id": log_file_id, "ip": entry.ip, "timestamp": entry.timestamp,
|
|
"path": entry.path, "rule_matched": threat.rule_matched, "severity": threat.severity,
|
|
})
|
|
|
|
if bot_hit_rows:
|
|
db.session.execute(insert(BotHit.__table__), bot_hit_rows)
|
|
if suspicious_rows:
|
|
db.session.execute(insert(SuspiciousEvent.__table__), suspicious_rows)
|
|
|
|
_upsert_ip_registry(ip_agg, verified_this_batch)
|
|
db.session.commit()
|
|
|
|
|
|
def _verify_with_cache(ip: str, bot_name: str) -> bool:
|
|
"""TTL-gated reverse/forward DNS check, cached via
|
|
ip_registry.last_verified_bot_result (Ch07 schema addition).
|
|
"""
|
|
row = db.session.get(IPRegistry, ip)
|
|
if row is not None and not bot_identifier.is_verification_stale(row.last_verified_at):
|
|
return bool(row.last_verified_bot_result)
|
|
return bot_identifier.verify_bot_ip(ip, bot_name)
|
|
|
|
|
|
def _upsert_ip_registry(ip_agg: dict[str, dict], verified_this_batch: dict[str, bool]) -> None:
|
|
now = datetime.utcnow()
|
|
for ip, agg in ip_agg.items():
|
|
values = {
|
|
"ip": ip, "first_seen": agg["first"], "last_seen": agg["last"],
|
|
"total_requests": agg["count"], "reputation_score": 0, "is_flagged": False,
|
|
}
|
|
if ip in verified_this_batch:
|
|
values["last_verified_at"] = now
|
|
values["last_verified_bot_result"] = verified_this_batch[ip]
|
|
|
|
stmt = sqlite_insert(IPRegistry.__table__).values(**values)
|
|
update_set = {
|
|
"last_seen": func.max(IPRegistry.last_seen, stmt.excluded.last_seen),
|
|
"first_seen": func.min(IPRegistry.first_seen, stmt.excluded.first_seen),
|
|
"total_requests": IPRegistry.total_requests + stmt.excluded.total_requests,
|
|
}
|
|
if ip in verified_this_batch:
|
|
update_set["last_verified_at"] = stmt.excluded.last_verified_at
|
|
update_set["last_verified_bot_result"] = stmt.excluded.last_verified_bot_result
|
|
stmt = stmt.on_conflict_do_update(index_elements=["ip"], set_=update_set)
|
|
db.session.execute(stmt)
|