telegram-rss-bot/rss_manager.py

#!/usr/bin/env python3
"""
RSS Feed Manager for Threat Intelligence Bot
Handles feed parsing, article extraction, and message formatting
"""

import feedparser
import json
import logging
import asyncio
import aiohttp
import sqlite3
from datetime import datetime, timezone, timedelta
import hashlib
from typing import Dict, List, Optional, Tuple
from bs4 import BeautifulSoup
from html import escape
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from pathlib import Path

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

class RSSFeedManager:
    def __init__(self, feeds_file: str = "feeds.json", seen_file: str = "seen_articles.db", feed_type: str = "general"):
        self.feeds_file = feeds_file
        self.seen_file = seen_file
        self.feed_type = feed_type  # 'daily', 'blogs', or 'general'
        self.feeds = {}
        self.seen_articles = {}  # Changed to dict to store timestamps
        self.session = None
        self.first_run = True
        self.load_feeds()
        self.load_seen_articles()
    
    async def __aenter__(self):
        """Async context manager entry"""
        self.session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=30),
            headers={'User-Agent': 'ThreatIntel-Bot/1.0'}
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        if self.session:
            await self.session.close()
    
    def load_feeds(self):
        """Load RSS feeds from JSON file"""
        if not self.feeds_file:
            self.feeds = {}
            return
        try:
            with open(self.feeds_file, 'r') as f:
                self.feeds = json.load(f)
            logger.info(f"Loaded {sum(len(category) for category in self.feeds.values())} feeds from {len(self.feeds)} categories")
        except FileNotFoundError:
            logger.error(f"Feeds file {self.feeds_file} not found")
            self.feeds = {}
        except json.JSONDecodeError as e:
            logger.error(f"Error parsing feeds file: {e}")
            self.feeds = {}
    
    def load_seen_articles(self):
        """Load seen article keys from SQLite, migrating from legacy JSON when present."""
        self._ensure_seen_db()
        self._migrate_legacy_seen_json()
        self._load_seen_from_sqlite()

    def _ensure_seen_db(self):
        """Create seen-article database and indexes if missing."""
        try:
            with sqlite3.connect(self.seen_file) as conn:
                conn.execute(
                    """
                    CREATE TABLE IF NOT EXISTS seen_articles (
                        article_key TEXT PRIMARY KEY,
                        seen_at TEXT NOT NULL
                    )
                    """
                )
                conn.execute(
                    "CREATE INDEX IF NOT EXISTS idx_seen_articles_seen_at ON seen_articles (seen_at)"
                )
                conn.commit()
        except Exception as e:
            logger.error(f"Error initializing seen-article database: {e}")

    def _load_seen_from_sqlite(self):
        """Load seen keys from SQLite into memory."""
        try:
            with sqlite3.connect(self.seen_file) as conn:
                cursor = conn.execute("SELECT article_key, seen_at FROM seen_articles")
                self.seen_articles = {row[0]: row[1] for row in cursor.fetchall()}
            self.first_run = len(self.seen_articles) == 0
            logger.info(f"Loaded {len(self.seen_articles)} seen articles (first_run: {self.first_run})")
        except Exception as e:
            logger.error(f"Error loading seen articles from database: {e}")
            self.seen_articles = {}
            self.first_run = True

    def _migrate_legacy_seen_json(self):
        """One-time migration from legacy seen_articles.json into SQLite."""
        if Path(self.seen_file).suffix.lower() == ".json":
            return

        json_path = Path(self.seen_file).with_suffix(".json")
        if not json_path.exists():
            return

        try:
            with sqlite3.connect(self.seen_file) as conn:
                existing_count = conn.execute("SELECT COUNT(1) FROM seen_articles").fetchone()[0]
            if existing_count > 0:
                return

            with open(json_path, "r") as f:
                seen_data = json.load(f)

            current_time = datetime.now(timezone.utc).isoformat()
            raw_seen = seen_data.get("seen", {})
            if isinstance(raw_seen, list):
                migrated = {str(k): current_time for k in raw_seen}
            elif isinstance(raw_seen, dict):
                migrated = {str(k): str(v) for k, v in raw_seen.items()}
            else:
                migrated = {}

            if migrated:
                with sqlite3.connect(self.seen_file) as conn:
                    conn.executemany(
                        "INSERT OR REPLACE INTO seen_articles (article_key, seen_at) VALUES (?, ?)",
                        [(k, v) for k, v in migrated.items()],
                    )
                    conn.commit()
                logger.info(f"Migrated {len(migrated)} seen articles from {json_path.name} to SQLite")
        except FileNotFoundError:
            return
        except json.JSONDecodeError:
            logger.error("Error parsing legacy seen_articles.json during migration")
        except Exception as e:
            logger.error(f"Error migrating legacy seen articles: {e}")

    def save_seen_articles(self):
        """Save seen article keys with retention in SQLite."""
        try:
            cutoff_time = datetime.now(timezone.utc) - timedelta(days=7)
            cutoff_iso = cutoff_time.isoformat()
            cleaned_articles = {}
            with sqlite3.connect(self.seen_file) as conn:
                conn.execute("DELETE FROM seen_articles WHERE seen_at < ?", (cutoff_iso,))
                conn.executemany(
                    "INSERT OR REPLACE INTO seen_articles (article_key, seen_at) VALUES (?, ?)",
                    [(k, v) for k, v in self.seen_articles.items()],
                )
                conn.commit()
                cursor = conn.execute("SELECT article_key, seen_at FROM seen_articles")
                cleaned_articles = {row[0]: row[1] for row in cursor.fetchall()}

            self.seen_articles = cleaned_articles
            logger.info(f"Saved {len(self.seen_articles)} seen articles (cleaned old entries)")
        except Exception as e:
            logger.error(f"Error saving seen articles: {e}")
    
    def generate_article_hash(self, article: Dict) -> str:
        """Generate unique hash for article based on URL only for better duplicate detection"""
        # Use only the URL for hash to catch duplicates across different RSS sources
        # Different sources may have same article with different titles/dates
        url = article.get('link', '')
        return hashlib.md5(url.encode()).hexdigest()

    @staticmethod
    def canonicalize_url(url: str) -> str:
        """Normalize URLs so tracking variants map to one dedup key."""
        if not url:
            return ""

        tracking_params = {"fbclid", "gclid", "mc_cid", "mc_eid", "igshid", "mkt_tok", "yclid"}
        parts = urlsplit(url.strip())

        filtered_query = []
        for key, value in parse_qsl(parts.query, keep_blank_values=True):
            lowered = key.lower()
            if lowered.startswith("utm_") or lowered in tracking_params:
                continue
            filtered_query.append((key, value))

        cleaned_query = urlencode(filtered_query, doseq=True)
        cleaned_path = parts.path.rstrip("/") if parts.path not in ("", "/") else parts.path

        return urlunsplit((parts.scheme.lower(), parts.netloc.lower(), cleaned_path, cleaned_query, ""))

    @staticmethod
    def normalize_text(value: str) -> str:
        """Normalize text for stable content fingerprinting."""
        return " ".join((value or "").lower().split())

    @classmethod
    def get_article_fingerprints(cls, article: Dict) -> Dict[str, str]:
        """Build URL/content fingerprints for dedup checks."""
        raw_url = article.get('url') or article.get('link') or ""
        canonical_url = cls.canonicalize_url(raw_url)

        title = cls.normalize_text(article.get('title', ''))
        description = cls.normalize_text(article.get('description', ''))
        content_blob = f"{title}|{description}"

        url_fp = hashlib.sha256(canonical_url.encode("utf-8")).hexdigest() if canonical_url else ""
        content_fp = hashlib.sha256(content_blob.encode("utf-8")).hexdigest()

        return {
            'url_key': f"url:{url_fp}" if url_fp else "",
            'content_key': f"content:{content_fp}",
            'canonical_url': canonical_url,
        }

    def is_article_recent(self, published_dt: datetime) -> bool:
        """Allow articles published within the last 48 hours."""
        if published_dt.tzinfo is None:
            published_dt = published_dt.replace(tzinfo=timezone.utc)
        age = datetime.now(timezone.utc) - published_dt.astimezone(timezone.utc)
        return age.total_seconds() <= 48 * 3600

    def is_duplicate_article(self, article: Dict) -> bool:
        """Check URL/content fingerprint duplicates against sent-history."""
        fingerprints = self.get_article_fingerprints(article)
        url_key = fingerprints['url_key']
        content_key = fingerprints['content_key']
        return (url_key and url_key in self.seen_articles) or (content_key in self.seen_articles)

    @classmethod
    def mark_article_as_sent(cls, article: Dict, seen_file: str = "seen_articles.db"):
        """Persist fingerprints only after successful delivery."""
        manager = cls(feeds_file="", seen_file=seen_file, feed_type=article.get('feed_type', 'general'))
        fingerprints = cls.get_article_fingerprints(article)
        now_iso = datetime.now(timezone.utc).isoformat()

        if fingerprints['url_key']:
            manager.seen_articles[fingerprints['url_key']] = now_iso
        manager.seen_articles[fingerprints['content_key']] = now_iso
        manager.save_seen_articles()

    def clean_html(self, text: str) -> str:
        """Clean HTML tags and decode entities"""
        if not text:
            return ""
        # Remove HTML tags
        soup = BeautifulSoup(text, 'html.parser')
        cleaned = soup.get_text()
        # Clean up whitespace
        cleaned = ' '.join(cleaned.split())
        return cleaned[:500] + "..." if len(cleaned) > 500 else cleaned
    
    def extract_thumbnail(self, article: Dict, content: str = "") -> Optional[str]:
        """Extract thumbnail URL from article"""
        # Try media content first
        media_content = getattr(article, 'media_content', [])
        if media_content and isinstance(media_content, list):
            for media in media_content:
                if media.get('type', '').startswith('image/'):
                    return media.get('url')
        
        # Try enclosures
        enclosures = getattr(article, 'enclosures', [])
        for enclosure in enclosures:
            if enclosure.get('type', '').startswith('image/'):
                return enclosure.get('href')
        
        # Try media thumbnail
        if hasattr(article, 'media_thumbnail'):
            thumbnails = getattr(article, 'media_thumbnail', [])
            if thumbnails:
                return thumbnails[0].get('url')
        
        # Try parsing content for images
        if content:
            soup = BeautifulSoup(content, 'html.parser')
            img_tags = soup.find_all('img')
            for img in img_tags:
                src = img.get('src')
                if src and any(ext in src.lower() for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp']):
                    return src
        
        return None
    
    async def fetch_feed(self, url: str, name: str, initial_run: bool = False) -> List[Dict]:
        """Fetch and parse single RSS feed"""
        try:
            async with self.session.get(url) as response:
                if response.status == 200:
                    content = await response.text()
                    feed = feedparser.parse(content)
                    
                    if feed.bozo and feed.bozo_exception:
                        logger.warning(f"Feed parsing warning for {name}: {feed.bozo_exception}")
                    
                    articles = []
                    for entry in feed.entries[:10]:  # Limit to 10 most recent
                        # Parse publication date. Skip undated items to prevent stale replays.
                        published = getattr(entry, 'published_parsed', None) or getattr(entry, 'updated_parsed', None)
                        if not published:
                            continue

                        published_dt = datetime(*published[:6], tzinfo=timezone.utc)

                        # Strict real-time gate: today's UTC articles only.
                        if not self.is_article_recent(published_dt):
                            continue

                        # Get content/description
                        content = ""
                        if hasattr(entry, 'content') and entry.content:
                            content = entry.content[0].value if isinstance(entry.content, list) else entry.content
                        elif hasattr(entry, 'summary'):
                            content = entry.summary
                        elif hasattr(entry, 'description'):
                            content = entry.description
                        
                        article = {
                            'hash': self.generate_article_hash(entry),
                            'title': self.clean_html(getattr(entry, 'title', 'No Title')),
                            'description': self.clean_html(content),
                            'url': self.canonicalize_url(getattr(entry, 'link', '')),
                            'published': published_dt.isoformat(),
                            'published_human': published_dt.strftime('%Y-%m-%d %H:%M UTC'),
                            'source': name,
                            'category': '',  # Will be set by caller
                            'thumbnail': self.extract_thumbnail(entry, content),
                            'feed_type': self.feed_type  # Track which feed type this came from
                        }

                        # Skip if article has already been sent recently.
                        if self.is_duplicate_article(article):
                            continue

                        articles.append(article)
                    
                    if initial_run and articles:
                        logger.info(f"Initial run - marked {len(articles)} articles from {name} as seen")
                    elif articles:
                        logger.info(f"Fetched {len(articles)} new articles from {name}")
                    
                    return articles
                else:
                    logger.error(f"HTTP {response.status} for feed {name}: {url}")
                    return []
        
        except asyncio.TimeoutError:
            logger.error(f"Timeout fetching feed {name}: {url}")
            return []
        except Exception as e:
            logger.error(f"Error fetching feed {name}: {e}")
            return []
    
    async def fetch_all_feeds(self, initial_run: bool = False) -> List[Dict]:
        """Fetch all RSS feeds concurrently"""
        if not self.session:
            raise RuntimeError("RSS Manager not initialized. Use 'async with' context manager.")
        
        tasks = []
        for category_name, feeds in self.feeds.items():
            for feed_name, feed_url in feeds.items():
                task = self.fetch_feed(feed_url, feed_name, initial_run)
                tasks.append((task, category_name))
        
        all_articles = []
        results = await asyncio.gather(*[task for task, _ in tasks], return_exceptions=True)
        
        for (_, category), result in zip(tasks, results):
            if isinstance(result, Exception):
                logger.error(f"Exception in category {category}: {result}")
                continue
            
            for article in result:
                article['category'] = category
                all_articles.append(article)
        
        # Sort by published date (newest first)
        all_articles.sort(key=lambda x: x['published'], reverse=True)
        
        if initial_run and all_articles:
            # On first run, mark today's currently-visible items as sent to avoid startup flood.
            now_iso = datetime.now(timezone.utc).isoformat()
            for article in all_articles:
                fingerprints = self.get_article_fingerprints(article)
                if fingerprints['url_key']:
                    self.seen_articles[fingerprints['url_key']] = now_iso
                self.seen_articles[fingerprints['content_key']] = now_iso
            self.save_seen_articles()
            logger.info(f"Initial run: Marked {len(all_articles)} existing articles as seen (no alerts sent)")
            return []  # Don't send alerts on first run
        
        logger.info(f"Total new articles: {len(all_articles)}")
        return all_articles
    
    async def check_feed_status(self) -> Dict[str, Dict]:
        """Check status of all feeds"""
        if not self.session:
            raise RuntimeError("RSS Manager not initialized. Use 'async with' context manager.")
        
        status = {}
        
        for category_name, feeds in self.feeds.items():
            status[category_name] = {}
            for feed_name, feed_url in feeds.items():
                try:
                    async with self.session.get(feed_url) as response:
                        if response.status == 200:
                            content = await response.text()
                            feed = feedparser.parse(content)
                            entry_count = len(feed.entries)
                            last_updated = "Unknown"
                            
                            if feed.entries:
                                latest = feed.entries[0]
                                if hasattr(latest, 'published_parsed') and latest.published_parsed:
                                    dt = datetime(*latest.published_parsed[:6], tzinfo=timezone.utc)
                                    last_updated = dt.strftime('%Y-%m-%d %H:%M UTC')
                            
                            status[category_name][feed_name] = {
                                'status': '✅ Online',
                                'entries': entry_count,
                                'last_updated': last_updated,
                                'http_status': response.status
                            }
                        else:
                            status[category_name][feed_name] = {
                                'status': '❌ HTTP Error',
                                'entries': 0,
                                'last_updated': 'Error',
                                'http_status': response.status
                            }
                
                except Exception as e:
                    status[category_name][feed_name] = {
                        'status': '❌ Error',
                        'entries': 0,
                        'last_updated': 'Error',
                        'error': str(e)[:100]
                    }
        
        return status
    
    @staticmethod
    def format_telegram_message(article: Dict) -> Tuple[str, Optional[str]]:
        """Format article for Telegram message"""
        category_emojis = {
            'news': '📰',
            'malware': '🦠',
            'threat_intel': '🛰️',
            'osint': '🕵️',
            'research': '🔬'
        }

        emoji = category_emojis.get(article.get('category', ''), '📰')
        title = escape(article.get('title', 'No Title'))
        description = escape(article.get('description', ''))
        source = escape(article.get('source', 'Unknown Source'))
        published_human = escape(article.get('published_human', 'Unknown'))
        url = escape(article.get('url', ''), quote=True)
        cves = article.get('cves', [])[:5]
        threat_actors = article.get('threat_actors', [])[:3]
        malware_families = article.get('malware_families', [])[:3]

        message = f"{emoji} <b>{title}</b>\n\n"

        if description:
            message += f"{description}\n\n"

        message += f"📡 {source}  ·  {published_human}\n"

        if cves:
            message += f"🆔 {', '.join(escape(c) for c in cves)}\n"

        if threat_actors:
            message += f"👤 {', '.join(escape(a.title()) for a in threat_actors)}\n"

        if malware_families:
            message += f"🦠 {', '.join(escape(f) for f in malware_families)}\n"

        message += f"\n🔗 <a href=\"{url}\">Read Full Article</a>"

        return message, article.get('thumbnail')

# Example usage for testing
async def main():
    async with RSSFeedManager() as rss_manager:
        # Test fetching feeds
        articles = await rss_manager.fetch_all_feeds()
        
        for article in articles[:3]:  # Show first 3 articles
            message, thumbnail = rss_manager.format_telegram_message(article)
            print("=" * 50)
            print(message)
            if thumbnail:
                print(f"Thumbnail: {thumbnail}")
            print()
        
        # Test status check
        print("Feed Status Check:")
        status = await rss_manager.check_feed_status()
        for category, feeds in status.items():
            print(f"\n{category}:")
            for feed_name, feed_status in feeds.items():
                print(f"  {feed_name}: {feed_status['status']} ({feed_status['entries']} entries)")

if __name__ == "__main__":
    asyncio.run(main())