315 lines
13 KiB
Python
Executable File
315 lines
13 KiB
Python
Executable File
"""
|
|
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))
|