Initial ThreatLens release
This commit is contained in:
Executable
+289
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
utils/helpers.py — small, dependency-light utilities.
|
||||
|
||||
Contains:
|
||||
* Severity helpers (label → colour, ordering).
|
||||
* CVE-ID validation/normalisation.
|
||||
* Datetime parsing/formatting for NVD timestamps.
|
||||
* A tiny on-disk JSON cache with TTL for rate-limit friendliness.
|
||||
* The async SQLite store for bot config + user notes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from config import settings
|
||||
|
||||
log = logging.getLogger("vulnforge.helpers")
|
||||
|
||||
# CVE IDs look like: CVE-2024-12345 (4-digit year, 4+ digit sequence).
|
||||
_CVE_RE = re.compile(r"^CVE-\d{4}-\d{4,}$", re.IGNORECASE)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
# Severity helpers
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
SEVERITY_ORDER = {"CRITICAL": 4, "HIGH": 3, "MEDIUM": 2, "LOW": 1, "NONE": 0, "UNKNOWN": 0}
|
||||
|
||||
|
||||
def severity_from_score(score: float | None) -> str:
|
||||
"""Map a CVSS base score (0-10) to a qualitative severity label."""
|
||||
if score is None:
|
||||
return "UNKNOWN"
|
||||
if score >= 9.0:
|
||||
return "CRITICAL"
|
||||
if score >= 7.0:
|
||||
return "HIGH"
|
||||
if score >= 4.0:
|
||||
return "MEDIUM"
|
||||
if score > 0.0:
|
||||
return "LOW"
|
||||
return "NONE"
|
||||
|
||||
|
||||
def severity_color(severity: str) -> int:
|
||||
"""Return the cyberpunk palette colour int for a severity label."""
|
||||
mapping = {
|
||||
"CRITICAL": settings.color_critical,
|
||||
"HIGH": settings.color_high,
|
||||
"MEDIUM": settings.color_medium,
|
||||
"LOW": settings.color_low,
|
||||
"NONE": settings.color_info,
|
||||
"UNKNOWN": settings.color_dark,
|
||||
}
|
||||
return mapping.get(severity.upper(), settings.color_info)
|
||||
|
||||
|
||||
def severity_emoji(severity: str) -> str:
|
||||
"""A glanceable emoji marker per severity tier."""
|
||||
return {
|
||||
"CRITICAL": "🟥",
|
||||
"HIGH": "🟧",
|
||||
"MEDIUM": "🟨",
|
||||
"LOW": "🟦",
|
||||
"NONE": "⬜",
|
||||
"UNKNOWN": "⬛",
|
||||
}.get(severity.upper(), "⬛")
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
# CVE-ID + datetime helpers
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
def is_valid_cve_id(cve_id: str) -> bool:
|
||||
"""True when the string matches the canonical CVE-YYYY-NNNN format."""
|
||||
return bool(_CVE_RE.match(cve_id.strip()))
|
||||
|
||||
|
||||
def normalize_cve_id(cve_id: str) -> str:
|
||||
"""Upper-case + strip a CVE id for consistent lookups/keys."""
|
||||
return cve_id.strip().upper()
|
||||
|
||||
|
||||
def parse_nvd_datetime(raw: str | None) -> datetime | None:
|
||||
"""Parse NVD ISO timestamps (e.g. '2024-05-01T12:00:00.000') as UTC."""
|
||||
if not raw:
|
||||
return None
|
||||
cleaned = raw.replace("Z", "+00:00")
|
||||
for fmt in ("%Y-%m-%dT%H:%M:%S.%f%z", "%Y-%m-%dT%H:%M:%S%z"):
|
||||
try:
|
||||
return datetime.strptime(cleaned, fmt)
|
||||
except ValueError:
|
||||
continue
|
||||
try:
|
||||
dt = datetime.fromisoformat(cleaned)
|
||||
return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
|
||||
except ValueError:
|
||||
log.debug("Could not parse NVD datetime: %s", raw)
|
||||
return None
|
||||
|
||||
|
||||
def discord_timestamp(dt: datetime | None, style: str = "R") -> str:
|
||||
"""Render a Discord dynamic timestamp tag, or 'N/A' if dt is None."""
|
||||
if dt is None:
|
||||
return "N/A"
|
||||
return f"<t:{int(dt.timestamp())}:{style}>"
|
||||
|
||||
|
||||
def truncate(text: str, limit: int) -> str:
|
||||
"""Trim text to `limit` chars (Discord field/desc safe), adding an ellipsis."""
|
||||
if text is None:
|
||||
return ""
|
||||
return text if len(text) <= limit else text[: limit - 1].rstrip() + "…"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
# JSON file cache with TTL
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
class JSONCache:
|
||||
"""A minimal on-disk JSON cache keyed by filename, with TTL expiry."""
|
||||
|
||||
def __init__(self, cache_dir: Path) -> None:
|
||||
self._dir = cache_dir
|
||||
self._dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _path(self, key: str) -> Path:
|
||||
safe = re.sub(r"[^A-Za-z0-9_.-]", "_", key)
|
||||
return self._dir / f"{safe}.json"
|
||||
|
||||
def get(self, key: str, ttl: int) -> Any | None:
|
||||
"""Return cached payload if present and younger than ttl seconds."""
|
||||
path = self._path(key)
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
if time.time() - path.stat().st_mtime > ttl:
|
||||
return None
|
||||
with path.open("r", encoding="utf-8") as fh:
|
||||
return json.load(fh)
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
log.warning("Cache read failed for %s: %s", key, exc)
|
||||
return None
|
||||
|
||||
def set(self, key: str, payload: Any) -> None:
|
||||
"""Persist payload to disk (best-effort; failures are logged only)."""
|
||||
path = self._path(key)
|
||||
try:
|
||||
with path.open("w", encoding="utf-8") as fh:
|
||||
json.dump(payload, fh)
|
||||
except OSError as exc:
|
||||
log.warning("Cache write failed for %s: %s", key, exc)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
# SQLite store (config + user notes)
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
class Store:
|
||||
"""
|
||||
Async SQLite wrapper.
|
||||
|
||||
Tables:
|
||||
guild_config(guild_id PK, digest_channel_id, notify_role_id, updated_at)
|
||||
cve_notes(id PK, guild_id, user_id, cve_id, note, created_at)
|
||||
kev_seen(cve_id PK, date_added) -- tracks KEV entries we've recorded
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self._db_path = db_path
|
||||
|
||||
async def init(self) -> None:
|
||||
"""Create tables if they do not exist."""
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS guild_config (
|
||||
guild_id INTEGER PRIMARY KEY,
|
||||
digest_channel_id INTEGER,
|
||||
notify_role_id INTEGER,
|
||||
updated_at TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
await db.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS cve_notes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
guild_id INTEGER,
|
||||
user_id INTEGER NOT NULL,
|
||||
cve_id TEXT NOT NULL,
|
||||
note TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
await db.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS kev_seen (
|
||||
cve_id TEXT PRIMARY KEY,
|
||||
date_added TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
await db.commit()
|
||||
log.info("SQLite store initialised at %s", self._db_path)
|
||||
|
||||
# ── guild config ─────────────────────────────────────────
|
||||
async def set_guild_config(
|
||||
self,
|
||||
guild_id: int,
|
||||
digest_channel_id: int | None,
|
||||
notify_role_id: int | None,
|
||||
) -> None:
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
INSERT INTO guild_config (guild_id, digest_channel_id, notify_role_id, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(guild_id) DO UPDATE SET
|
||||
digest_channel_id = excluded.digest_channel_id,
|
||||
notify_role_id = excluded.notify_role_id,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(
|
||||
guild_id,
|
||||
digest_channel_id,
|
||||
notify_role_id,
|
||||
datetime.now(timezone.utc).isoformat(),
|
||||
),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get_guild_config(self, guild_id: int) -> dict[str, Any] | None:
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
async with db.execute(
|
||||
"SELECT * FROM guild_config WHERE guild_id = ?", (guild_id,)
|
||||
) as cur:
|
||||
row = await cur.fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
async def all_guild_configs(self) -> list[dict[str, Any]]:
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
async with db.execute(
|
||||
"SELECT * FROM guild_config WHERE digest_channel_id IS NOT NULL"
|
||||
) as cur:
|
||||
return [dict(r) for r in await cur.fetchall()]
|
||||
|
||||
# ── user notes ───────────────────────────────────────────
|
||||
async def add_note(
|
||||
self, guild_id: int | None, user_id: int, cve_id: str, note: str
|
||||
) -> None:
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
INSERT INTO cve_notes (guild_id, user_id, cve_id, note, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
guild_id,
|
||||
user_id,
|
||||
normalize_cve_id(cve_id),
|
||||
note,
|
||||
datetime.now(timezone.utc).isoformat(),
|
||||
),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get_notes(self, cve_id: str, limit: int = 5) -> list[dict[str, Any]]:
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
async with db.execute(
|
||||
"""
|
||||
SELECT * FROM cve_notes
|
||||
WHERE cve_id = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(normalize_cve_id(cve_id), limit),
|
||||
) as cur:
|
||||
return [dict(r) for r in await cur.fetchall()]
|
||||
Reference in New Issue
Block a user