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:
bot
2026-08-25 12:20:08 +03:00
parent 1b2eb74cf1
commit e8866ce4e6
+106 -53
View File
@@ -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:
# 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 initial feed scan: {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 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))
)