telegram-rss-bot/threat_intel_bot.py

#!/usr/bin/env python3
"""
Threat Intelligence RSS Feed Telegram Bot
Main bot script handling commands, subscriptions, and real-time alerts
"""

import asyncio
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any, Dict, List
from pathlib import Path
from functools import partial

from telegram import Update
from telegram.ext import (
    ApplicationBuilder, 
    CommandHandler, 
    ContextTypes, 
    MessageHandler,
    filters
)
from telegram.constants import ParseMode
from telegram.error import TelegramError

import aiohttp
from rss_manager import RSSFeedManager
from content_classifier import ContentClassifier
from ransomware_fetcher import RansomwareFetcher
from bot_config import CATEGORY_LABELS, CATEGORY_FEEDS, EXTRA_CATEGORY_LABELS, POLL_INTERVAL_SECONDS

# Configure logging
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

CATEGORY_CONFIG = {
    key: {"label": label, "emoji": emoji, "feeds_file": CATEGORY_FEEDS[key]}
    for key, (label, emoji) in CATEGORY_LABELS.items()
}

# Categories that aren't RSS-fed (no feeds_file) and so are excluded from the
# RSS polling/status loops in monitor_feeds() and stats_command(), but are
# still subscribable via /on_<category> and shown in /start help.
EXTRA_CATEGORIES = {
    key: {"label": label, "emoji": emoji}
    for key, (label, emoji) in EXTRA_CATEGORY_LABELS.items()
}

ALL_CATEGORIES = {**CATEGORY_CONFIG, **EXTRA_CATEGORIES}

