Initial ThreatLens release

2026-06-27 21:24:57 -04:00
commit d96e326d4d
27 changed files with 3059 additions and 0 deletions
@@ -0,0 +1,160 @@
+"""
+main.py — VulnForge entrypoint.
+
+Boots the Nextcord bot, wires up the shared SecurityAPI client + SQLite store,
+loads cogs, configures logging, and handles graceful startup/shutdown.
+
+Run:
+    python main.py
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import signal
+import sys
+
+import nextcord
+from nextcord.ext import commands
+
+from config import settings
+from utils.api import SecurityAPI
+from utils.helpers import Store
+
+# ──────────────────────────────────────────────────────────────
+#  Logging
+# ──────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=getattr(logging, settings.log_level, logging.INFO),
+    format="%(asctime)s │ %(levelname)-8s │ %(name)s │ %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+log = logging.getLogger("vulnforge")
+
+
+class VulnForge(commands.Bot):
+    """The VulnForge bot — owns the shared API client and data store."""
+
+    def __init__(self) -> None:
+        intents = nextcord.Intents.default()
+        # When GUILD_ID is set we scope ALL commands to that guild so they appear
+        # instantly (and we don't depend on Nextcord auto-discovering guild_ids
+        # from individual decorators, which was failing to push the guild sync).
+        kwargs: dict = {"intents": intents}
+        if settings.guild_id:
+            kwargs["default_guild_ids"] = [settings.guild_id]
+        super().__init__(**kwargs)
+        self.api = SecurityAPI()
+        self.store = Store(settings.db_path)
+
+    async def on_ready(self) -> None:
+        log.info("Logged in as %s (id: %s)", self.user, self.user.id)
+        log.info("Connected to %s guild(s)", len(self.guilds))
+        # Force a fresh sync so command renames/additions show up immediately.
+        try:
+            if settings.guild_id:
+                await self.sync_application_commands(guild_id=settings.guild_id)
+                log.info("Slash commands synced to guild %s", settings.guild_id)
+            else:
+                await self.sync_application_commands()
+                log.info("Slash commands synced globally")
+        except Exception:  # noqa: BLE001
+            log.exception("sync_application_commands failed")
+        await self.change_presence(
+            activity=nextcord.Activity(
+                type=nextcord.ActivityType.watching,
+                name="the CVE & KEV feeds ☣",
+            )
+        )
+
+    async def setup(self) -> None:
+        """Load extensions and register cog commands BEFORE we connect.
+
+        We can't use Nextcord's `setup_hook` here — in 2.6 it isn't invoked by
+        bot.start() in this configuration, so cogs never got loaded and the
+        bot's command list stayed empty (sync was a no-op). Call this manually
+        from the runner before bot.start().
+        """
+        settings.ensure_dirs()
+        await self.store.init()
+        await self.api.start()
+
+        for ext in ("cogs.cve", "cogs.weekly", "cogs.enrich"):
+            try:
+                self.load_extension(ext)
+                log.info("Loaded extension: %s", ext)
+            except Exception:  # noqa: BLE001
+                log.exception("Failed to load extension %s", ext)
+
+        # Nextcord 2.6 quirk: add_cog() attaches slash commands to the cog
+        # object (cog.application_commands) but does NOT push them into the
+        # bot's registration list. add_all_cog_commands() also no-ops here.
+        # Manually push each cog command onto the bot so sync sees them.
+        for cog_name, cog in self.cogs.items():
+            count = 0
+            for cmd in getattr(cog, "application_commands", []):
+                if settings.guild_id and not cmd.guild_ids_to_rollout:
+                    cmd.guild_ids_to_rollout.add(settings.guild_id)
+                self.add_application_command(cmd, use_rollout=True)
+                count += 1
+            log.info("Registered %d command(s) from cog %s", count, cog_name)
+        # get_all_application_commands() returns ALL commands (guild + global);
+        # get_application_commands() only returns globals so it's misleading here.
+        log.info(
+            "Total application commands on bot: %d",
+            len(list(self.get_all_application_commands())),
+        )
+
+    async def close(self) -> None:
+        log.info("Shutting down — closing API session…")
+        await self.api.close()
+        await super().close()
+
+
+def main() -> None:
+    if not settings.discord_token:
+        log.error("DISCORD_TOKEN is not set. Copy .env.example to .env and fill it in.")
+        sys.exit(1)
+
+    async def runner() -> None:
+        # IMPORTANT: construct the bot *inside* the running loop. Nextcord binds
+        # its gateway/heartbeat to the active event loop; building it before the
+        # loop exists causes a loop mismatch where the heartbeat is scheduled on
+        # a dead loop ("heartbeat blocked for N seconds" while the loop sits idle
+        # in selector.poll). Creating it here guarantees a single, correct loop.
+        bot = VulnForge()
+        loop = asyncio.get_running_loop()
+
+        def _handle_signal(*_args: object) -> None:
+            log.info("Signal received — stopping bot…")
+            loop.create_task(bot.close())
+
+        # Graceful Ctrl+C / SIGTERM (e.g. `docker stop`) handling.
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                loop.add_signal_handler(sig, _handle_signal)
+            except NotImplementedError:
+                # add_signal_handler is unavailable on Windows; KeyboardInterrupt covers SIGINT.
+                pass
+
+        # Nextcord's Bot is not an async context manager (unlike discord.py 2.x),
+        # so start it directly and ensure a clean close on exit.
+        try:
+            # Load extensions + register cog commands BEFORE connecting. Doing
+            # this in setup_hook didn't run in this Nextcord version.
+            await bot.setup()
+            await bot.start(settings.discord_token)
+        finally:
+            if not bot.is_closed():
+                await bot.close()
+
+    try:
+        asyncio.run(runner())
+    except KeyboardInterrupt:
+        log.info("KeyboardInterrupt — exiting.")
+
+
+if __name__ == "__main__":
+    main()