ThreatLens/main.py

"""
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()