262 lines
9.3 KiB
Python
Executable File
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))
|