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