Fix restart data loss; isolate per-article failures in the poll loop
monitor_feeds() unconditionally ran the startup scan with initial_run=True, which marks everything unseen as seen WITHOUT alerting — not just on the bot's very first run, but on every restart. Anything published while the service was stopped for a deploy was silently dropped. RSSFeedManager.first_run and the new RansomwareFetcher.is_first_run() were already available to tell the two cases apart; monitor_feeds() now actually checks them, and does a real fetch-and-alert pass for anything missed on a plain restart. Also splits the per-article classify+send loop out into _classify_and_send()/_send_ransomware_victims(), each wrapping a single article/victim in its own try/except — previously one bad article could raise past the whole batch and skip ransomware polling for that cycle too. Category/poll-interval config now comes from bot_config.py instead of a local copy. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
+112
-59
@@ -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_<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,
|
||||
@@ -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_<category> 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_<category> 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))
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user