Initial ThreatLens release
This commit is contained in:
Executable
+261
@@ -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))
|
||||
Reference in New Issue
Block a user