Files
ThreatLens/cogs/cve.py
T
2026-06-27 21:24:57 -04:00

262 lines
9.3 KiB
Python
Executable File

"""
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))