Initial ThreatLens release
This commit is contained in:
Executable
+1
@@ -0,0 +1 @@
|
||||
"""cogs package — VulnForge feature modules (weekly digest, CVE lookups)."""
|
||||
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))
|
||||
Executable
+314
@@ -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
@@ -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))
|
||||
Reference in New Issue
Block a user