class ThreatIntelBot:
    def __init__(
        self,
        token: str,
        subscribers_file: str = "subscribers.json",
    ):
        self.token = token
        self.subscribers_file = subscribers_file
        self.subscribers: Dict[str, Dict[str, Any]] = {}
        self.classifier = ContentClassifier()
        self.application = None
        self.monitoring_task = None
        ransomware_key = os.getenv("RANSOMWARE_LIVE_API_KEY", "")
        self.ransomware_fetcher = RansomwareFetcher(ransomware_key) if ransomware_key else None
        self.load_subscribers()
    
    def load_subscribers(self):
        """Load subscriber list from file with feed type support"""
        try:
            with open(self.subscribers_file, 'r') as f:
                data = json.load(f)
                # Migrate old format to new format
                if 'subscribers' in data and isinstance(data['subscribers'], list):
                    # Old format: convert list to dict and preserve "all feeds" intent.
                    self.subscribers = {
                        str(sub_id): {'topic_id': None, 'feed_types': list(CATEGORY_CONFIG.keys())}
                        for sub_id in data['subscribers']
                    }
                else:
                    # New format: dict with topic info and feed types.
                    raw_subscribers = data.get('subscribers', {})
                    normalized: Dict[str, Dict[str, Any]] = {}
                    for chat_key, info in raw_subscribers.items():
                        if not isinstance(info, dict):
                            continue

                        topic_id = info.get('topic_id')
                        raw_feed_types = info.get('feed_types', [])
                        if isinstance(raw_feed_types, list):
                            mapped = []
                            for feed in raw_feed_types:
                                if feed in ALL_CATEGORIES:
                                    mapped.append(feed)
                                elif feed == "daily":
                                    mapped.extend(["news", "threat_intel", "osint", "malware"])
                                elif feed == "blogs":
                                    mapped.extend(["research", "threat_intel", "malware"])
                            feed_types = sorted(set(mapped))
                        else:
                            feed_types = []

                        normalized[str(chat_key)] = {
                            'topic_id': topic_id,
                            'feed_types': feed_types
                        }

                    self.subscribers = normalized
            logger.info(f"Loaded {len(self.subscribers)} subscribers")
        except FileNotFoundError:
            logger.info("No subscribers file found, starting with empty list")
            self.subscribers = {}
        except json.JSONDecodeError:
            logger.error("Error parsing subscribers file")
            self.subscribers = {}
    
    def save_subscribers(self):
        """Save subscriber list to file"""
        try:
            with open(self.subscribers_file, 'w') as f:
                json.dump({'subscribers': self.subscribers}, f, indent=2)
            logger.info(f"Saved {len(self.subscribers)} subscribers")
        except Exception as e:
            logger.error(f"Error saving subscribers: {e}")
    
    async def start_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /start command"""
        user_name = update.effective_user.first_name or "User"
        
        command_lines = []
        for key, cfg in ALL_CATEGORIES.items():
            command_lines.append(
                f"{cfg['emoji']} `{('/on_' + key)}` / `{('/off_' + key)}` - {cfg['label']}"
            )

        welcome_message = (
            f"🔴 **Welcome to Threat Intelligence Bot, {user_name}!** 🔴\n\n"
            "This bot monitors cybersecurity RSS feeds and delivers alerts to Telegram topics.\n\n"
            "**📂 Category Commands:**\n"
            + "\n".join(f"• {line}" for line in command_lines)
            + "\n\n**🛠️ General Commands:**\n"
            "• `/stats` - Check feed status and statistics\n"
            "• `/help` - Show this help message\n\n"
            "**💡 Tip for Telegram Groups with Topics:**\n"
            "Run each `/on_<category>` command inside the topic where that category should post."
        )
        
        await update.message.reply_text(
            welcome_message,
            parse_mode=ParseMode.MARKDOWN
        )
    
    async def help_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /help command"""
        await self.start_command(update, context)
    
    async def category_on_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE, category: str):
        """Handle /on_<category> command."""
        if category not in ALL_CATEGORIES:
            return

        chat_id = update.effective_chat.id
        message_thread_id = getattr(update.message, 'message_thread_id', None)
        subscriber_key = f"{chat_id}_{message_thread_id}" if message_thread_id else str(chat_id)
        is_forum_topic = message_thread_id is not None
        cfg = ALL_CATEGORIES[category]

        if subscriber_key not in self.subscribers:
            self.subscribers[subscriber_key] = {'topic_id': message_thread_id, 'feed_types': []}

        if category not in self.subscribers[subscriber_key]['feed_types']:
            self.subscribers[subscriber_key]['feed_types'].append(category)
            self.save_subscribers()
            chat_type = "forum topic" if is_forum_topic else ("group" if update.effective_chat.type in ['group', 'supergroup'] else "private chat")
            await update.message.reply_text(
                f"{cfg['emoji']} **{cfg['label']} Activated!**\n\n"
                f"You will now receive **{cfg['label'].lower()}** alerts in this {chat_type}.\n"
                f"Type `/off_{category}` to unsubscribe.",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if is_forum_topic else None
            )
            logger.info(f"Subscriber {subscriber_key} enabled {category.upper()} - Topic ID: {message_thread_id}")
        else:
            await update.message.reply_text(
                f"✅ {cfg['label']} is already enabled for this {('topic' if is_forum_topic else 'chat')}.",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if is_forum_topic else None
            )

    async def category_off_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE, category: str):
        """Handle /off_<category> command."""
        if category not in ALL_CATEGORIES:
            return

        chat_id = update.effective_chat.id
        message_thread_id = getattr(update.message, 'message_thread_id', None)
        subscriber_key = f"{chat_id}_{message_thread_id}" if message_thread_id else str(chat_id)
        is_forum_topic = message_thread_id is not None
        cfg = ALL_CATEGORIES[category]

        if subscriber_key in self.subscribers and category in self.subscribers[subscriber_key].get('feed_types', []):
            self.subscribers[subscriber_key]['feed_types'].remove(category)
            if not self.subscribers[subscriber_key]['feed_types']:
                del self.subscribers[subscriber_key]
            self.save_subscribers()

            chat_type = "topic" if is_forum_topic else "chat"
            await update.message.reply_text(
                f"✋ **{cfg['label']} disabled.**\n\n"
                f"This {chat_type} will no longer receive {cfg['label'].lower()} alerts.\n"
                f"Type `/on_{category}` to re-subscribe.",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if is_forum_topic else None
            )
            logger.info(f"Subscriber {subscriber_key} disabled {category.upper()} - Topic ID: {message_thread_id}")
        else:
            await update.message.reply_text(
                f"❌ This chat is not subscribed to {cfg['label'].lower()} alerts.",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if is_forum_topic else None
            )

    async def stats_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /stats command - show feed status for all categories."""
        chat_id = update.effective_chat.id
        message_thread_id = getattr(update.message, 'message_thread_id', None)

        subscriber_key = f"{chat_id}_{message_thread_id}" if message_thread_id else str(chat_id)

        if subscriber_key not in self.subscribers:
            await update.message.reply_text(
                "This topic is not subscribed yet. Use /on_news, /on_threat_intel, /on_malware, /on_osint, /on_research, or /on_ransomware first.",
                message_thread_id=message_thread_id if message_thread_id else None
            )
            return

        await update.message.reply_text(
            "🔍 **Checking feed status...** Please wait...",
            parse_mode=ParseMode.MARKDOWN,
            message_thread_id=message_thread_id if message_thread_id else None
        )

        try:
            total_feeds = 0
            online_feeds = 0

            # Send header message
            await update.message.reply_text(
                "📊 **Feed Status Report**\n",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if message_thread_id else None
            )
            await asyncio.sleep(0.5)

            for category, cfg in CATEGORY_CONFIG.items():
                async with RSSFeedManager(feeds_file=cfg["feeds_file"], feed_type=category) as manager:
                    status_data = await manager.check_feed_status()

                await update.message.reply_text(
                    f"{cfg['emoji']} **{cfg['label'].upper()} FEEDS:**\n",
                    parse_mode=ParseMode.MARKDOWN,
                    message_thread_id=message_thread_id if message_thread_id else None
                )
                await asyncio.sleep(0.2)

                for feed_category, feeds in status_data.items():
                    category_display = feed_category.replace('_', ' ').title()
                    category_message = f"**{category_display}:**\n"

                    for feed_name, feed_info in feeds.items():
                        total_feeds += 1
                        status_icon = feed_info['status']
                        entries = feed_info['entries']
                        last_updated = feed_info['last_updated']

                        if '✅' in status_icon:
                            online_feeds += 1

                        feed_line = f"• {feed_name}: {status_icon}\n"
                        if entries > 0:
                            feed_line += f"  └ {entries} articles, last: {last_updated}\n"
                        if 'error' in feed_info:
                            feed_line += f"  └ Error: {feed_info['error'][:50]}\n"
                        category_message += feed_line

                    await update.message.reply_text(
                        category_message,
                        parse_mode=ParseMode.MARKDOWN,
                        message_thread_id=message_thread_id if message_thread_id else None
                    )
                    await asyncio.sleep(0.2)

            # Send summary as final message
            summary = f"📈 **Summary:**\n"
            summary += f"• Total Feeds: {total_feeds}\n"
            summary += f"• Online: {online_feeds}\n"
            summary += f"• Offline: {total_feeds - online_feeds}\n"
            summary += f"• Subscribers: {len(self.subscribers)}\n"
            summary += f"• Last Check: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M UTC')}"

            await update.message.reply_text(
                summary,
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if message_thread_id else None
            )

        except Exception as e:
            logger.error(f"Error in stats command: {e}")
            await update.message.reply_text(
                "❌ Error checking feed status. Please try again later.",
                parse_mode=ParseMode.MARKDOWN,
                message_thread_id=message_thread_id if message_thread_id else None
            )
    async def handle_non_subscribers(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle messages from non-subscribers - completely silent for groups"""
        # Bot remains silent for all non-command messages
        # Only responds to slash commands, ignores everything else
        pass
    
    async def send_alert(self, article: Dict, context: ContextTypes.DEFAULT_TYPE = None) -> bool:
        """Send threat intel alert to subscribers based on feed type."""
        if not self.subscribers:
            return False

        try:
            message, thumbnail = RSSFeedManager.format_telegram_message(article)

            failed_sends = []
            article_feed_type = article.get('feed_type', 'general')
            successful_sends = 0

            for subscriber_key, sub_info in self.subscribers.copy().items():
                # Check if subscriber wants this feed type
                feed_types = sub_info.get('feed_types', [])
                if article_feed_type not in feed_types:
                    continue  # Skip if subscriber doesn't want this feed type

                # Parse chat_id from subscriber_key (format: "chatid_topicid" or "chatid")
                chat_id_part = subscriber_key.split('_')[0]
                chat_id = int(chat_id_part)
                topic_id = sub_info.get('topic_id') if isinstance(sub_info, dict) else None
                
                try:
                    if thumbnail:
                        # Try to send with photo
                        try:
                            await self.application.bot.send_photo(
                                chat_id=chat_id,
                                photo=thumbnail,
                                caption=message,
                                parse_mode=ParseMode.HTML,
                                message_thread_id=topic_id
                            )
                            successful_sends += 1
                        except TelegramError:
                            # Fallback to text if photo fails
                            await self.application.bot.send_message(
                                chat_id=chat_id,
                                text=message,
                                parse_mode=ParseMode.HTML,
                                disable_web_page_preview=False,
                                message_thread_id=topic_id
                            )
                            successful_sends += 1
                    else:
                        await self.application.bot.send_message(
                            chat_id=chat_id,
                            text=message,
                            parse_mode=ParseMode.HTML,
                            disable_web_page_preview=False,
                            message_thread_id=topic_id
                        )
                        successful_sends += 1
                
                except Exception as e:
                    logger.warning(f"Failed to send to {chat_id}: {e}")
                    if "bot was blocked" in str(e).lower() or "chat not found" in str(e).lower():
                        failed_sends.append(subscriber_key)

            # Remove blocked/invalid subscribers
            for failed_key in failed_sends:
                self.subscribers.pop(failed_key, None)
                logger.info(f"Removed blocked/invalid subscriber: {failed_key}")
            
            if failed_sends:
                self.save_subscribers()
            
            if successful_sends > 0:
                # Mark as sent only after at least one Telegram delivery succeeded.
                RSSFeedManager.mark_article_as_sent(article, seen_file="seen_articles.db")
                logger.info(f"Alert sent to {successful_sends} subscribers: {article['title'][:50]}...")
                return True

            logger.info(f"No eligible subscribers for article: {article['title'][:50]}...")
            return False
        
        except Exception as e:
            logger.error(f"Error sending alert: {e}")
            return False
    
    async def _classify_and_send(self, articles: List[Dict]) -> int:
        """Deduplicate a batch (exact + fuzzy cross-source title match), classify, and send.

        A single bad article is logged and skipped rather than aborting the rest
        of the batch (and, since ransomware polling runs after this in the caller,
        rather than aborting that too).
        """
        unique_articles = []
        batch_seen_url = set()
        batch_seen_content = set()
        batch_titles: List[str] = []
        for article in articles:
            fp = RSSFeedManager.get_article_fingerprints(article)
            url_key = fp['url_key']
            content_key = fp['content_key']
            if (url_key and url_key in batch_seen_url) or (content_key in batch_seen_content):
                continue
            title_norm = RSSFeedManager.normalize_text(article.get('title', ''))
            if title_norm and any(
                RSSFeedManager.titles_are_near_duplicate(title_norm, other) for other in batch_titles
            ):
                continue
            if url_key:
                batch_seen_url.add(url_key)
            batch_seen_content.add(content_key)
            if title_norm:
                batch_titles.append(title_norm)
            unique_articles.append(article)

        sent_count = 0
        for article in unique_articles:
            try:
                classified = self.classifier.classify_article(article)
                sent = await self.send_alert(classified)
                if sent:
                    sent_count += 1
            except Exception as e:
                logger.error(f"Error processing article {article.get('title', '')[:50]!r}: {e}")
            await asyncio.sleep(1)
        return sent_count

    async def _send_ransomware_victims(self, victims: List[Dict]) -> int:
        """Convert and send ransomware.live victims; one bad record doesn't stop the rest."""
        victim_count = 0
        for victim in victims:
            try:
                article = RansomwareFetcher.to_article(victim)
                sent = await self.send_alert(article)
                if sent:
                    self.ransomware_fetcher.mark_seen(victim.get("id", ""))
                    victim_count += 1
            except Exception as e:
                logger.error(f"Error processing ransomware victim {victim.get('id', '')}: {e}")
            await asyncio.sleep(1)
        return victim_count

    async def monitor_feeds(self):
        """Background task to monitor RSS feeds by category."""
        logger.info("Starting RSS feed monitoring for category feeds...")

        # Only mark-without-alerting on a genuinely fresh seen-DB. On a restart of an
        # already-running bot, do a real fetch instead so articles published while the
        # service was stopped (e.g. during a deploy) still get delivered.
        startup_articles: List[Dict] = []
        for category, cfg in CATEGORY_CONFIG.items():
            try:
                async with RSSFeedManager(feeds_file=cfg["feeds_file"], feed_type=category) as manager:
                    if manager.first_run:
                        logger.info(f"Initial {category.upper()} feed scan - marking existing articles as seen...")
                        await manager.fetch_all_feeds(initial_run=True)
                    else:
                        articles = await manager.fetch_all_feeds(initial_run=False)
                        if articles:
                            logger.info(f"Restart {category.upper()} scan - {len(articles)} article(s) missed while stopped")
                        startup_articles.extend(articles)
            except Exception as e:
                logger.error(f"Error in startup {category.upper()} scan: {e}")

        if startup_articles:
            sent = await self._classify_and_send(startup_articles)
            logger.info(f"Delivered {sent} article(s) missed during downtime")

        if self.ransomware_fetcher:
            try:
                async with aiohttp.ClientSession() as session:
                    if self.ransomware_fetcher.is_first_run():
                        await self.ransomware_fetcher.fetch_new_victims(session, initial_run=True)
                        logger.info("Initial ransomware scan complete — existing victims marked as seen")
                    else:
                        missed_victims = await self.ransomware_fetcher.fetch_new_victims(session)
                        if missed_victims:
                            logger.info(f"Restart ransomware scan - {len(missed_victims)} victim(s) missed while stopped")
                            await self._send_ransomware_victims(missed_victims)
            except Exception as e:
                logger.error(f"Error in startup ransomware scan: {e}")

        while True:
            try:
                all_articles = []

                for category, cfg in CATEGORY_CONFIG.items():
                    async with RSSFeedManager(feeds_file=cfg["feeds_file"], feed_type=category) as manager:
                        category_articles = await manager.fetch_all_feeds(initial_run=False)
                        all_articles.extend(category_articles)
                        if category_articles:
                            logger.info(f"Found {len(category_articles)} new {category.upper()} articles")

                sent_count = await self._classify_and_send(all_articles)
                if all_articles:
                    logger.info(f"Processed {len(all_articles)} candidate articles, delivered {sent_count}")
                else:
                    logger.info("No new articles found")

                # Poll ransomware.live for new victims
                if self.ransomware_fetcher:
                    async with aiohttp.ClientSession() as session:
                        new_victims = await self.ransomware_fetcher.fetch_new_victims(session)
                    if new_victims:
                        victim_count = await self._send_ransomware_victims(new_victims)
                        logger.info(f"Ransomware victims: {len(new_victims)} new, {victim_count} delivered")

            except Exception as e:
                logger.error(f"Error in feed monitoring: {e}")

            await asyncio.sleep(POLL_INTERVAL_SECONDS)
    
    def build_application(self):
        """Build the Telegram application"""
        # Build application
        self.application = ApplicationBuilder().token(self.token).build()
        
        # Add command handlers
        self.application.add_handler(CommandHandler("start", self.start_command))
        self.application.add_handler(CommandHandler("help", self.help_command))
        for category in ALL_CATEGORIES:
            self.application.add_handler(
                CommandHandler(f"on_{category}", partial(self.category_on_command, category=category))
            )
            self.application.add_handler(
                CommandHandler(f"off_{category}", partial(self.category_off_command, category=category))
            )
        self.application.add_handler(CommandHandler("stats", self.stats_command))
        
        # Handle all other messages (for non-subscribers)
        self.application.add_handler(
            MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_non_subscribers)
        )
        
        logger.info("Telegram bot application built successfully")
    
    async def start_monitoring(self):
        """Start the background monitoring task"""
        if not self.monitoring_task:
            self.monitoring_task = asyncio.create_task(self.monitor_feeds())
            logger.info("RSS monitoring task started")
    
    async def stop_monitoring(self):
        """Stop the background monitoring task"""
        if self.monitoring_task:
            self.monitoring_task.cancel()
            try:
                await self.monitoring_task
            except asyncio.CancelledError:
                pass
            self.monitoring_task = None
            logger.info("RSS monitoring task stopped")
    
    async def run(self):
        """Run the bot"""
        self.build_application()
        
        try:
            # Initialize and start the application
            await self.application.initialize()
            await self.application.start()
            
            # Start background monitoring
            await self.start_monitoring()
            
            # Start polling
            await self.application.updater.start_polling()
            
            logger.info("🚀 Threat Intel Bot is running!")
            logger.info(f"📊 {len(self.subscribers)} subscribers loaded")
            
            # Keep running until interrupted
            while True:
                await asyncio.sleep(1)
                
        except KeyboardInterrupt:
            logger.info("Received interrupt, shutting down...")
        
        except Exception as e:
            logger.error(f"Error running bot: {e}")
            raise
        
        finally:
            # Cleanup
            await self.stop_monitoring()
            await self.application.stop()
            await self.application.shutdown()

def main():
    """Main function"""
    # Load .env file if it exists
    env_file = Path('.env')
    if env_file.exists():
        with open(env_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip().strip('"\'')
    
    # Get bot token from environment variable
    bot_token = os.getenv('BOT_TOKEN')
    if not bot_token:
        logger.error("BOT_TOKEN not found!")
        logger.info("Either:")
        logger.info("1. Set environment variable: export BOT_TOKEN='your_token'")
        logger.info("2. Create .env file with: BOT_TOKEN=your_token")
        return
    
    bot = ThreatIntelBot(bot_token)
    
    try:
        asyncio.run(bot.run())
    except KeyboardInterrupt:
        logger.info("Bot stopped by user")
    except Exception as e:
        logger.error(f"Bot crashed: {e}")

if __name__ == "__main__":
    main()