#!/usr/bin/env python3
"""
Content Classifier for Threat Intelligence RSS Bot
Extracts CVEs, threat actors, malware families, and MITRE technique IDs.
"""

import re
from typing import Dict, List
from datetime import datetime, timezone
import logging

logger = logging.getLogger(__name__)


class ContentClassifier:

    THREAT_ACTORS = {
        'apt1', 'apt28', 'apt29', 'apt32', 'apt33', 'apt34', 'apt35', 'apt37', 'apt38', 'apt39', 'apt40', 'apt41',
        'lazarus', 'kimsuky', 'andariel', 'fancy bear', 'cozy bear', 'sandworm',
        'turla', 'equation group', 'carbanak', 'fin7', 'fin6', 'fin8',
        'nobelium', 'hafnium', 'phosphorus', 'holmium', 'strontium',
        'volt typhoon', 'flax typhoon', 'mustang panda', 'winnti',
        'scattered spider', 'lapsus', 'lapsus$', 'unc2452', 'unc3944',
        'ta505', 'ta577', 'ta558', 'gold southfield', 'gold dupont',
    }

    # Stored as display names; matched case-insensitively
    MALWARE_FAMILIES = {
        # Ransomware
        'LockBit', 'REvil', 'BlackCat', 'ALPHV', 'Cl0p', 'Conti', 'DarkSide',
        'Ryuk', 'BlackMatter', 'Akira', 'Black Basta', 'RansomHub', 'Rhysida',
        'Medusa Locker', 'Cactus Ransomware', '8Base', 'Play Ransomware',
        'Hunters International', 'Inc Ransom', 'Monti', 'Nokoyawa',
        # C2 frameworks / RATs
        'Cobalt Strike', 'Mimikatz', 'Sliver', 'Brute Ratel', 'Havoc',
        'AsyncRAT', 'Remcos', 'njRAT', 'NanoCore', 'XWorm', 'QuasarRAT',
        'DarkComet', 'NetWire', 'Metasploit',
        # Loaders / droppers
        'Emotet', 'TrickBot', 'QakBot', 'IcedID', 'BazarLoader', 'Dridex',
        'GootLoader', 'BumbleBee', 'PikaBot', 'DarkGate',
        # Stealers
        'AgentTesla', 'FormBook', 'RedLine', 'Vidar', 'Raccoon', 'Lumma',
        'Rhadamanthys', 'StealC', 'Meduza', 'Aurora',
        # APT tooling
        'PlugX', 'ShadowPad', 'Gh0stRAT', 'PoisonIvy',
    }

    MITRE_PATTERN    = re.compile(r'T\d{4}(?:\.\d{3})?', re.IGNORECASE)
    CVE_PATTERN      = re.compile(r'CVE-\d{4}-\d{4,7}', re.IGNORECASE)

    def extract_cves(self, text: str) -> List[str]:
        return list(set(cve.upper() for cve in self.CVE_PATTERN.findall(text)))

    def extract_mitre_techniques(self, text: str) -> List[str]:
        return list(set(t.upper() for t in self.MITRE_PATTERN.findall(text)))

    def extract_threat_actors(self, text: str) -> List[str]:
        text_lower = text.lower()
        return list(set(
            actor.upper()
            for actor in self.THREAT_ACTORS
            if re.search(r'\b' + re.escape(actor) + r'\b', text_lower)
        ))

    def extract_malware_families(self, text: str) -> List[str]:
        text_lower = text.lower()
        return list(set(
            family
            for family in self.MALWARE_FAMILIES
            if re.search(r'\b' + re.escape(family.lower()) + r'\b', text_lower)
        ))

    def classify_article(self, article: Dict) -> Dict:
        combined_text = f"{article.get('title', '')} {article.get('description', '')}"

        cves             = self.extract_cves(combined_text)
        mitre_techniques = self.extract_mitre_techniques(combined_text)
        threat_actors    = self.extract_threat_actors(combined_text)
        malware_families = self.extract_malware_families(combined_text)

        article['cves']             = cves
        article['mitre_techniques'] = mitre_techniques
        article['threat_actors']    = threat_actors
        article['malware_families'] = malware_families

        logger.info(
            f"Classified: {article.get('title', '')[:50]}... | "
            f"CVEs: {len(cves)} | Actors: {len(threat_actors)} | Malware: {len(malware_families)}"
        )
        return article


if __name__ == "__main__":
    test_article = {
        'title': 'APT28 Deploys Cobalt Strike via CVE-2024-12345 Zero-Day in Apache Struts',
        'description': 'Fancy Bear exploited a critical RCE flaw. LockBit ransomware also observed. '
                       'Hash: e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855',
        'source': 'watchTowr Labs',
        'url': 'https://example.com',
        'category': 'research',
        'published': datetime.now(timezone.utc).isoformat(),
        'published_human': 'test',
    }
    c = ContentClassifier()
    result = c.classify_article(test_article)
    print(f"CVEs:    {result['cves']}")
    print(f"Actors:  {result['threat_actors']}")
    print(f"Malware: {result['malware_families']}")
    print(f"MITRE:   {result['mitre_techniques']}")