Initial ThreatLens release

This commit is contained in:
2026-06-27 21:24:57 -04:00
commit d96e326d4d
27 changed files with 3059 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""cogs package — VulnForge feature modules (weekly digest, CVE lookups)."""
Executable
+261
View File
@@ -0,0 +1,261 @@
"""
cogs/cve.py — interactive CVE lookups + KEV listing + setup.
Slash commands:
/cve <cve_id> → 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))
Executable
+314
View File
@@ -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/<guild>/<channel>/<message>
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/.../.../<id>`).",
),
)
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))
Executable
+159
View File
@@ -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))