From d96e326d4d3c714392e4c638428ac816f2c121b0 Mon Sep 17 00:00:00 2001 From: subinacls Date: Sat, 27 Jun 2026 21:24:57 -0400 Subject: [PATCH] Initial ThreatLens release --- .dockerignore | 29 ++ .env.example | 24 ++ .gitignore | 10 + Dockerfile | 36 ++ README.md | 149 ++++++++ check_logs.sh | 2 + cogs/__init__.py | 1 + cogs/cve.py | 261 ++++++++++++++ cogs/enrich.py | 314 +++++++++++++++++ cogs/weekly.py | 159 +++++++++ config.py | 85 +++++ deep_inspect.sh | 33 ++ deploy_check.sh | 10 + diag.sh | 21 ++ docker-compose.yml | 24 ++ inspect_cmds.sh | 26 ++ main.py | 160 +++++++++ requirements.txt | 7 + scripts/__init__.py | 1 + scripts/enrich_report.py | 49 +++ utils/__init__.py | 1 + utils/api.py | 329 ++++++++++++++++++ utils/embeds.py | 293 ++++++++++++++++ utils/enricher.py | 728 +++++++++++++++++++++++++++++++++++++++ utils/helpers.py | 289 ++++++++++++++++ verify.sh | 6 + verify2.sh | 12 + 27 files changed, 3059 insertions(+) create mode 100755 .dockerignore create mode 100755 .env.example create mode 100755 .gitignore create mode 100755 Dockerfile create mode 100755 README.md create mode 100755 check_logs.sh create mode 100755 cogs/__init__.py create mode 100755 cogs/cve.py create mode 100755 cogs/enrich.py create mode 100755 cogs/weekly.py create mode 100755 config.py create mode 100755 deep_inspect.sh create mode 100755 deploy_check.sh create mode 100755 diag.sh create mode 100755 docker-compose.yml create mode 100755 inspect_cmds.sh create mode 100755 main.py create mode 100755 requirements.txt create mode 100755 scripts/__init__.py create mode 100755 scripts/enrich_report.py create mode 100755 utils/__init__.py create mode 100755 utils/api.py create mode 100755 utils/embeds.py create mode 100755 utils/enricher.py create mode 100755 utils/helpers.py create mode 100755 verify.sh create mode 100755 verify2.sh diff --git a/.dockerignore b/.dockerignore new file mode 100755 index 0000000..1d71d7a --- /dev/null +++ b/.dockerignore @@ -0,0 +1,29 @@ +# ── Secrets & local data (never bake into the image) ── +.env +data/ +*.db +*.log + +# ── Python ── +__pycache__/ +*.py[cod] +*$py.class +.venv/ +venv/ +env/ +*.egg-info/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ + +# ── Git / CI / docs ── +.git/ +.gitignore +.github/ +README.md + +# ── OS / editor ── +.DS_Store +Thumbs.db +.idea/ +.vscode/ diff --git a/.env.example b/.env.example new file mode 100755 index 0000000..aa9060c --- /dev/null +++ b/.env.example @@ -0,0 +1,24 @@ +# ────────────────────────────────────────────────────────────── +# VulnForge — Environment variables +# Copy this file to ".env" and fill in your own values. +# ────────────────────────────────────────────────────────────── + +# REQUIRED: Your Discord bot token (Discord Developer Portal → Bot → Token) +DISCORD_TOKEN=your-discord-bot-token-here + +# OPTIONAL: Restrict slash-command sync to a single guild for instant updates. +# Leave blank to sync globally (can take up to 1 hour to propagate). +GUILD_ID= + +# OPTIONAL: NVD API key. Greatly increases rate limits (50 req / 30s vs 5). +# Request one for free: https://nvd.nist.gov/developers/request-an-api-key +NVD_API_KEY= + +# OPTIONAL: Default channel ID for the weekly digest. Can also be set via /setup. +DIGEST_CHANNEL_ID= + +# OPTIONAL: Role ID to ping on the weekly digest. Can also be set via /setup. +NOTIFY_ROLE_ID= + +# OPTIONAL: Logging verbosity (DEBUG, INFO, WARNING, ERROR) +LOG_LEVEL=INFO diff --git a/.gitignore b/.gitignore new file mode 100755 index 0000000..07f2c6f --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +.env +data/ +*.db +*.sqlite +*.sqlite3 +__pycache__/ +*.pyc +*.pyo +.venv/ +venv/ diff --git a/Dockerfile b/Dockerfile new file mode 100755 index 0000000..7fdb278 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,36 @@ +# ────────────────────────────────────────────────────────────── +# VulnForge — production Docker image +# ────────────────────────────────────────────────────────────── +FROM python:3.12-slim AS base + +# Faster, cleaner Python in containers. +ENV PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + TZ=UTC + +WORKDIR /app + +# Install dependencies first to leverage Docker layer caching. +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Copy the rest of the application source. +COPY . . + +# Run as a non-root user for security; give it ownership of /app/data +# so the SQLite DB + cache are writable (also mounted as a volume). +RUN useradd --create-home --uid 10001 vulnforge \ + && mkdir -p /app/data/cache \ + && chown -R vulnforge:vulnforge /app +USER vulnforge + +# Persist SQLite DB + on-disk cache across restarts. +VOLUME ["/app/data"] + +# Lightweight liveness check: confirms the Python runtime + deps import OK. +HEALTHCHECK --interval=5m --timeout=10s --start-period=30s --retries=3 \ + CMD python -c "import nextcord, aiohttp, aiosqlite" || exit 1 + +CMD ["python", "main.py"] diff --git a/README.md b/README.md new file mode 100755 index 0000000..96b557e --- /dev/null +++ b/README.md @@ -0,0 +1,149 @@ +# ☣ VulnForge + +> **Security Intelligence Discord Bot** — weekly CVE/KEV threat digests and on-demand vulnerability lookups, wrapped in a dark cyberpunk aesthetic. + +VulnForge fuses two authoritative sources — the **NVD CVE API v2** and the **CISA Known Exploited Vulnerabilities (KEV)** catalog — into a single Discord experience built for security researchers and blue/red teams. + +--- + +## ⚡ Features + +### 1. Automated Weekly Digest +Every **Monday @ 09:00 UTC**, VulnForge posts a rich digest to your configured channel: +- 🆕 Total new CVEs in the last 7 days +- ☣ How many are already in CISA KEV / newly added to KEV +- ▶ Top **8** most critical CVEs (sorted by CVSS, color-coded by severity) +- 🔗 Direct links to NVD and CISA + +### 2. Slash Commands +| Command | Description | +|---------|-------------| +| `/cve ` | Opens a **modal** to pick output sections (details, KEV, refs, CWE) + add a researcher note, then replies with a rich embed. | +| `/kev [limit]` | Latest CISA KEV entries (actively exploited in the wild). | +| `/weekly` | **Admin-only** — force-post the weekly digest now. | +| `/setup` | **Admin-only** — set the digest channel + notification role. | + +### 3. Under the hood +- **Nextcord** (slash commands, modals, views) +- **aiohttp** with retry/backoff + on-disk **TTL caching** (rate-limit friendly) +- **aiosqlite** for guild config + persistent researcher notes +- Full **logging** and graceful shutdown +- Severity-coded **cyberpunk embeds** (neon reds/oranges/cyans on near-black) + +--- + +## 📁 Project Structure + +``` +vulnforge/ +├── main.py # Entrypoint: boots bot, wires API + store, loads cogs +├── config.py # Env + endpoints + colour palette + tunables +├── .env.example # Copy to .env and fill in +├── requirements.txt +├── README.md +├── Dockerfile # Hardened, non-root production image +├── docker-compose.yml # One-command deploy with persistent volume +├── .dockerignore +├── cogs/ +│ ├── __init__.py +│ ├── weekly.py # Scheduled digest + /weekly +│ └── cve.py # /cve modal, /kev, /setup +├── utils/ +│ ├── __init__.py +│ ├── api.py # NVD + CISA KEV async client (cache + retry) +│ ├── embeds.py # Cyberpunk embed builders +│ └── helpers.py # Severity maps, validation, JSON cache, SQLite store +└── data/ # Caches + SQLite DB (auto-created) +``` + +--- + +## 🚀 Setup + +### 1. Create a Discord application +1. Go to the [Discord Developer Portal](https://discord.com/developers/applications). +2. **New Application** → **Bot** → copy the **token**. +3. Under **Installation / OAuth2**, enable the `bot` and `applications.commands` scopes. +4. Invite the bot to your server with the **Send Messages** + **Embed Links** permissions. + +### 2. Install +```powershell +cd vulnforge +python -m venv .venv +.\.venv\Scripts\Activate.ps1 # Windows +# source .venv/bin/activate # macOS/Linux +pip install -r requirements.txt +``` + +### 3. Configure +```powershell +copy .env.example .env # cp on macOS/Linux +``` +Edit `.env`: +- `DISCORD_TOKEN` — **required** +- `GUILD_ID` — optional; set it for **instant** slash-command sync during development +- `NVD_API_KEY` — optional but **strongly recommended** ([request one free](https://nvd.nist.gov/developers/request-an-api-key)) to lift rate limits + +### 4. Run +```powershell +python main.py +``` + +### 5. Configure in Discord +Run `/setup` in your server (needs **Manage Server**) to choose the digest channel and an optional notification role. Then test with `/weekly`. + +--- + +## 🐳 Docker + +VulnForge ships with a hardened image (non-root user, healthcheck, persistent volume). It only makes **outbound** calls, so no ports are exposed. + +### Compose (recommended) +```powershell +copy .env.example .env # fill in DISCORD_TOKEN +docker compose up -d --build +docker compose logs -f # tail logs +docker compose down # stop +``` + +### Plain Docker +```powershell +docker build -t vulnforge:latest . +docker run -d --name vulnforge --restart unless-stopped ` + --env-file .env ` + -v vulnforge_data:/app/data ` + vulnforge:latest +``` + +The named volume `vulnforge_data` persists the SQLite DB (guild config + researcher notes) and the on-disk API cache across rebuilds. Update `.env` and re-run `docker compose up -d --build` to apply config changes. + +--- + +## 🎨 Severity Colour Legend + +| Severity | Marker | Colour | +|----------|--------|--------| +| Critical | 🟥 | Neon Red `#FF003C` | +| High | 🟧 | Hot Orange `#FF6B00` | +| Medium | 🟨 | Amber `#FFB300` | +| Low | 🟦 | Cyan `#00E5FF` | +| KEV / Exploited | 🩸 | Ruby `#E0115F` | + +--- + +## 🔒 Notes & Limits +- **NVD rate limits**: 5 requests / 30s without a key, 50 / 30s with one. The weekly sweep paginates the full 7-day window — a key makes it noticeably faster. Responses are cached on disk (`data/cache/`). +- **KEV feed** is cached for 6 hours by default (`config.py → kev_cache_ttl`). +- Never commit your real `.env` — only `.env.example` belongs in version control. +- This tool is for **defensive research and awareness**. Use responsibly. + +--- + +## 🛠 Customisation +- Change the digest day/time in [cogs/weekly.py](cogs/weekly.py#L34) (`RUN_AT` + the Monday guard). +- Tune `digest_top_n`, cache TTLs, and the colour palette in [config.py](config.py). +- Embed styling lives entirely in [utils/embeds.py](utils/embeds.py). + +--- + +*VulnForge // Security Intelligence — built with Nextcord.* diff --git a/check_logs.sh b/check_logs.sh new file mode 100755 index 0000000..915d984 --- /dev/null +++ b/check_logs.sh @@ -0,0 +1,2 @@ +#!/bin/bash +docker logs vulnforge 2>&1 | grep -iE 'sync|command|extension|error|exception|guild|cog|loaded' | tail -60 diff --git a/cogs/__init__.py b/cogs/__init__.py new file mode 100755 index 0000000..262abdd --- /dev/null +++ b/cogs/__init__.py @@ -0,0 +1 @@ +"""cogs package — VulnForge feature modules (weekly digest, CVE lookups).""" diff --git a/cogs/cve.py b/cogs/cve.py new file mode 100755 index 0000000..44d5866 --- /dev/null +++ b/cogs/cve.py @@ -0,0 +1,261 @@ +""" +cogs/cve.py — interactive CVE lookups + KEV listing + setup. + +Slash commands: + /cve → opens a modal letting the user toggle which sections to show, + then replies with a rich severity-coloured embed. + /kev → latest CISA KEV additions. + /setup → admin-only; configure digest channel + notification role. + +The modal pattern: the slash command pre-validates the CVE id, then sends a +modal so the user can pick options (full details / KEV status / references / +CWE / add a note). On submit we fetch from NVD + KEV and render the embed. +""" + +from __future__ import annotations + +import logging + +import nextcord +from nextcord import Interaction, SlashOption +from nextcord.ext import commands + +from config import settings +from utils.embeds import ( + build_cve_embed, + build_error_embed, + build_info_embed, + build_kev_list_embed, +) +from utils.helpers import is_valid_cve_id, normalize_cve_id + +log = logging.getLogger("vulnforge.cogs.cve") + +# Restrict slash sync to a guild when GUILD_ID is set (instant updates in dev). +GUILD_IDS = [settings.guild_id] if settings.guild_id else None + + +class CVEOptionsModal(nextcord.ui.Modal): + """Modal shown after /cve — lets the user choose output sections + a note.""" + + def __init__(self, cog: "CVECog", cve_id: str) -> None: + super().__init__(title=f"Query {cve_id}", timeout=300) + self._cog = cog + self._cve_id = cve_id + + # Free-text "options" field — keeps it works-everywhere simple: + # comma-separated toggles, defaulting to "all". + self.sections = nextcord.ui.TextInput( + label="Sections (details, kev, refs, cwe, all)", + placeholder="all —or— details, kev, refs", + default_value="all", + required=False, + max_length=100, + style=nextcord.TextInputStyle.short, + ) + self.add_item(self.sections) + + # Optional researcher note persisted to SQLite and shown on the embed. + self.note = nextcord.ui.TextInput( + label="Add a researcher note (optional)", + placeholder="e.g. Confirmed exploitable on internal asset X", + required=False, + max_length=300, + style=nextcord.TextInputStyle.paragraph, + ) + self.add_item(self.note) + + async def callback(self, interaction: Interaction) -> None: + await interaction.response.defer(ephemeral=False) + + raw = (self.sections.value or "all").lower() + selected = {s.strip() for s in raw.split(",") if s.strip()} + want_all = "all" in selected or not selected + + show_details = want_all or "details" in selected + show_kev = want_all or "kev" in selected + show_refs = want_all or any(k in selected for k in ("refs", "references")) + show_cwe = want_all or "cwe" in selected + + api = self._cog.bot.api + cve = await api.fetch_cve(self._cve_id) + if cve is None: + await interaction.followup.send( + embed=build_error_embed( + "CVE Not Found", + f"`{self._cve_id}` was not found on NVD, or the API is " + "unavailable. Double-check the identifier and try again.", + ) + ) + return + + kev = await api.get_kev_entry(self._cve_id) + + # Persist + collect notes for this CVE. + note_text = (self.note.value or "").strip() + if note_text: + try: + await self._cog.bot.store.add_note( + interaction.guild_id, + interaction.user.id, + self._cve_id, + note_text, + ) + except Exception as exc: # noqa: BLE001 — best-effort persistence + log.warning("Failed saving note for %s: %s", self._cve_id, exc) + notes = await self._cog.bot.store.get_notes(self._cve_id) + + embed = build_cve_embed( + cve, + kev, + show_details=show_details, + show_kev=show_kev, + show_references=show_refs, + show_cwe=show_cwe, + notes=notes, + ) + await interaction.followup.send(embed=embed) + + +class SetupModal(nextcord.ui.Modal): + """Admin modal to set digest channel + notification role by ID.""" + + def __init__(self, cog: "CVECog") -> None: + super().__init__(title="VulnForge Setup", timeout=300) + self._cog = cog + + self.channel_id = nextcord.ui.TextInput( + label="Digest channel ID", + placeholder="Right-click a channel → Copy Channel ID", + required=True, + max_length=25, + style=nextcord.TextInputStyle.short, + ) + self.add_item(self.channel_id) + + self.role_id = nextcord.ui.TextInput( + label="Notification role ID (optional)", + placeholder="Role to ping on the weekly digest", + required=False, + max_length=25, + style=nextcord.TextInputStyle.short, + ) + self.add_item(self.role_id) + + async def callback(self, interaction: Interaction) -> None: + await interaction.response.defer(ephemeral=True) + + def _parse(value: str) -> int | None: + value = (value or "").strip() + return int(value) if value.isdigit() else None + + channel_id = _parse(self.channel_id.value) + role_id = _parse(self.role_id.value) + + if channel_id is None: + await interaction.followup.send( + embed=build_error_embed("Invalid Channel", "That channel ID is not numeric."), + ephemeral=True, + ) + return + + channel = interaction.guild.get_channel(channel_id) if interaction.guild else None + if channel is None: + await interaction.followup.send( + embed=build_error_embed( + "Channel Not Found", + "I couldn't find that channel in this server. " + "Make sure the ID is correct and I can see the channel.", + ), + ephemeral=True, + ) + return + + await self._cog.bot.store.set_guild_config( + interaction.guild_id, channel_id, role_id + ) + role_note = f" • notifying <@&{role_id}>" if role_id else "" + await interaction.followup.send( + embed=build_info_embed( + "Setup Complete", + f"Weekly digest will post to {channel.mention}{role_note}.\n" + "Use `/threatlens_weekly` to fire a test report.", + ), + ephemeral=True, + ) + + +class CVECog(commands.Cog): + """Interactive CVE/KEV slash commands.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + + @nextcord.slash_command( + name="threatlens_cve", + description="Look up a CVE and choose what details to display.", + guild_ids=GUILD_IDS, + ) + async def cve( + self, + interaction: Interaction, + cve_id: str = SlashOption( + description="CVE identifier, e.g. CVE-2024-3094", + required=True, + ), + ) -> None: + cve_id = normalize_cve_id(cve_id) + if not is_valid_cve_id(cve_id): + await interaction.response.send_message( + embed=build_error_embed( + "Invalid CVE ID", + f"`{cve_id}` is not a valid identifier. Expected format: " + "`CVE-YYYY-NNNN` (e.g. `CVE-2024-3094`).", + ), + ephemeral=True, + ) + return + # Open the options modal (must be the first response to the interaction). + await interaction.response.send_modal(CVEOptionsModal(self, cve_id)) + + @nextcord.slash_command( + name="threatlens_kev", + description="Show the latest CISA Known Exploited Vulnerabilities.", + guild_ids=GUILD_IDS, + ) + async def kev( + self, + interaction: Interaction, + limit: int = SlashOption( + description="How many entries to show (1-15)", + required=False, + default=10, + min_value=1, + max_value=15, + ), + ) -> None: + await interaction.response.defer() + entries = await self.bot.api.latest_kev(limit=limit) + await interaction.followup.send(embed=build_kev_list_embed(entries, limit)) + + @nextcord.slash_command( + name="threatlens_setup", + description="Admin: configure the weekly digest channel and role.", + guild_ids=GUILD_IDS, + ) + async def setup_cmd(self, interaction: Interaction) -> None: + perms = interaction.user.guild_permissions if interaction.guild else None + if not perms or not perms.manage_guild: + await interaction.response.send_message( + embed=build_error_embed( + "Permission Denied", + "You need the **Manage Server** permission to run setup.", + ), + ephemeral=True, + ) + return + await interaction.response.send_modal(SetupModal(self)) + + +def setup(bot: commands.Bot) -> None: + bot.add_cog(CVECog(bot)) diff --git a/cogs/enrich.py b/cogs/enrich.py new file mode 100755 index 0000000..5d96756 --- /dev/null +++ b/cogs/enrich.py @@ -0,0 +1,314 @@ +""" +cogs/enrich.py — /threatlens_enrich slash command. + +Accepts a markdown report (uploaded as an attachment OR pasted via a follow-up +URL parameter), runs it through utils.enricher, and posts the enriched +markdown back into the channel as a file attachment along with a stats embed +summarising what was found. + +Why a file attachment instead of an embed body? + - Discord caps embed descriptions at 4 096 chars and total message at 6 000 + chars; a real pentest report is 30-150 KB. The enriched output (with all + inline links + an INTEL APPENDIX) easily breaks that. A .md attachment + preserves every clickable link and is downloadable by the team. +""" + +from __future__ import annotations + +import io +import logging +from datetime import datetime, timezone + +import aiohttp +import nextcord +from nextcord import Interaction, SlashOption +from nextcord.ext import commands + +from config import settings +from utils.embeds import build_error_embed +from utils.enricher import enrich_markdown + +log = logging.getLogger("vulnforge.cogs.enrich") + +GUILD_IDS = [settings.guild_id] if settings.guild_id else None + +MAX_INPUT_BYTES = 5 * 1024 * 1024 # 5 MB — protects the bot from huge uploads +ALLOWED_EXTS = (".md", ".markdown", ".txt") + + +def _build_stats_embed(filename: str, byte_count: int, stats_summary: dict[str, int]) -> nextcord.Embed: + """Compact embed showing what the enricher found.""" + nonzero = {k: v for k, v in stats_summary.items() if v} + total = sum(nonzero.values()) + + embed = nextcord.Embed( + title=f"📎 Report Enriched — `{filename}`", + description=( + f"**{total}** artifacts linked across **{len(nonzero)}** categories.\n" + f"Output size: **{byte_count:,} bytes**.\n" + "Every artifact below appears in the report inline as a clickable " + "link AND in the new **INTEL APPENDIX** at the bottom of the file " + "with full Shodan / Censys / NVD / VirusTotal / GitHub pivots." + ), + color=settings.color_success, + ) + + if nonzero: + # Two columns of artifact categories with counts. + items = list(nonzero.items()) + mid = (len(items) + 1) // 2 + col_a = "\n".join(f"• **{k}** — `{v}`" for k, v in items[:mid]) + col_b = "\n".join(f"• **{k}** — `{v}`" for k, v in items[mid:]) + embed.add_field(name="Artifacts", value=col_a or "—", inline=True) + if col_b: + embed.add_field(name="\u200b", value=col_b, inline=True) + + embed.set_footer(text="⛓ VulnForge // Report Enricher") + embed.timestamp = datetime.now(timezone.utc) + return embed + + +async def _fetch_url(session: aiohttp.ClientSession, url: str) -> tuple[bytes, str]: + """Fetch a remote markdown file. Returns (bytes, derived_filename).""" + async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp: + resp.raise_for_status() + raw = await resp.read() + # Derive a filename from the URL path; fall back to "report.md". + name = url.rsplit("/", 1)[-1].split("?", 1)[0] or "report.md" + if not name.lower().endswith(ALLOWED_EXTS): + name = name + ".md" + return raw, name + + +def _parse_message_ref(ref: str) -> int | None: + """Accept a numeric ID or a full Discord message link and return the id.""" + ref = ref.strip() + if ref.isdigit(): + return int(ref) + # discord.com/channels/// + parts = ref.rstrip("/").split("/") + if len(parts) >= 1 and parts[-1].isdigit(): + return int(parts[-1]) + return None + + +def _message_to_markdown(msg: nextcord.Message) -> str: + """Flatten a Discord message (content + embeds) into a markdown blob. + + Embeds become `## title\n description\n ### field name\n field value`, + so the enricher sees every CVE / IP / URL the bot itself rendered. + """ + chunks: list[str] = [] + if msg.content: + chunks.append(msg.content) + for emb in msg.embeds: + if emb.title: + chunks.append(f"## {emb.title}") + if emb.url: + chunks.append(emb.url) + if emb.description: + chunks.append(emb.description) + for f in emb.fields: + if f.name and f.name not in ("\u200b",): + chunks.append(f"### {f.name}") + if f.value: + chunks.append(f.value) + if emb.footer and emb.footer.text: + chunks.append(f"_{emb.footer.text}_") + return "\n\n".join(chunks) + + +class EnrichCog(commands.Cog): + """Markdown report enrichment via slash command.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + + @nextcord.slash_command( + name="threatlens_enrich", + description="Enrich a pentest markdown report with clickable artifact / PoC / source links.", + guild_ids=GUILD_IDS, + ) + async def enrich( + self, + interaction: Interaction, + file: nextcord.Attachment | None = SlashOption( + name="file", + description="Markdown report (.md / .txt) to enrich.", + required=False, + ), + url: str | None = SlashOption( + name="url", + description="Or fetch a public raw-markdown URL (e.g. GitHub raw, gist).", + required=False, + ), + message: str | None = SlashOption( + name="message", + description="Or enrich a prior message: paste its ID or full Discord link.", + required=False, + ), + recent: bool | None = SlashOption( + name="recent", + description="Or auto-grab my last bot message in this channel and enrich it.", + required=False, + ), + ) -> None: + await interaction.response.defer(ephemeral=False) + + if not file and not url and not message and not recent: + await interaction.followup.send( + embed=build_error_embed( + "Nothing to enrich", + "Pick one source: attach a `file`, pass a `url`, pass a " + "`message` ID/link, or set `recent: true` to enrich the " + "most recent bot message in this channel.", + ), + ) + return + + # Resolve the source bytes + a filename. + try: + if file: + if file.size and file.size > MAX_INPUT_BYTES: + await interaction.followup.send( + embed=build_error_embed( + "File too large", + f"`{file.filename}` is {file.size:,} bytes; the limit " + f"is {MAX_INPUT_BYTES:,}.", + ), + ) + return + if file.filename and not file.filename.lower().endswith(ALLOWED_EXTS): + await interaction.followup.send( + embed=build_error_embed( + "Unsupported file type", + f"`{file.filename}` is not a markdown file. " + f"Allowed extensions: {', '.join(ALLOWED_EXTS)}.", + ), + ) + return + raw = await file.read() + src_name = file.filename or "report.md" + elif message or recent: + channel = interaction.channel + if channel is None: + await interaction.followup.send( + embed=build_error_embed( + "No channel", + "Could not resolve the channel for this interaction.", + ), + ) + return + target_msg: nextcord.Message | None = None + if message: + mid = _parse_message_ref(message) + if mid is None: + await interaction.followup.send( + embed=build_error_embed( + "Bad message reference", + "Pass a numeric message ID or a full Discord " + "message link (`discord.com/channels/.../.../`).", + ), + ) + return + target_msg = await channel.fetch_message(mid) + else: + # recent=True — find the most recent bot-authored message + # in this channel that has content or embeds. + me = self.bot.user + async for m in channel.history(limit=25): + if me and m.author.id == me.id and (m.content or m.embeds): + # skip the interaction's own deferred response if + # discord happens to surface it; we want command output. + target_msg = m + break + if target_msg is None: + await interaction.followup.send( + embed=build_error_embed( + "No recent bot message", + "Couldn't find a recent VulnForge message in this " + "channel to enrich. Run a `/threatlens_*` command first.", + ), + ) + return + md = _message_to_markdown(target_msg) + if not md.strip(): + await interaction.followup.send( + embed=build_error_embed( + "Empty message", + "That message has no text or embed content to enrich.", + ), + ) + return + raw = md.encode("utf-8") + src_name = f"discord-msg-{target_msg.id}.md" + else: + # Reuse the bot's existing aiohttp session if present so we + # don't open another connection pool per call. + session = getattr(getattr(self.bot, "api", None), "_session", None) + if session is None or session.closed: + async with aiohttp.ClientSession() as fresh: + raw, src_name = await _fetch_url(fresh, url) # type: ignore[arg-type] + else: + raw, src_name = await _fetch_url(session, url) # type: ignore[arg-type] + + if len(raw) > MAX_INPUT_BYTES: + await interaction.followup.send( + embed=build_error_embed( + "Fetched file too large", + f"Fetched {len(raw):,} bytes from the URL; the limit " + f"is {MAX_INPUT_BYTES:,}.", + ), + ) + return + except aiohttp.ClientError as exc: + log.warning("URL fetch failed: %s", exc) + await interaction.followup.send( + embed=build_error_embed("Fetch failed", f"`{exc}`"), + ) + return + except Exception: # noqa: BLE001 — surface any unexpected error to chat + log.exception("enrich: failed to read source") + await interaction.followup.send( + embed=build_error_embed("Read failed", "Could not read the input."), + ) + return + + try: + text = raw.decode("utf-8", errors="replace") + except Exception: # noqa: BLE001 + await interaction.followup.send( + embed=build_error_embed( + "Decode failed", + "Could not decode the file as UTF-8. Make sure it is a text/markdown file.", + ), + ) + return + + try: + enriched, stats = enrich_markdown(text) + except Exception: # noqa: BLE001 — never crash the bot on bad input + log.exception("enrich: enrichment pass crashed") + await interaction.followup.send( + embed=build_error_embed( + "Enrichment failed", + "The enrichment pass threw an unexpected error. Check the bot logs.", + ), + ) + return + + # Build output filename: foo.md -> foo.enriched.md + stem = src_name.rsplit(".", 1)[0] if "." in src_name else src_name + out_name = f"{stem}.enriched.md" + + out_bytes = enriched.encode("utf-8") + buf = io.BytesIO(out_bytes) + attachment = nextcord.File(buf, filename=out_name) + + embed = _build_stats_embed(out_name, len(out_bytes), stats.as_summary()) + + await interaction.followup.send(embed=embed, file=attachment) + + +def setup(bot: commands.Bot) -> None: + bot.add_cog(EnrichCog(bot)) diff --git a/cogs/weekly.py b/cogs/weekly.py new file mode 100755 index 0000000..bc5e3dd --- /dev/null +++ b/cogs/weekly.py @@ -0,0 +1,159 @@ +""" +cogs/weekly.py — automated weekly threat digest. + +A background loop fires every Monday at 09:00 UTC and posts a rich digest to +each configured guild channel: + * total new CVEs (last 7 days) + * how many are already in CISA KEV / newly added to KEV + * the top 8 most critical CVEs by CVSS + * source links + +Admins can force a run with /weekly. The schedule uses nextcord.ext.tasks with +an explicit UTC time, and a guard so it only sends on Mondays. +""" + +from __future__ import annotations + +import datetime as dt +import logging + +import nextcord +from nextcord import Interaction +from nextcord.ext import commands, tasks + +from config import settings +from utils.api import CVERecord, KEVEntry +from utils.embeds import build_error_embed, build_info_embed, build_weekly_embed +from utils.helpers import SEVERITY_ORDER + +log = logging.getLogger("vulnforge.cogs.weekly") + +GUILD_IDS = [settings.guild_id] if settings.guild_id else None + +# 09:00 UTC trigger. The loop runs every 24h aligned to this time; we then +# check the weekday so it only actually posts on Mondays. +RUN_AT = dt.time(hour=9, minute=0, tzinfo=dt.timezone.utc) + + +class WeeklyCog(commands.Cog): + """Scheduled weekly digest + manual trigger.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + self.weekly_digest.start() + + def cog_unload(self) -> None: + self.weekly_digest.cancel() + + # ── scheduled loop ─────────────────────────────────────── + @tasks.loop(time=RUN_AT) + async def weekly_digest(self) -> None: + # Only fire on Mondays (weekday() == 0). + if dt.datetime.now(dt.timezone.utc).weekday() != 0: + return + log.info("Weekly digest schedule triggered (Monday 09:00 UTC)") + await self._run_for_all_guilds() + + @weekly_digest.before_loop + async def _before(self) -> None: + await self.bot.wait_until_ready() + + # ── data assembly ──────────────────────────────────────── + async def _build_digest_embed(self, days: int = 7) -> nextcord.Embed: + api = self.bot.api + recent: list[CVERecord] = await api.fetch_recent_cves(days=days) + newly_added_kev: list[KEVEntry] = await api.kev_added_since(days=days) + + # Count how many recent CVEs are present in KEV. + in_kev = 0 + for cve in recent: + if await api.is_in_kev(cve.cve_id): + in_kev += 1 + + # Top N by CVSS score (then by severity rank as tiebreaker). + scored = [c for c in recent if c.cvss_score is not None] + scored.sort( + key=lambda c: (c.cvss_score or 0.0, SEVERITY_ORDER.get(c.cvss_severity, 0)), + reverse=True, + ) + top = scored[: settings.digest_top_n] + + return build_weekly_embed( + total_new=len(recent), + in_kev_count=in_kev, + newly_added_kev=newly_added_kev, + top_cves=top, + days=days, + ) + + async def _run_for_all_guilds(self) -> None: + embed = await self._build_digest_embed() + configs = await self.bot.store.all_guild_configs() + + # Fall back to env-configured channel if no guild has run /setup. + if not configs and settings.digest_channel_id: + configs = [ + { + "digest_channel_id": settings.digest_channel_id, + "notify_role_id": settings.notify_role_id, + } + ] + + for cfg in configs: + channel = self.bot.get_channel(cfg["digest_channel_id"]) + if channel is None: + log.warning("Digest channel %s not found/visible", cfg["digest_channel_id"]) + continue + content = f"<@&{cfg['notify_role_id']}>" if cfg.get("notify_role_id") else None + try: + await channel.send(content=content, embed=embed) + log.info("Weekly digest posted to channel %s", channel.id) + except nextcord.DiscordException as exc: + log.error("Failed posting digest to %s: %s", channel.id, exc) + + # ── admin command ──────────────────────────────────────── + @nextcord.slash_command( + name="threatlens_weekly", + description="Admin: force-post the weekly threat digest now.", + guild_ids=GUILD_IDS, + ) + async def weekly(self, interaction: Interaction) -> None: + perms = interaction.user.guild_permissions if interaction.guild else None + if not perms or not perms.manage_guild: + await interaction.response.send_message( + embed=build_error_embed( + "Permission Denied", + "You need the **Manage Server** permission to force a digest.", + ), + ephemeral=True, + ) + return + + await interaction.response.defer(ephemeral=True) + embed = await self._build_digest_embed() + + cfg = await self.bot.store.get_guild_config(interaction.guild_id) + target_id = (cfg or {}).get("digest_channel_id") or settings.digest_channel_id + role_id = (cfg or {}).get("notify_role_id") or settings.notify_role_id + + channel = self.bot.get_channel(target_id) if target_id else interaction.channel + if channel is None: + await interaction.followup.send( + embed=build_error_embed( + "No Channel Configured", + "Run `/setup` first to choose a digest channel.", + ), + ephemeral=True, + ) + return + + content = f"<@&{role_id}>" if role_id else None + await channel.send(content=content, embed=embed) + await interaction.followup.send( + embed=build_info_embed("Digest Sent", f"Weekly digest posted to {channel.mention}."), + ephemeral=True, + ) + + +def setup(bot: commands.Bot) -> None: + bot.add_cog(WeeklyCog(bot)) diff --git a/config.py b/config.py new file mode 100755 index 0000000..3a22181 --- /dev/null +++ b/config.py @@ -0,0 +1,85 @@ +""" +config.py — Centralised configuration for VulnForge. + +Loads environment variables, defines API endpoints, severity colour maps, +cache TTLs and other tunables. Import `settings` everywhere instead of reading +os.environ directly so configuration stays in one place. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path + +from dotenv import load_dotenv + +# Load .env from the project root (the directory this file lives in). +PROJECT_ROOT = Path(__file__).resolve().parent +load_dotenv(PROJECT_ROOT / ".env") + + +def _get_int(name: str) -> int | None: + """Parse an optional integer env var, returning None when unset/blank.""" + raw = os.getenv(name, "").strip() + if not raw: + return None + try: + return int(raw) + except ValueError: + return None + + +@dataclass(frozen=True) +class Settings: + """Immutable runtime settings sourced from environment + defaults.""" + + # ── Discord ────────────────────────────────────────────── + discord_token: str = os.getenv("DISCORD_TOKEN", "").strip() + guild_id: int | None = _get_int("GUILD_ID") + digest_channel_id: int | None = _get_int("DIGEST_CHANNEL_ID") + notify_role_id: int | None = _get_int("NOTIFY_ROLE_ID") + + # ── External APIs ──────────────────────────────────────── + nvd_api_key: str = os.getenv("NVD_API_KEY", "").strip() + nvd_base_url: str = "https://services.nvd.nist.gov/rest/json/cves/2.0" + cisa_kev_url: str = ( + "https://www.cisa.gov/sites/default/files/feeds/" + "known_exploited_vulnerabilities.json" + ) + nvd_cve_web: str = "https://nvd.nist.gov/vuln/detail" + cisa_kev_web: str = "https://www.cisa.gov/known-exploited-vulnerabilities-catalog" + + # ── Paths ──────────────────────────────────────────────── + data_dir: Path = PROJECT_ROOT / "data" + db_path: Path = PROJECT_ROOT / "data" / "vulnforge.db" + cache_dir: Path = PROJECT_ROOT / "data" / "cache" + + # ── Behaviour ──────────────────────────────────────────── + log_level: str = os.getenv("LOG_LEVEL", "INFO").strip().upper() + # Cache TTL (seconds): KEV feed is large + slow-changing. + kev_cache_ttl: int = 60 * 60 * 6 # 6 hours + cve_cache_ttl: int = 60 * 30 # 30 minutes + http_timeout: int = 30 # seconds per request + digest_top_n: int = 8 # CVEs shown in weekly digest + request_retries: int = 3 # NVD transient retry attempts + + # ── Cyberpunk / hacker colour palette (hex → int) ───────── + # Severity-coded so embeds read at a glance on mobile + desktop. + color_critical: int = 0xFF003C # neon red + color_high: int = 0xFF6B00 # hot orange + color_medium: int = 0xFFB300 # amber + color_low: int = 0x00E5FF # cyan + color_info: int = 0x0A84FF # electric blue + color_kev: int = 0xE0115F # ruby (actively exploited) + color_dark: int = 0x0D1117 # near-black backdrop accent + color_success: int = 0x00FF9C # matrix green + user_agent: str = "VulnForge/1.0 (+https://github.com/your-org/vulnforge)" + + def ensure_dirs(self) -> None: + """Create data/cache directories if they do not yet exist.""" + self.data_dir.mkdir(parents=True, exist_ok=True) + self.cache_dir.mkdir(parents=True, exist_ok=True) + + +settings = Settings() diff --git a/deep_inspect.sh b/deep_inspect.sh new file mode 100755 index 0000000..77f18b2 --- /dev/null +++ b/deep_inspect.sh @@ -0,0 +1,33 @@ +#!/bin/bash +echo "=== GLOBAL COMMANDS ===" +bash /tmp/diag.sh 2>&1 | sed -n '/GLOBAL COMMANDS/,/GUILD COMMANDS/p' +echo +echo "=== GUILD COMMANDS ===" +bash /tmp/diag.sh 2>&1 | sed -n '/GUILD COMMANDS/,/BOT MEMBER/p' +echo +echo "=== inspect bot internals ===" +docker exec vulnforge python -c " +import sys; sys.path.insert(0, '/app') +from main import VulnForge +b = VulnForge() +b.load_extension('cogs.cve') +b.load_extension('cogs.weekly') +print('Before manual add:') +print(' get_application_commands():', len(b.get_application_commands())) +print(' get_all_application_commands():', len(list(b.get_all_application_commands()))) +print(' _application_commands_to_add:', b._application_commands_to_add) +print() +# Manual add +for cog_name, cog in b.cogs.items(): + for cmd in getattr(cog, 'application_commands', []): + cmd.guild_ids_to_rollout.add(1502477191093420135) + b.add_application_command(cmd, use_rollout=True) +print('After manual add:') +print(' get_application_commands():', len(b.get_application_commands())) +print(' get_all_application_commands():', len(list(b.get_all_application_commands()))) +print(' _application_commands_to_add:', len(b._application_commands_to_add) if b._application_commands_to_add else 0) +print() +print('Commands details:') +for c in b.get_all_application_commands(): + print(' -', c.name, 'guild_ids_to_rollout:', c.guild_ids_to_rollout, 'is_global:', c.is_global) +" diff --git a/deploy_check.sh b/deploy_check.sh new file mode 100755 index 0000000..93eb87f --- /dev/null +++ b/deploy_check.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -e +cd ~/vulnforge +docker compose up -d --build 2>&1 | tail -4 +sleep 18 +echo "=== bot logs ===" +docker logs vulnforge 2>&1 | grep -E 'vulnforge .' | head -30 +echo +echo "=== diag ===" +bash /tmp/diag.sh | tail -40 diff --git a/diag.sh b/diag.sh new file mode 100755 index 0000000..12f197d --- /dev/null +++ b/diag.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Diagnose why VulnForge slash commands aren't showing in Discord. +cd ~/vulnforge +# Strip Windows CRLF before sourcing .env (otherwise vars carry trailing \r). +sed -i 's/\r$//' .env +set -a; . ./.env; set +a +AUTH="Authorization: Bot $DISCORD_TOKEN" +APP_ID=$(curl -s -H "$AUTH" https://discord.com/api/v10/oauth2/applications/@me | python3 -c "import sys,json;print(json.load(sys.stdin)['id'])") +echo "=== APP_ID=$APP_ID GUILD_ID=$GUILD_ID ===" +echo +echo "=== BOT GUILDS ===" +curl -s -H "$AUTH" https://discord.com/api/v10/users/@me/guilds | python3 -m json.tool +echo +echo "=== GLOBAL COMMANDS ===" +curl -s -H "$AUTH" "https://discord.com/api/v10/applications/$APP_ID/commands" | python3 -m json.tool +echo +echo "=== GUILD COMMANDS for $GUILD_ID ===" +curl -s -H "$AUTH" "https://discord.com/api/v10/applications/$APP_ID/guilds/$GUILD_ID/commands" | python3 -m json.tool +echo +echo "=== BOT MEMBER IN GUILD ===" +curl -s -H "$AUTH" "https://discord.com/api/v10/guilds/$GUILD_ID/members/$APP_ID" | python3 -m json.tool diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100755 index 0000000..4a09403 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,24 @@ +# ────────────────────────────────────────────────────────────── +# VulnForge — docker compose +# Usage: +# 1. copy .env.example .env (and fill in DISCORD_TOKEN) +# 2. docker compose up -d --build +# 3. docker compose logs -f +# ────────────────────────────────────────────────────────────── +services: + vulnforge: + build: . + image: vulnforge:latest + container_name: vulnforge + restart: unless-stopped + # Loads DISCORD_TOKEN, NVD_API_KEY, GUILD_ID, etc. from .env + env_file: + - .env + # Persist SQLite DB + cached API responses across container rebuilds. + volumes: + - vulnforge_data:/app/data + # No inbound ports needed — the bot only makes outbound calls + # (Discord gateway, NVD, CISA). + +volumes: + vulnforge_data: diff --git a/inspect_cmds.sh b/inspect_cmds.sh new file mode 100755 index 0000000..b1cab85 --- /dev/null +++ b/inspect_cmds.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Inspect what slash commands nextcord actually knows about. +docker exec vulnforge python -c " +import sys; sys.path.insert(0, '/app') +from main import VulnForge +bot = VulnForge() +# Manually load the extensions so we can see the result. +bot.load_extension('cogs.cve') +bot.load_extension('cogs.weekly') +print('=== get_application_commands ===') +for c in bot.get_application_commands(): + print(' -', type(c).__name__, repr(getattr(c, 'name', None)), 'guild_ids=', getattr(c, 'guild_ids', None)) +print() +print('=== cogs ===') +for name, cog in bot.cogs.items(): + print(' -', name, 'application_commands:', list(getattr(cog, 'application_commands', []))) +print() +print('=== _application_commands_to_add (internal) ===') +internal = getattr(bot, '_application_commands_to_add', None) or getattr(bot, '_pending_application_commands', None) +print(' ', internal) +print() +print('=== all attributes containing \"command\" ===') +for a in dir(bot): + if 'command' in a.lower() and not a.startswith('__'): + print(' -', a) +" diff --git a/main.py b/main.py new file mode 100755 index 0000000..b86e938 --- /dev/null +++ b/main.py @@ -0,0 +1,160 @@ +""" +main.py — VulnForge entrypoint. + +Boots the Nextcord bot, wires up the shared SecurityAPI client + SQLite store, +loads cogs, configures logging, and handles graceful startup/shutdown. + +Run: + python main.py +""" + +from __future__ import annotations + +import asyncio +import logging +import signal +import sys + +import nextcord +from nextcord.ext import commands + +from config import settings +from utils.api import SecurityAPI +from utils.helpers import Store + +# ────────────────────────────────────────────────────────────── +# Logging +# ────────────────────────────────────────────────────────────── +logging.basicConfig( + level=getattr(logging, settings.log_level, logging.INFO), + format="%(asctime)s │ %(levelname)-8s │ %(name)s │ %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], +) +log = logging.getLogger("vulnforge") + + +class VulnForge(commands.Bot): + """The VulnForge bot — owns the shared API client and data store.""" + + def __init__(self) -> None: + intents = nextcord.Intents.default() + # When GUILD_ID is set we scope ALL commands to that guild so they appear + # instantly (and we don't depend on Nextcord auto-discovering guild_ids + # from individual decorators, which was failing to push the guild sync). + kwargs: dict = {"intents": intents} + if settings.guild_id: + kwargs["default_guild_ids"] = [settings.guild_id] + super().__init__(**kwargs) + self.api = SecurityAPI() + self.store = Store(settings.db_path) + + async def on_ready(self) -> None: + log.info("Logged in as %s (id: %s)", self.user, self.user.id) + log.info("Connected to %s guild(s)", len(self.guilds)) + # Force a fresh sync so command renames/additions show up immediately. + try: + if settings.guild_id: + await self.sync_application_commands(guild_id=settings.guild_id) + log.info("Slash commands synced to guild %s", settings.guild_id) + else: + await self.sync_application_commands() + log.info("Slash commands synced globally") + except Exception: # noqa: BLE001 + log.exception("sync_application_commands failed") + await self.change_presence( + activity=nextcord.Activity( + type=nextcord.ActivityType.watching, + name="the CVE & KEV feeds ☣", + ) + ) + + async def setup(self) -> None: + """Load extensions and register cog commands BEFORE we connect. + + We can't use Nextcord's `setup_hook` here — in 2.6 it isn't invoked by + bot.start() in this configuration, so cogs never got loaded and the + bot's command list stayed empty (sync was a no-op). Call this manually + from the runner before bot.start(). + """ + settings.ensure_dirs() + await self.store.init() + await self.api.start() + + for ext in ("cogs.cve", "cogs.weekly", "cogs.enrich"): + try: + self.load_extension(ext) + log.info("Loaded extension: %s", ext) + except Exception: # noqa: BLE001 + log.exception("Failed to load extension %s", ext) + + # Nextcord 2.6 quirk: add_cog() attaches slash commands to the cog + # object (cog.application_commands) but does NOT push them into the + # bot's registration list. add_all_cog_commands() also no-ops here. + # Manually push each cog command onto the bot so sync sees them. + for cog_name, cog in self.cogs.items(): + count = 0 + for cmd in getattr(cog, "application_commands", []): + if settings.guild_id and not cmd.guild_ids_to_rollout: + cmd.guild_ids_to_rollout.add(settings.guild_id) + self.add_application_command(cmd, use_rollout=True) + count += 1 + log.info("Registered %d command(s) from cog %s", count, cog_name) + # get_all_application_commands() returns ALL commands (guild + global); + # get_application_commands() only returns globals so it's misleading here. + log.info( + "Total application commands on bot: %d", + len(list(self.get_all_application_commands())), + ) + + async def close(self) -> None: + log.info("Shutting down — closing API session…") + await self.api.close() + await super().close() + + +def main() -> None: + if not settings.discord_token: + log.error("DISCORD_TOKEN is not set. Copy .env.example to .env and fill it in.") + sys.exit(1) + + async def runner() -> None: + # IMPORTANT: construct the bot *inside* the running loop. Nextcord binds + # its gateway/heartbeat to the active event loop; building it before the + # loop exists causes a loop mismatch where the heartbeat is scheduled on + # a dead loop ("heartbeat blocked for N seconds" while the loop sits idle + # in selector.poll). Creating it here guarantees a single, correct loop. + bot = VulnForge() + loop = asyncio.get_running_loop() + + def _handle_signal(*_args: object) -> None: + log.info("Signal received — stopping bot…") + loop.create_task(bot.close()) + + # Graceful Ctrl+C / SIGTERM (e.g. `docker stop`) handling. + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, _handle_signal) + except NotImplementedError: + # add_signal_handler is unavailable on Windows; KeyboardInterrupt covers SIGINT. + pass + + # Nextcord's Bot is not an async context manager (unlike discord.py 2.x), + # so start it directly and ensure a clean close on exit. + try: + # Load extensions + register cog commands BEFORE connecting. Doing + # this in setup_hook didn't run in this Nextcord version. + await bot.setup() + await bot.start(settings.discord_token) + finally: + if not bot.is_closed(): + await bot.close() + + try: + asyncio.run(runner()) + except KeyboardInterrupt: + log.info("KeyboardInterrupt — exiting.") + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100755 index 0000000..580bc4c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +# ────────────────────────────────────────────────────────────── +# VulnForge — Python dependencies +# ────────────────────────────────────────────────────────────── +nextcord>=2.6.0 # Modern Discord API wrapper (modals, slash, views) +aiohttp>=3.9.0 # Async HTTP client for NVD / CISA feeds +aiosqlite>=0.20.0 # Async SQLite for config + user notes +python-dotenv>=1.0.1 # Load secrets from .env diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100755 index 0000000..5c42c41 --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""scripts/ — VulnForge utility scripts (CLI wrappers).""" diff --git a/scripts/enrich_report.py b/scripts/enrich_report.py new file mode 100755 index 0000000..d3257d8 --- /dev/null +++ b/scripts/enrich_report.py @@ -0,0 +1,49 @@ +""" +scripts/enrich_report.py — CLI wrapper for the markdown enricher. + +Usage: + python -m scripts.enrich_report [-o ] + +If `-o` is omitted, writes alongside the input as `.enriched.md`. +Prints a one-line summary of artifact counts to stdout on success. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# Allow running both as `python -m scripts.enrich_report` and as a plain +# script from the vulnforge/ folder. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from utils.enricher import enrich_markdown # noqa: E402 + + +def main() -> int: + p = argparse.ArgumentParser(description="Enrich a markdown report with clickable artifact links.") + p.add_argument("input", type=Path, help="Path to the input markdown file.") + p.add_argument( + "-o", "--output", type=Path, default=None, + help="Output path. Defaults to .enriched.md alongside the input.", + ) + args = p.parse_args() + + src: Path = args.input + if not src.is_file(): + print(f"error: {src} does not exist or is not a file", file=sys.stderr) + return 2 + + dst: Path = args.output or src.with_suffix(".enriched.md") + text = src.read_text(encoding="utf-8") + enriched, stats = enrich_markdown(text) + dst.write_text(enriched, encoding="utf-8") + + summary = ", ".join(f"{k}={v}" for k, v in stats.as_summary().items() if v) + print(f"wrote {dst} ({summary or 'no artifacts found'})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100755 index 0000000..2d1adba --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1 @@ +"""utils package — shared helpers for VulnForge (API client, embeds, helpers).""" diff --git a/utils/api.py b/utils/api.py new file mode 100755 index 0000000..3e2956b --- /dev/null +++ b/utils/api.py @@ -0,0 +1,329 @@ +""" +utils/api.py — async data layer for NVD CVE API v2 and CISA KEV feed. + +Responsibilities: + * Single shared aiohttp.ClientSession (created on bot startup, closed on shutdown). + * Rate-limit-friendly: on-disk TTL caching + retry/backoff for NVD 403/429/5xx. + * Normalises raw API JSON into simple dataclasses (CVERecord, KEVEntry) so the + cogs/embeds never touch the messy upstream schema directly. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import Any + +import aiohttp + +from config import settings +from utils.helpers import ( + JSONCache, + parse_nvd_datetime, + severity_from_score, +) + +log = logging.getLogger("vulnforge.api") + + +# ────────────────────────────────────────────────────────────── +# Data models +# ────────────────────────────────────────────────────────────── +@dataclass +class CVERecord: + """A normalised view of a single NVD CVE entry.""" + + cve_id: str + description: str + published: datetime | None + last_modified: datetime | None + cvss_score: float | None + cvss_severity: str + cvss_vector: str | None + cvss_version: str | None + cwes: list[str] = field(default_factory=list) + references: list[str] = field(default_factory=list) + status: str = "Unknown" + + @property + def nvd_url(self) -> str: + return f"{settings.nvd_cve_web}/{self.cve_id}" + + +@dataclass +class KEVEntry: + """A normalised CISA Known-Exploited-Vulnerability entry.""" + + cve_id: str + vendor: str + product: str + name: str + date_added: str + short_description: str + required_action: str + due_date: str + ransomware: str + + @property + def date_added_dt(self) -> datetime | None: + try: + return datetime.strptime(self.date_added, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + except (ValueError, TypeError): + return None + + +# ────────────────────────────────────────────────────────────── +# Client +# ────────────────────────────────────────────────────────────── +class SecurityAPI: + """Async client for NVD + CISA KEV with caching and retry/backoff.""" + + def __init__(self) -> None: + self._session: aiohttp.ClientSession | None = None + self._cache = JSONCache(settings.cache_dir) + # In-memory KEV index so per-CVE lookups don't reparse the whole feed. + self._kev_index: dict[str, KEVEntry] = {} + self._kev_loaded_at: datetime | None = None + self._kev_lock = asyncio.Lock() + + # ── lifecycle ──────────────────────────────────────────── + async def start(self) -> None: + if self._session is None or self._session.closed: + headers = {"User-Agent": settings.user_agent} + if settings.nvd_api_key: + headers["apiKey"] = settings.nvd_api_key + timeout = aiohttp.ClientTimeout(total=settings.http_timeout) + self._session = aiohttp.ClientSession(headers=headers, timeout=timeout) + log.info("aiohttp session started (NVD key: %s)", bool(settings.nvd_api_key)) + + async def close(self) -> None: + if self._session and not self._session.closed: + await self._session.close() + log.info("aiohttp session closed") + + def _require_session(self) -> aiohttp.ClientSession: + if self._session is None or self._session.closed: + raise RuntimeError("SecurityAPI session not started — call start() first.") + return self._session + + # ── low-level GET with retry/backoff ───────────────────── + async def _get_json( + self, url: str, params: dict[str, Any] | None = None + ) -> dict[str, Any] | None: + session = self._require_session() + delay = 2.0 + for attempt in range(1, settings.request_retries + 1): + try: + async with session.get(url, params=params) as resp: + if resp.status == 200: + return await resp.json() + if resp.status in (403, 429, 503, 502, 500): + log.warning( + "GET %s → %s (attempt %s/%s), backing off %.1fs", + url, resp.status, attempt, settings.request_retries, delay, + ) + await asyncio.sleep(delay) + delay *= 2 + continue + log.error("GET %s → unexpected status %s", url, resp.status) + return None + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + log.warning( + "GET %s failed (attempt %s/%s): %s", + url, attempt, settings.request_retries, exc, + ) + await asyncio.sleep(delay) + delay *= 2 + log.error("GET %s exhausted retries", url) + return None + + # ── NVD: single CVE ────────────────────────────────────── + async def fetch_cve(self, cve_id: str) -> CVERecord | None: + """Fetch a single CVE by id (cached for cve_cache_ttl).""" + cache_key = f"cve_{cve_id}" + cached = self._cache.get(cache_key, settings.cve_cache_ttl) + if cached is not None: + return self._parse_cve_item(cached) + + data = await self._get_json(settings.nvd_base_url, {"cveId": cve_id}) + if not data: + return None + vulns = data.get("vulnerabilities") or [] + if not vulns: + return None + item = vulns[0] + self._cache.set(cache_key, item) + return self._parse_cve_item(item) + + # ── NVD: recent CVEs (last N days) ─────────────────────── + async def fetch_recent_cves(self, days: int = 7) -> list[CVERecord]: + """ + Fetch CVEs published in the last `days` days. + + NVD limits the pubStartDate/pubEndDate window to 120 days and paginates + with resultsPerPage / startIndex. We page through everything in the + window and cache the assembled list. + """ + cache_key = f"recent_{days}d" + cached = self._cache.get(cache_key, settings.cve_cache_ttl) + if cached is not None: + return [self._parse_cve_item(i) for i in cached] + + end = datetime.now(timezone.utc) + start = end - timedelta(days=days) + iso = "%Y-%m-%dT%H:%M:%S.000" + results_per_page = 2000 + start_index = 0 + collected: list[dict[str, Any]] = [] + + while True: + params = { + "pubStartDate": start.strftime(iso), + "pubEndDate": end.strftime(iso), + "resultsPerPage": results_per_page, + "startIndex": start_index, + } + data = await self._get_json(settings.nvd_base_url, params) + if not data: + break + vulns = data.get("vulnerabilities") or [] + collected.extend(vulns) + total = int(data.get("totalResults", 0)) + start_index += results_per_page + if start_index >= total or not vulns: + break + # Be gentle on the API between pages. + await asyncio.sleep(0.7 if settings.nvd_api_key else 6.0) + + self._cache.set(cache_key, collected) + log.info("Fetched %s CVEs from the last %s days", len(collected), days) + return [self._parse_cve_item(i) for i in collected] + + # ── CISA KEV feed ──────────────────────────────────────── + async def _load_kev(self, force: bool = False) -> dict[str, KEVEntry]: + """Load + index the KEV feed (cached on disk + in memory).""" + async with self._kev_lock: + fresh = ( + self._kev_index + and self._kev_loaded_at + and datetime.now(timezone.utc) - self._kev_loaded_at + < timedelta(seconds=settings.kev_cache_ttl) + ) + if fresh and not force: + return self._kev_index + + payload = None if force else self._cache.get("kev_feed", settings.kev_cache_ttl) + if payload is None: + payload = await self._get_json(settings.cisa_kev_url) + if payload: + self._cache.set("kev_feed", payload) + + index: dict[str, KEVEntry] = {} + for v in (payload or {}).get("vulnerabilities", []): + entry = KEVEntry( + cve_id=v.get("cveID", "").upper(), + vendor=v.get("vendorProject", "Unknown"), + product=v.get("product", "Unknown"), + name=v.get("vulnerabilityName", ""), + date_added=v.get("dateAdded", ""), + short_description=v.get("shortDescription", ""), + required_action=v.get("requiredAction", ""), + due_date=v.get("dueDate", ""), + ransomware=v.get("knownRansomwareCampaignUse", "Unknown"), + ) + if entry.cve_id: + index[entry.cve_id] = entry + + self._kev_index = index + self._kev_loaded_at = datetime.now(timezone.utc) + log.info("KEV feed loaded: %s entries", len(index)) + return index + + async def get_kev_entry(self, cve_id: str) -> KEVEntry | None: + index = await self._load_kev() + return index.get(cve_id.upper()) + + async def is_in_kev(self, cve_id: str) -> bool: + return (await self.get_kev_entry(cve_id)) is not None + + async def latest_kev(self, limit: int = 10) -> list[KEVEntry]: + """Return the most recently added KEV entries (newest first).""" + index = await self._load_kev() + entries = sorted( + index.values(), + key=lambda e: e.date_added or "", + reverse=True, + ) + return entries[:limit] + + async def kev_added_since(self, days: int = 7) -> list[KEVEntry]: + """KEV entries whose dateAdded falls within the last `days` days.""" + index = await self._load_kev() + cutoff = datetime.now(timezone.utc) - timedelta(days=days) + out: list[KEVEntry] = [] + for e in index.values(): + dt = e.date_added_dt + if dt and dt >= cutoff: + out.append(e) + out.sort(key=lambda e: e.date_added or "", reverse=True) + return out + + # ── parsing ────────────────────────────────────────────── + @staticmethod + def _parse_cve_item(item: dict[str, Any]) -> CVERecord: + """Convert a raw NVD `vulnerabilities[]` item into a CVERecord.""" + cve = item.get("cve", {}) + cve_id = cve.get("id", "UNKNOWN") + + # Description: prefer English. + description = "" + for d in cve.get("descriptions", []): + if d.get("lang") == "en": + description = d.get("value", "") + break + if not description and cve.get("descriptions"): + description = cve["descriptions"][0].get("value", "") + + # CVSS: prefer v3.1 → v3.0 → v2 (highest available). + score = severity = vector = version = None + metrics = cve.get("metrics", {}) + for key in ("cvssMetricV31", "cvssMetricV30", "cvssMetricV2"): + arr = metrics.get(key) + if arr: + data = arr[0].get("cvssData", {}) + score = data.get("baseScore") + vector = data.get("vectorString") + version = data.get("version") + severity = ( + data.get("baseSeverity") + or arr[0].get("baseSeverity") + or severity_from_score(score) + ) + break + + cwes: list[str] = [] + for weakness in cve.get("weaknesses", []): + for desc in weakness.get("description", []): + val = desc.get("value") + if val and val.upper() != "NVD-CWE-NOINFO" and val not in cwes: + cwes.append(val) + + refs = [r.get("url") for r in cve.get("references", []) if r.get("url")] + + return CVERecord( + cve_id=cve_id, + description=description or "No description available.", + published=parse_nvd_datetime(cve.get("published")), + last_modified=parse_nvd_datetime(cve.get("lastModified")), + cvss_score=score, + cvss_severity=(severity or severity_from_score(score)).upper(), + cvss_vector=vector, + cvss_version=version, + cwes=cwes, + references=refs, + status=cve.get("vulnStatus", "Unknown"), + ) diff --git a/utils/embeds.py b/utils/embeds.py new file mode 100755 index 0000000..b2fe31f --- /dev/null +++ b/utils/embeds.py @@ -0,0 +1,293 @@ +""" +utils/embeds.py — cyberpunk / hacker-aesthetic embed builders. + +Every embed in VulnForge is built here so styling stays consistent: +dark backdrops, neon severity colours, monospace code blocks, and tidy +mobile + desktop formatting. Functions accept normalised dataclasses from +utils.api and return ready-to-send nextcord.Embed objects. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import nextcord + +from config import settings +from utils.api import CVERecord, KEVEntry +from utils.enricher import cve_links +from utils.helpers import ( + discord_timestamp, + severity_color, + severity_emoji, + truncate, +) + +# Visual furniture reused across embeds. +DIVIDER = "▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰▰" +BRAND = "VulnForge // Security Intelligence" +BRAND_ICON = "https://www.cisa.gov/profiles/cisad8_gov/themes/custom/gesso/logo.svg" + + +def cve_pivots(cve_id: str, *, compact: bool = False) -> str: + """One-line markdown pivot links for a CVE — PoC search / exploit-db / etc. + + Used to enrich every CVE-bearing embed (cve, weekly, kev) with quick links + to GitHub PoCs, exploit-db, Nuclei templates, trickest/cve and Vulncheck. + Pass `compact=True` for tighter rendering in list-style embeds. + """ + links = {name: url for name, url in cve_links(cve_id)} + if compact: + order = ("GitHub PoC search", "exploit-db", "Nuclei templates") + labels = {"GitHub PoC search": "PoC", "exploit-db": "ExDB", "Nuclei templates": "Nuclei"} + else: + order = ("GitHub PoC search", "exploit-db", "Nuclei templates", + "trickest/cve", "Vulncheck", "GitHub Advisory") + labels = { + "GitHub PoC search": "GH PoC", + "exploit-db": "Exploit-DB", + "Nuclei templates": "Nuclei", + "trickest/cve": "trickest/cve", + "Vulncheck": "Vulncheck", + "GitHub Advisory": "GH Advisory", + } + return " · ".join(f"[{labels[k]}]({links[k]})" for k in order if k in links) + + +def _footer(embed: nextcord.Embed) -> nextcord.Embed: + embed.set_footer(text=f"⛓ {BRAND}") + embed.timestamp = datetime.now(timezone.utc) + return embed + + +# ────────────────────────────────────────────────────────────── +# CVE detail embed +# ────────────────────────────────────────────────────────────── +def build_cve_embed( + cve: CVERecord, + kev: KEVEntry | None, + *, + show_details: bool = True, + show_kev: bool = True, + show_references: bool = True, + show_cwe: bool = True, + notes: list[dict] | None = None, +) -> nextcord.Embed: + """Build the rich, option-driven CVE embed returned by /cve.""" + sev = cve.cvss_severity + color = settings.color_kev if kev else severity_color(sev) + exploited = " ☣ ACTIVELY EXPLOITED" if kev else "" + + embed = nextcord.Embed( + title=f"{severity_emoji(sev)} {cve.cve_id}{exploited}", + url=cve.nvd_url, + description=truncate(cve.description, 1024) if show_details else None, + color=color, + ) + + # Score line — always shown; it's the headline metric. + score_str = f"`{cve.cvss_score}`" if cve.cvss_score is not None else "`N/A`" + ver = f" (CVSS v{cve.cvss_version})" if cve.cvss_version else "" + embed.add_field( + name="◈ Severity", + value=f"{severity_emoji(sev)} **{sev}**{ver}", + inline=True, + ) + embed.add_field(name="◈ Base Score", value=score_str, inline=True) + embed.add_field(name="◈ Status", value=f"`{cve.status}`", inline=True) + + if show_details: + embed.add_field( + name="◈ Published", + value=discord_timestamp(cve.published, "D"), + inline=True, + ) + embed.add_field( + name="◈ Modified", + value=discord_timestamp(cve.last_modified, "R"), + inline=True, + ) + if cve.cvss_vector: + embed.add_field( + name="◈ Vector", + value=f"```{truncate(cve.cvss_vector, 240)}```", + inline=False, + ) + + if show_cwe and cve.cwes: + embed.add_field( + name="◈ Weaknesses (CWE)", + value=" ".join(f"`{c}`" for c in cve.cwes[:8]), + inline=False, + ) + + # KEV block — the most security-critical signal. + if show_kev: + if kev: + ransom = kev.ransomware if kev.ransomware else "Unknown" + embed.add_field(name=DIVIDER, value="**☣ CISA KEV — ACTIVE EXPLOITATION**", inline=False) + embed.add_field(name="▸ Vendor / Product", value=f"`{kev.vendor}` / `{kev.product}`", inline=True) + embed.add_field(name="▸ Added to KEV", value=discord_timestamp(kev.date_added_dt, "D"), inline=True) + embed.add_field(name="▸ Remediation Due", value=f"`{kev.due_date or 'N/A'}`", inline=True) + embed.add_field(name="▸ Ransomware Use", value=f"`{ransom}`", inline=True) + if kev.required_action: + embed.add_field( + name="▸ Required Action", + value=truncate(kev.required_action, 1024), + inline=False, + ) + else: + embed.add_field( + name="◈ CISA KEV Status", + value="✅ Not currently listed in the Known Exploited Vulnerabilities catalog.", + inline=False, + ) + + if show_references and cve.references: + links = "\n".join(f"• [{truncate(u, 70)}]({u})" for u in cve.references[:6]) + embed.add_field(name="◈ References", value=truncate(links, 1024), inline=False) + + if notes: + rendered = "\n".join( + f"• <@{n['user_id']}>: {truncate(n['note'], 150)}" for n in notes[:5] + ) + embed.add_field(name="◈ Researcher Notes", value=truncate(rendered, 1024), inline=False) + + pivots = cve_pivots(cve.cve_id) + if pivots: + embed.add_field(name="◈ Recon Pivots", value=pivots, inline=False) + + embed.add_field( + name="◈ Links", + value=f"[🛰 NVD]({cve.nvd_url}) • [🛡 CISA KEV]({settings.cisa_kev_web})", + inline=False, + ) + return _footer(embed) + + +# ────────────────────────────────────────────────────────────── +# Weekly digest embed +# ────────────────────────────────────────────────────────────── +def build_weekly_embed( + *, + total_new: int, + in_kev_count: int, + newly_added_kev: list[KEVEntry], + top_cves: list[CVERecord], + days: int = 7, +) -> nextcord.Embed: + """Build the automated weekly threat-intelligence digest embed.""" + embed = nextcord.Embed( + title="☣ VULNFORGE WEEKLY THREAT DIGEST", + description=( + f"```ansi\n\u001b[1;31m> Intelligence sweep — last {days} days\u001b[0m\n```\n" + f"Aggregated from **NVD** and the **CISA KEV** catalog." + ), + color=settings.color_critical if newly_added_kev else settings.color_info, + url=settings.cisa_kev_web, + ) + + # Headline metrics block. + embed.add_field(name="🆕 New CVEs", value=f"**{total_new:,}**", inline=True) + embed.add_field(name="☣ Already in KEV", value=f"**{in_kev_count:,}**", inline=True) + embed.add_field(name="⚠ Newly Added to KEV", value=f"**{len(newly_added_kev):,}**", inline=True) + + embed.add_field(name=DIVIDER, value="**▶ TOP CRITICAL VULNERABILITIES**", inline=False) + + if top_cves: + for cve in top_cves: + sev = cve.cvss_severity + score = cve.cvss_score if cve.cvss_score is not None else "N/A" + pivots = cve_pivots(cve.cve_id, compact=True) + link_line = f"[NVD]({cve.nvd_url})" + if pivots: + link_line += f" · {pivots}" + embed.add_field( + name=f"{severity_emoji(sev)} {cve.cve_id} — {score}", + value=f"{truncate(cve.description, 140)}\n{link_line}", + inline=False, + ) + else: + embed.add_field( + name="—", + value="No scored CVEs found in the window.", + inline=False, + ) + + if newly_added_kev: + kev_lines = "\n".join( + f"• `{e.cve_id}` — {truncate(e.product, 38)} ({e.vendor}) — {cve_pivots(e.cve_id, compact=True)}" + for e in newly_added_kev[:8] + ) + embed.add_field( + name="☣ Newly Weaponized (Added to KEV)", + value=truncate(kev_lines, 1024), + inline=False, + ) + + embed.add_field( + name="◈ Sources", + value=( + "[🛰 NVD Feed](https://nvd.nist.gov/vuln) • " + f"[🛡 CISA KEV]({settings.cisa_kev_web})" + ), + inline=False, + ) + return _footer(embed) + + +# ────────────────────────────────────────────────────────────── +# KEV list embed +# ────────────────────────────────────────────────────────────── +def build_kev_list_embed(entries: list[KEVEntry], limit: int) -> nextcord.Embed: + """Build the embed for the /kev command (latest KEV additions).""" + embed = nextcord.Embed( + title="☣ CISA KEV — LATEST KNOWN EXPLOITED VULNERABILITIES", + description=( + "```ansi\n\u001b[1;31m> Actively exploited in the wild\u001b[0m\n```" + f"Showing the **{min(limit, len(entries))}** most recent additions." + ), + color=settings.color_kev, + url=settings.cisa_kev_web, + ) + for e in entries[:limit]: + pivots = cve_pivots(e.cve_id, compact=True) + link_line = f"[NVD]({settings.nvd_cve_web}/{e.cve_id})" + if pivots: + link_line += f" · {pivots}" + embed.add_field( + name=f"🩸 {e.cve_id} — {truncate(e.product, 50)}", + value=truncate( + f"**Vendor:** {e.vendor}\n" + f"**Added:** {discord_timestamp(e.date_added_dt, 'D')} • " + f"**Ransomware:** `{e.ransomware}`\n" + f"{truncate(e.short_description, 140)}\n" + f"{link_line}", + 1024, + ), + inline=False, + ) + if not entries: + embed.description = "No KEV entries could be retrieved right now." + return _footer(embed) + + +# ────────────────────────────────────────────────────────────── +# Error / utility embeds +# ────────────────────────────────────────────────────────────── +def build_error_embed(title: str, message: str) -> nextcord.Embed: + embed = nextcord.Embed( + title=f"⛔ {title}", + description=message, + color=settings.color_critical, + ) + return _footer(embed) + + +def build_info_embed(title: str, message: str) -> nextcord.Embed: + embed = nextcord.Embed( + title=f"✅ {title}", + description=message, + color=settings.color_success, + ) + return _footer(embed) \ No newline at end of file diff --git a/utils/enricher.py b/utils/enricher.py new file mode 100755 index 0000000..aa22fd5 --- /dev/null +++ b/utils/enricher.py @@ -0,0 +1,728 @@ +""" +utils/enricher.py — Markdown artifact enrichment. + +Scans a markdown report for security-relevant artifacts (CVEs, IPv4 addresses, +hostnames, emails, GitHub repos / commits, raw URLs, vendor product names) and: + + 1. Rewrites each occurrence inline as a clickable markdown link pointing at + the most useful primary lookup (NVD for CVEs, Shodan for IPs, github.com + for repos, etc.). + 2. Appends an "INTEL APPENDIX" section grouping every unique artifact with + the FULL set of enrichment links (Shodan + Censys + VirusTotal for IPs, + NVD + GitHub PoC search + exploit-db + nuclei templates for CVEs, etc.). + +Code fences, existing markdown links and inline code spans are preserved +unmodified so the original report stays readable. + +Public API: + enrich_markdown(text: str) -> tuple[str, Stats] +""" + +from __future__ import annotations + +import ipaddress +import re +from dataclasses import dataclass, field +from typing import Callable +from urllib.parse import quote, quote_plus + +# ────────────────────────────────────────────────────────────────────────── +# Regex patterns +# ────────────────────────────────────────────────────────────────────────── +# CVE: CVE-YYYY-NNNN(N+) — case insensitive, normalised to uppercase. +CVE_RE = re.compile(r"\bCVE-\d{4}-\d{4,7}\b", re.IGNORECASE) + +# IPv4 with optional CIDR; we filter private/invalid in post-processing. +IPV4_RE = re.compile( + r"\b(?:25[0-5]|2[0-4]\d|[01]?\d\d?)" + r"(?:\.(?:25[0-5]|2[0-4]\d|[01]?\d\d?)){3}" + r"(?:/\d{1,2})?\b" +) + +# Emails — RFC-ish, good enough for report text. +EMAIL_RE = re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]{2,}\b") + +# Raw http(s) URLs — used both for "wrap as link" and for hostname extraction. +URL_RE = re.compile(r"https?://[^\s<>()\[\]\"']+", re.IGNORECASE) + +# GitHub repo refs like org/repo, optionally followed by @sha or #commit-sha. +# Restrictive: 1st segment 1-39 chars (GitHub login limit), 2nd 1-100 chars, +# both must contain at least one letter to avoid matching numeric paths. +GH_REPO_RE = re.compile( + r"(? list[tuple[str, str]]: + """Return (label, url) tuples for a CVE id.""" + q = quote_plus(cve) + return [ + ("NVD", f"https://nvd.nist.gov/vuln/detail/{cve}"), + ("MITRE", f"https://cve.mitre.org/cgi-bin/cvename.cgi?name={cve}"), + ("GitHub Advisory", f"https://github.com/advisories?query={q}"), + ("GitHub PoC search", f"https://github.com/search?q={q}&type=repositories"), + ("exploit-db", f"https://www.exploit-db.com/search?cve={cve.removeprefix('CVE-')}"), + ("Nuclei templates", f"https://github.com/projectdiscovery/nuclei-templates/search?q={q}"), + ("trickest/cve", f"https://github.com/trickest/cve/tree/main/{cve.split('-')[1]}"), + ("Vulncheck", f"https://vulncheck.com/browse/cve/{cve}"), + ] + + +def ipv4_links(ip: str) -> list[tuple[str, str]]: + """Return (label, url) tuples for an IPv4 address (no CIDR suffix).""" + return [ + ("Shodan", f"https://www.shodan.io/host/{ip}"), + ("Censys", f"https://search.censys.io/hosts/{ip}"), + ("VirusTotal", f"https://www.virustotal.com/gui/ip-address/{ip}"), + ("GreyNoise", f"https://viz.greynoise.io/ip/{ip}"), + ("AbuseIPDB", f"https://www.abuseipdb.com/check/{ip}"), + ("ipinfo.io", f"https://ipinfo.io/{ip}"), + ] + + +def hostname_links(host: str) -> list[tuple[str, str]]: + """Return (label, url) tuples for a hostname / FQDN.""" + host_q = quote(host) + return [ + ("Open", f"https://{host}"), + ("crt.sh", f"https://crt.sh/?q=%25.{host_q}"), + ("urlscan.io", f"https://urlscan.io/domain/{host_q}"), + ("Shodan", f"https://www.shodan.io/search?query=hostname%3A{host_q}"), + ("VirusTotal", f"https://www.virustotal.com/gui/domain/{host_q}"), + ("SecurityTrails", f"https://securitytrails.com/domain/{host_q}/dns"), + ("Wayback", f"https://web.archive.org/web/*/{host_q}/*"), + ] + + +def email_links(email: str) -> list[tuple[str, str]]: + """Return (label, url) tuples for an email address.""" + local, _, domain = email.partition("@") + return [ + ("mailto", f"mailto:{email}"), + ("HIBP", f"https://haveibeenpwned.com/account/{quote(email)}"), + ("Hunter (domain)", f"https://hunter.io/search/{quote(domain)}"), + ("Epieos", f"https://epieos.com/?q={quote(email)}"), + ] + + +def github_links(org: str, repo: str, sha: str | None) -> list[tuple[str, str]]: + """Return (label, url) tuples for a GitHub repo, optionally at a commit.""" + base = f"https://github.com/{org}/{repo}" + out = [ + ("Repo", base), + ("Commits", f"{base}/commits"), + ("Issues", f"{base}/issues"), + ("Secret scan (TruffleHog)", "https://github.com/trufflesecurity/trufflehog"), + ("GH advisories", f"https://github.com/{org}/{repo}/security/advisories"), + ] + if sha: + out.insert(1, (f"Commit {sha[:7]}", f"{base}/commit/{sha}")) + return out + + +# ────────────────────────────────────────────────────────────────────────── +# Stats container +# ────────────────────────────────────────────────────────────────────────── +@dataclass +class Stats: + """Counts + unique artifact sets for the enrichment appendix.""" + cves: set[str] = field(default_factory=set) + ipv4s: set[str] = field(default_factory=set) + cidrs: set[str] = field(default_factory=set) + hostnames: set[str] = field(default_factory=set) + emails: set[str] = field(default_factory=set) + repos: set[tuple[str, str]] = field(default_factory=set) # (org, repo) + commits: set[tuple[str, str, str]] = field(default_factory=set) # (org, repo, sha) + urls: set[str] = field(default_factory=set) + vendors: set[str] = field(default_factory=set) + + def total(self) -> int: + return ( + len(self.cves) + len(self.ipv4s) + len(self.cidrs) + len(self.hostnames) + + len(self.emails) + len(self.repos) + len(self.commits) + len(self.urls) + + len(self.vendors) + ) + + def as_summary(self) -> dict[str, int]: + return { + "CVEs": len(self.cves), + "IPv4s": len(self.ipv4s), + "CIDRs": len(self.cidrs), + "Hostnames": len(self.hostnames), + "Emails": len(self.emails), + "GitHub repos": len(self.repos), + "Git commits": len(self.commits), + "URLs": len(self.urls), + "Vendors": len(self.vendors), + } + + +# ────────────────────────────────────────────────────────────────────────── +# Protect/restore code spans + existing links +# ────────────────────────────────────────────────────────────────────────── +_PLACEHOLDER = "\x00ENRICH{}\x00" + + +def _stash(text: str, stats: Stats) -> tuple[str, list[str]]: + """Replace code fences / inline code / existing links with placeholders. + + Returns (modified_text, stash) where placeholders look like \\x00ENRICH\\x00 + so the rewrite passes won't touch them. + + Inline backticks (`x`) get *special* handling: if `x` is a single, clean + artifact (CVE, IP, IPv4/CIDR, email, FQDN, org/repo, URL), the WHOLE + backtick span is replaced with `[\`x\`](url)` so the artifact is clickable + while still rendering as code. Otherwise the span is stashed unchanged. + """ + stash: list[str] = [] + + def stash_repl(m: re.Match[str]) -> str: + stash.append(m.group(0)) + return _PLACEHOLDER.format(len(stash) - 1) + + def inline_code_repl(m: re.Match[str]) -> str: + raw = m.group(0) # e.g. "`devcolor/codebenders-datathon`" + inner = raw[1:-1].strip() # strip the backticks + link = _single_artifact_link(inner, stats) + if link is None: + # Plain code — keep the original span out of further passes. + stash.append(raw) + return _PLACEHOLDER.format(len(stash) - 1) + # Replace with a markdown link wrapping the code, then stash the + # whole thing so subsequent passes don't try to re-match the inner. + replacement = f"[{raw}]({link})" + stash.append(replacement) + return _PLACEHOLDER.format(len(stash) - 1) + + text = FENCED_BLOCK_RE.sub(stash_repl, text) + text = INLINE_CODE_RE.sub(inline_code_repl, text) + text = EXISTING_LINK_RE.sub(stash_repl, text) + return text, stash + + +def _single_artifact_link(content: str, stats: Stats) -> str | None: + """If `content` is exactly one recognisable artifact, return its primary + enrichment URL and update stats. Otherwise return None. + """ + s = content.strip() + if not s: + return None + + # CVE + if CVE_RE.fullmatch(s): + cve = s.upper() + stats.cves.add(cve) + return f"https://nvd.nist.gov/vuln/detail/{cve}" + + # URL + if URL_RE.fullmatch(s): + url = s.rstrip(".,;:)") + stats.urls.add(url) + host_match = re.match(r"https?://([^/:]+)", url, re.IGNORECASE) + if host_match and _looks_like_hostname(host_match.group(1).lower()): + stats.hostnames.add(host_match.group(1).lower()) + return url + + # Email + if EMAIL_RE.fullmatch(s): + stats.emails.add(s) + return f"mailto:{s}" + + # GitHub org/repo (optionally @sha) — full-string match is strong signal + # that the user intended it as a repo ref. + gh = GH_REPO_RE.fullmatch(s) + if gh: + org, repo, sha = gh.group(1), gh.group(2), gh.group(3) + if _looks_like_repo(org, repo) and "." not in org: + stats.repos.add((org, repo)) + base = f"https://github.com/{org}/{repo}" + if sha: + stats.commits.add((org, repo, sha)) + return f"{base}/commit/{sha}" + return base + + # IPv4 (with optional CIDR) + if IPV4_RE.fullmatch(s): + ip, _, cidr = s.partition("/") + if cidr: + stats.cidrs.add(s) + return f"https://www.shodan.io/search?query=net%3A{quote(s)}" + if _is_public_ipv4(ip): + stats.ipv4s.add(ip) + return f"https://www.shodan.io/host/{ip}" + return None + + # Hostname + if HOSTNAME_RE.fullmatch(s) and _looks_like_hostname(s): + host = s.lower() + stats.hostnames.add(host) + return f"https://{host}" + + return None + + +def _restore(text: str, stash: list[str]) -> str: + for i, original in enumerate(stash): + text = text.replace(_PLACEHOLDER.format(i), original) + return text + + +# ────────────────────────────────────────────────────────────────────────── +# Hostname filtering — skip noisy false positives +# ────────────────────────────────────────────────────────────────────────── +_HOST_SKIP_TLDS = { + # Common file extensions that look like TLDs in our regex. + "md", "txt", "log", "csv", "json", "yml", "yaml", "ini", "cfg", + "py", "js", "ts", "go", "rs", "java", "html", "css", "xml", "sh", + "ps1", "bat", "exe", "dll", "so", "dylib", "png", "jpg", "jpeg", + "gif", "svg", "pdf", "zip", "tar", "gz", "rar", "7z", "iso", + "bak", "old", "tmp", "swp", "db", "sql", + # Common code/JS identifiers mis-detected as TLDs. + "redirect", "then", "catch", "call", "apply", "bind", + "prototype", "constructor", "length", "value", "name", +} + +_REPO_SKIP_SUFFIXES = ( + ".py", ".js", ".ts", ".java", ".rb", ".go", ".rs", ".c", ".cpp", + ".h", ".hpp", ".cs", ".php", ".html", ".css", ".json", ".yml", + ".yaml", ".xml", ".sh", ".md", ".txt", ".log", ".sql", ".bak", +) + + +# Recognised gTLDs / common ccTLDs. Anything outside this set OR a 2-letter +# all-alpha ccTLD gets rejected as not-a-real-hostname (catches JS method +# names like jquery.parseHtml, internal aliases like kctcsfin.ps, etc.). +_KNOWN_TLDS = { + # gTLDs + "com", "net", "org", "edu", "gov", "mil", "int", "info", "biz", + "name", "pro", "mobi", "aero", "coop", "museum", + # Modern / brand + "io", "co", "dev", "app", "ai", "cloud", "tech", "xyz", "online", + "site", "store", "shop", "blog", "page", "wiki", "news", "today", + "news", "live", "world", "group", "club", "fun", "studio", "design", + "art", "video", "tv", "fm", "gg", "ly", "sh", "me", "to", "vc", + "cc", "ws", "asia", "jobs", "travel", "post", "tel", "cat", + # Software / infra + "software", "services", "agency", "network", "social", "life", + "academy", "consulting", "solutions", "systems", "support", "ventures", + "finance", "money", "security", "wtf", +} + + +def _looks_like_hostname(host: str) -> bool: + """Heuristic: must look like a real internet hostname, not a code + identifier (`jquery.parseHtml`) or internal alias (`kctcsfin.ps`). + """ + host = host.strip(".") + if "." not in host: + return False + parts = host.split(".") + if any(p == "" for p in parts): + return False + tld = parts[-1].lower() + if tld in _HOST_SKIP_TLDS: + return False + if not tld.isalpha(): + return False + # TLD must be either a known gTLD/brand OR a 2-letter all-letter ccTLD. + if not (tld in _KNOWN_TLDS or len(tld) == 2): + return False + # Reject if any segment has internal mixed case (camelCase identifier). + for seg in parts: + if re.search(r"[a-z][A-Z]", seg): + return False + return True + + +def _looks_like_repo(org: str, repo: str) -> bool: + """Heuristic for GH org/repo refs (used by both inline and backtick passes).""" + if not (re.search(r"[A-Za-z]", org) and re.search(r"[A-Za-z]", repo)): + return False + if org.isdigit() or repo.isdigit(): + return False + rl = repo.lower() + if any(rl.endswith(suf) for suf in _REPO_SKIP_SUFFIXES): + return False + return True + + +def _is_public_ipv4(ip: str) -> bool: + """True if the IP is globally routable (skip RFC1918, loopback, etc.).""" + try: + addr = ipaddress.IPv4Address(ip) + except (ipaddress.AddressValueError, ValueError): + return False + return not (addr.is_private or addr.is_loopback or addr.is_link_local + or addr.is_multicast or addr.is_unspecified or addr.is_reserved) + + +# ────────────────────────────────────────────────────────────────────────── +# Inline rewrite passes +# ────────────────────────────────────────────────────────────────────────── +def _rewrite_cves(text: str, stats: Stats) -> str: + def repl(m: re.Match[str]) -> str: + cve = m.group(0).upper() + stats.cves.add(cve) + return f"[{cve}](https://nvd.nist.gov/vuln/detail/{cve})" + return CVE_RE.sub(repl, text) + + +def _rewrite_ipv4s(text: str, stats: Stats) -> str: + def repl(m: re.Match[str]) -> str: + raw = m.group(0) + ip, _, cidr = raw.partition("/") + if cidr: + stats.cidrs.add(raw) + # Link CIDR to Shodan net: query. + return f"[{raw}](https://www.shodan.io/search?query=net%3A{quote(raw)})" + if not _is_public_ipv4(ip): + return raw + stats.ipv4s.add(ip) + return f"[{ip}](https://www.shodan.io/host/{ip})" + return IPV4_RE.sub(repl, text) + + +def _rewrite_emails(text: str, stats: Stats) -> str: + def repl(m: re.Match[str]) -> str: + email = m.group(0) + stats.emails.add(email) + return f"[{email}](mailto:{email})" + return EMAIL_RE.sub(repl, text) + + +def _rewrite_urls(text: str, stats: Stats) -> str: + def repl(m: re.Match[str]) -> str: + url = m.group(0).rstrip(".,;:)") + trailing = m.group(0)[len(url):] + stats.urls.add(url) + # Extract hostname for the appendix too. + host_match = re.match(r"https?://([^/:]+)", url, re.IGNORECASE) + if host_match: + host = host_match.group(1).lower() + if _looks_like_hostname(host): + stats.hostnames.add(host) + return f"[{url}]({url}){trailing}" + return URL_RE.sub(repl, text) + + +def _rewrite_hostnames(text: str, stats: Stats) -> str: + def repl(m: re.Match[str]) -> str: + host = m.group(0).lower() + if not _looks_like_hostname(host): + return m.group(0) + stats.hostnames.add(host) + return f"[{host}](https://{host})" + return HOSTNAME_RE.sub(repl, text) + + +def _rewrite_github(text: str, stats: Stats) -> str: + """Rewrite `org/repo` and `org/repo@sha` refs as GitHub links. + + `org/repo` is an extremely ambiguous shape ("voice/data", "OWA/Exchange", + "443/tcp", "Dev/test"). To avoid false positives we require ALL of: + - both segments contain at least one letter + - neither segment is a known protocol/port-style stopword + - both segments contain a hyphen (a strong GitHub-style convention), AND + - a context word (github / git history / repo / commit) appears within + 80 chars before the match. + Single-hyphen / no-hyphen refs that are genuine repos (e.g. `shadcn/ui` or + `devcolor/codebenders-datathon`) are expected to appear in backticks in + the source; those are already handled by the inline-code pass in _stash(). + """ + org_skip = {"http", "https", "ftp", "mailto", "tcp", "udp", "icmp"} + + def repl(m: re.Match[str]) -> str: + org, repo, sha = m.group(1), m.group(2), m.group(3) + if org.lower() in org_skip or repo.lower() in org_skip: + return m.group(0) + if not _looks_like_repo(org, repo): + return m.group(0) + if not ("-" in org and "-" in repo): + return m.group(0) + start = m.start() + window = text[max(0, start - 80):start].lower() + if not re.search(r"github|git history|\bgit \b|\brepo\b|\bcommit\b", window): + return m.group(0) + stats.repos.add((org, repo)) + base = f"https://github.com/{org}/{repo}" + if sha: + stats.commits.add((org, repo, sha)) + return f"[{org}/{repo}@{sha[:7]}]({base}/commit/{sha})" + return f"[{org}/{repo}]({base})" + return GH_REPO_RE.sub(repl, text) + + +def _rewrite_vendors(text: str, stats: Stats) -> str: + """Wrap vendor/product mentions with their docs link (first occurrence only). + + To avoid noisy repetition, we only link the FIRST hit per vendor in the + document; subsequent occurrences stay plain text. + """ + seen: set[str] = set() + for pat, display, url in _VENDOR_PATTERNS: + if display in seen: + continue + # Only the first occurrence: use sub with count=1 and skip if already done. + def repl(m: re.Match[str], _disp=display, _url=url) -> str: + stats.vendors.add(_disp) + seen.add(_disp) + return f"[{m.group(0)}]({_url})" + text, n = pat.subn(repl, text, count=1) + if n: + seen.add(display) + return text + + +# ────────────────────────────────────────────────────────────────────────── +# Appendix builder +# ────────────────────────────────────────────────────────────────────────── +def _link_orphan_commits(text: str, stats: Stats) -> str: + """After other passes have run, find lines that contain a GitHub repo + link plus one or more unlinked backticked SHAs, and link those SHAs to + that repo's commit URL. + + Matches both `[\`abc1234\`]` (already wrapped via single-artifact pass that + failed because no repo context) and a bare ` `abc1234` ` code span. + """ + repo_link_re = re.compile(r"https?://github\.com/([\w.-]+)/([\w.-]+?)(?:[)/#?]|$)") + sha_code_re = re.compile(r"`([0-9a-f]{7,40})`") + + out_lines: list[str] = [] + for line in text.splitlines(keepends=True): + m = repo_link_re.search(line) + if not m: + out_lines.append(line) + continue + org, repo = m.group(1), m.group(2) + + def repl(sm: re.Match[str], _org=org, _repo=repo) -> str: + sha = sm.group(1) + stats.commits.add((_org, _repo, sha)) + return f"[`{sha}`](https://github.com/{_org}/{_repo}/commit/{sha})" + + out_lines.append(sha_code_re.sub(repl, line)) + return "".join(out_lines) + + +def _format_link_row(label: str, links: list[tuple[str, str]]) -> str: + bits = " · ".join(f"[{name}]({url})" for name, url in links) + return f"- **{label}** — {bits}" + + +def _build_appendix(stats: Stats) -> str: + if stats.total() == 0: + return "" + + lines: list[str] = [] + lines.append("\n---\n") + lines.append("## INTEL APPENDIX — Artifact Enrichment\n") + lines.append( + "_Auto-generated by VulnForge enricher. Each artifact below is paired with " + "external recon / threat-intel lookups for rapid pivot._\n" + ) + + if stats.cves: + lines.append("### CVEs\n") + for cve in sorted(stats.cves): + lines.append(_format_link_row(cve, cve_links(cve))) + lines.append("") + + if stats.ipv4s: + lines.append("### IPv4 Hosts\n") + for ip in sorted(stats.ipv4s, key=lambda x: tuple(int(p) for p in x.split("."))): + lines.append(_format_link_row(ip, ipv4_links(ip))) + lines.append("") + + if stats.cidrs: + lines.append("### IPv4 Ranges (CIDR)\n") + for cidr in sorted(stats.cidrs): + net_q = quote(cidr) + links = [ + ("Shodan net:", f"https://www.shodan.io/search?query=net%3A{net_q}"), + ("Censys", f"https://search.censys.io/search?resource=hosts&q=ip%3A%5B{quote(cidr.split('/')[0])}%2F{cidr.split('/')[1]}%5D"), + ("BGP.tools", f"https://bgp.tools/prefix/{net_q}"), + ] + lines.append(_format_link_row(cidr, links)) + lines.append("") + + if stats.hostnames: + lines.append("### Hostnames / Domains\n") + for host in sorted(stats.hostnames): + lines.append(_format_link_row(host, hostname_links(host))) + lines.append("") + + if stats.emails: + lines.append("### Emails\n") + for email in sorted(stats.emails): + lines.append(_format_link_row(email, email_links(email))) + lines.append("") + + if stats.repos or stats.commits: + lines.append("### Source Code (GitHub)\n") + # First emit repos; commits get their own bullet under the parent repo. + repo_commits: dict[tuple[str, str], list[str]] = {} + for org, repo, sha in stats.commits: + repo_commits.setdefault((org, repo), []).append(sha) + for org, repo in sorted(stats.repos): + lines.append(_format_link_row(f"{org}/{repo}", github_links(org, repo, None))) + for sha in sorted(repo_commits.get((org, repo), [])): + commit_url = f"https://github.com/{org}/{repo}/commit/{sha}" + lines.append(f" - commit [`{sha[:12]}`]({commit_url})") + # Orphan commits (commit seen but repo wasn't matched) — rare. + for (org, repo), shas in repo_commits.items(): + if (org, repo) not in stats.repos: + lines.append(_format_link_row(f"{org}/{repo}", github_links(org, repo, None))) + for sha in sorted(shas): + commit_url = f"https://github.com/{org}/{repo}/commit/{sha}" + lines.append(f" - commit [`{sha[:12]}`]({commit_url})") + lines.append("") + + if stats.vendors: + lines.append("### Vendors / Products\n") + for vendor in sorted(stats.vendors): + # Find the URL we used for this vendor. + for pat, disp, url in _VENDOR_PATTERNS: + if disp == vendor: + lines.append(f"- **{vendor}** — [docs]({url})") + break + lines.append("") + + return "\n".join(lines) + + +# ────────────────────────────────────────────────────────────────────────── +# Public entrypoint +# ────────────────────────────────────────────────────────────────────────── +def enrich_markdown(text: str) -> tuple[str, Stats]: + """Rewrite artifacts inline and append an INTEL APPENDIX. + + Code fences, inline code, and existing markdown links are preserved. + Returns the new text and a Stats object describing what was found. + """ + stats = Stats() + body, stash = _stash(text, stats) + + # Order matters: URLs before hostnames (URL pass also collects hostnames + # from its URL strings into stats so they appear in the appendix). CVEs + # before any other text replacement so the "CVE-..." literal becomes the + # link text. GitHub repo pass before hostnames so "org/repo" isn't eaten. + passes: list[Callable[[str, Stats], str]] = [ + _rewrite_urls, + _rewrite_cves, + _rewrite_emails, + _rewrite_github, + _rewrite_ipv4s, + _rewrite_hostnames, + _rewrite_vendors, + ] + for fn in passes: + body = fn(body, stats) + + body = _restore(body, stash) + # Post-pass: any backticked commit SHA on a line that also contains a + # github.com// link gets rewritten to that commit URL. + body = _link_orphan_commits(body, stats) + appendix = _build_appendix(stats) + return body + appendix, stats diff --git a/utils/helpers.py b/utils/helpers.py new file mode 100755 index 0000000..82ffdd2 --- /dev/null +++ b/utils/helpers.py @@ -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"" + + +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()] diff --git a/verify.sh b/verify.sh new file mode 100755 index 0000000..27813e9 --- /dev/null +++ b/verify.sh @@ -0,0 +1,6 @@ +#!/bin/bash +echo "=== main.py contains add_all_cog_commands? ===" +docker exec vulnforge grep -n add_all_cog_commands /app/main.py +echo +echo "=== relevant log entries ===" +docker logs vulnforge 2>&1 | grep -iE 'loaded|sync|locally|extension|exception|error' | tail -40 diff --git a/verify2.sh b/verify2.sh new file mode 100755 index 0000000..a2dd994 --- /dev/null +++ b/verify2.sh @@ -0,0 +1,12 @@ +#!/bin/bash +echo "=== Full vulnforge INFO logs ===" +docker logs vulnforge 2>&1 | grep -E 'vulnforge' | head -30 +echo +echo "=== Check if 'setup_hook' string appears anywhere ===" +docker logs vulnforge 2>&1 | grep -i setup_hook | head -10 +echo +echo "=== Search for any 'Loaded' or 'Failed' ===" +docker logs vulnforge 2>&1 | grep -iE 'loaded|failed' | head -20 +echo +echo "=== First 30 nextcord-level log lines (chronological) ===" +docker logs vulnforge 2>&1 | grep -E 'nextcord' | head -30