""" 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/// 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/.../.../`).", ), ) 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))