diff --git a/threat_intel_bot.py b/threat_intel_bot.py index bd53daa..3ab95da 100644 --- a/threat_intel_bot.py +++ b/threat_intel_bot.py @@ -9,7 +9,7 @@ import json import logging import os from datetime import datetime, timezone -from typing import Any, Dict +from typing import Any, Dict, List from pathlib import Path from functools import partial @@ -28,6 +28,7 @@ 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( @@ -37,13 +38,20 @@ logging.basicConfig( logger = logging.getLogger(__name__) CATEGORY_CONFIG = { - "news": {"label": "News", "emoji": "📰", "feeds_file": "feeds/news_feeds.json"}, - "malware": {"label": "Malware", "emoji": "🦠", "feeds_file": "feeds/malware_feeds.json"}, - "threat_intel": {"label": "Threat Intel", "emoji": "🛰️", "feeds_file": "feeds/threat_intel_feeds.json"}, - "osint": {"label": "OSINT", "emoji": "🕵️", "feeds_file": "feeds/osint_feeds.json"}, - "research": {"label": "Research", "emoji": "🔬", "feeds_file": "feeds/research_feeds.json"}, + 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_ 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, @@ -85,7 +93,7 @@ class ThreatIntelBot: if isinstance(raw_feed_types, list): mapped = [] for feed in raw_feed_types: - if feed in CATEGORY_CONFIG: + if feed in ALL_CATEGORIES: mapped.append(feed) elif feed == "daily": mapped.extend(["news", "threat_intel", "osint", "malware"]) @@ -123,7 +131,7 @@ class ThreatIntelBot: user_name = update.effective_user.first_name or "User" command_lines = [] - for key, cfg in CATEGORY_CONFIG.items(): + for key, cfg in ALL_CATEGORIES.items(): command_lines.append( f"{cfg['emoji']} `{('/on_' + key)}` / `{('/off_' + key)}` - {cfg['label']}" ) @@ -151,14 +159,14 @@ class ThreatIntelBot: async def category_on_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE, category: str): """Handle /on_ command.""" - if category not in CATEGORY_CONFIG: + 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 = CATEGORY_CONFIG[category] + cfg = ALL_CATEGORIES[category] if subscriber_key not in self.subscribers: self.subscribers[subscriber_key] = {'topic_id': message_thread_id, 'feed_types': []} @@ -184,14 +192,14 @@ class ThreatIntelBot: async def category_off_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE, category: str): """Handle /off_ command.""" - if category not in CATEGORY_CONFIG: + 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 = CATEGORY_CONFIG[category] + 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) @@ -224,7 +232,7 @@ class ThreatIntelBot: 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, or /on_research first.", + "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 @@ -393,26 +401,101 @@ class ThreatIntelBot: 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...") - # First run - mark existing articles and victims as seen, don't send alerts - try: - for category, cfg in CATEGORY_CONFIG.items(): + # 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: - logger.info(f"Initial {category.upper()} feed scan - marking existing articles as seen...") - await manager.fetch_all_feeds(initial_run=True) - except Exception as e: - logger.error(f"Error in initial feed scan: {e}") + 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: - await self.ransomware_fetcher.fetch_new_victims(session, initial_run=True) - logger.info("Initial ransomware scan complete — existing victims marked as seen") + 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 initial ransomware scan: {e}") + logger.error(f"Error in startup ransomware scan: {e}") while True: try: @@ -425,32 +508,9 @@ class ThreatIntelBot: if category_articles: logger.info(f"Found {len(category_articles)} new {category.upper()} articles") - # Deduplicate within this polling batch - unique_articles = [] - batch_seen_url = set() - batch_seen_content = set() - for article in all_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 - if url_key: - batch_seen_url.add(url_key) - batch_seen_content.add(content_key) - unique_articles.append(article) - - # Classify and send RSS articles - sent_count = 0 - for article in unique_articles: - classified = self.classifier.classify_article(article) - sent = await self.send_alert(classified) - if sent: - sent_count += 1 - await asyncio.sleep(1) - - if unique_articles: - logger.info(f"Processed {len(unique_articles)} unique articles, delivered {sent_count}") + 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") @@ -458,21 +518,14 @@ class ThreatIntelBot: if self.ransomware_fetcher: async with aiohttp.ClientSession() as session: new_victims = await self.ransomware_fetcher.fetch_new_victims(session) - victim_count = 0 - for victim in new_victims: - article = RansomwareFetcher.to_article(victim) - sent = await self.send_alert(article) - if sent: - self.ransomware_fetcher.mark_seen(victim.get("id", "")) - victim_count += 1 - await asyncio.sleep(1) 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(300) + await asyncio.sleep(POLL_INTERVAL_SECONDS) def build_application(self): """Build the Telegram application""" @@ -482,7 +535,7 @@ class ThreatIntelBot: # Add command handlers self.application.add_handler(CommandHandler("start", self.start_command)) self.application.add_handler(CommandHandler("help", self.help_command)) - for category in CATEGORY_CONFIG: + for category in ALL_CATEGORIES: self.application.add_handler( CommandHandler(f"on_{category}", partial(self.category_on_command, category=category)) )