#!/usr/bin/env python3 """ APT Classification System - Version corrigée et synchronisée Correction des incohérences entre Streamlit et Gradio """ import gradio as gr import torch import torch.nn as nn from transformers import AutoTokenizer, AutoModel, AutoConfig import numpy as np import json import time from datetime import datetime import plotly.graph_objects as go import re import requests import os import io from typing import Dict, List, Optional import logging from dataclasses import dataclass # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @dataclass class ClassificationResult: predicted_class: str confidence: float top5_probabilities: Dict[str, float] processing_time: float extracted_features: Dict[str, List[str]] attribution_factors: List[str] timestamp: str class CySecBERTMaxPerformance(nn.Module): """Version EXACTEMENT identique à Streamlit""" def __init__( self, model_name: str = "markusbayer/CySecBERT", num_classes: int = 12, # ⚠️ IMPORTANT: Doit correspondre au modèle sauvegardé max_length: int = 384, dropout_rate: float = 0.15 ): super(CySecBERTMaxPerformance, self).__init__() self.model_name = model_name self.num_classes = num_classes self.max_length = max_length # CySecBERT specialized for cybersecurity self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.config = AutoConfig.from_pretrained(model_name) self.bert = AutoModel.from_pretrained(model_name) # EXPANDED architecture for maximum capacity self.dropout = nn.Dropout(dropout_rate) self.intermediate1 = nn.Linear(self.config.hidden_size, 512) self.intermediate_dropout1 = nn.Dropout(dropout_rate * 0.6) self.intermediate2 = nn.Linear(512, 256) self.intermediate_dropout2 = nn.Dropout(dropout_rate * 0.7) # Batch normalization for stability self.batch_norm1 = nn.BatchNorm1d(512) self.batch_norm2 = nn.BatchNorm1d(256) self.classifier = nn.Linear(256, num_classes) # Optimized activations self.relu = nn.ReLU() self.gelu = nn.GELU() def forward(self, input_ids, attention_mask): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) # [CLS] token with minimal dropout cls_output = outputs.last_hidden_state[:, 0] cls_output = self.dropout(cls_output) # First LARGE intermediate layer intermediate1 = self.gelu(self.intermediate1(cls_output)) intermediate1 = self.intermediate_dropout1(intermediate1) if intermediate1.size(0) > 1: intermediate1 = self.batch_norm1(intermediate1) # Second intermediate layer intermediate2 = self.relu(self.intermediate2(intermediate1)) intermediate2 = self.intermediate_dropout2(intermediate2) if intermediate2.size(0) > 1: intermediate2 = self.batch_norm2(intermediate2) # Final classification logits = self.classifier(intermediate2) return { 'logits': logits, 'probabilities': torch.softmax(logits, dim=-1) } class APTClassifier: def __init__(self): self.device = torch.device('cpu') self.model = None self.class_names = [] self.label_encoder = None # ✅ PROFILS APT IDENTIQUES À STREAMLIT (corrigés) self.apt_profiles = { 'APT1': { 'country': 'China', 'flag': '🇨🇳', 'aliases': ['Comment Crew', 'Comment Group', 'PLA Unit 61398', 'Shanghai Group'], 'description': 'Chinese cyber espionage group attributed to the People\'s Liberation Army Unit 61398. Known for large-scale intellectual property theft and targeting of over 140 organizations across 20 industries.', 'first_observed': '2006', 'attribution_confidence': 'High', 'sponsor': 'State-sponsored (PLA Unit 61398)', 'malware': ['WEBC2', 'BACKDOOR.BARKIOFORK', 'AURIGA', 'BANGAT', 'BISCUIT'], 'tools': ['HTRAN', 'GSECDUMP', 'GETMAIL', 'MAPIGET'], 'targets': ['Intellectual property', 'Government agencies', 'Industrial companies', 'Legal services', 'IT companies'], 'sectors': ['Information Technology', 'Energy', 'Financial Services', 'Government', 'Healthcare'], 'regions': ['United States', 'Canada', 'United Kingdom', 'India'], 'ttps': ['T1566.001', 'T1059.003', 'T1071.001', 'T1083', 'T1005'], 'mitre_groups': ['G0006'], 'notable_campaigns': ['Operation Aurora (2009)', 'RSA SecurID breach (2011)', 'Elderwood campaigns'], 'motivations': ['Espionage', 'Intellectual property theft'], 'sophistication': 'Medium to High' }, 'APT28': { 'country': 'Russia', 'flag': '🇷🇺', 'aliases': ['Fancy Bear', 'Sofacy', 'Sednit', 'STRONTIUM', 'Pawn Storm', 'Swallowtail'], 'description': 'Russian military intelligence cyber operations unit attributed to GRU Unit 26165. Highly sophisticated group known for targeting government, military, and security organizations worldwide.', 'first_observed': '2007', 'attribution_confidence': 'High', 'sponsor': 'State-sponsored (GRU Unit 26165)', 'malware': ['X-Agent', 'Sofacy', 'GAMEFISH', 'Zebrocy', 'CHOPSTICK', 'EVILTOSS'], 'tools': ['Responder', 'Mimikatz', 'Compiled HTML Help', 'PowerShell Empire'], 'targets': ['Government agencies', 'Military organizations', 'Defense contractors', 'Aerospace', 'Media'], 'sectors': ['Government', 'Defense', 'Aerospace', 'Media', 'Think Tanks'], 'regions': ['United States', 'Europe', 'Asia-Pacific', 'Middle East'], 'ttps': ['T1566.001', 'T1059.001', 'T1055', 'T1027', 'T1083', 'T1203'], 'mitre_groups': ['G0007'], 'notable_campaigns': ['DNC hack (2016)', 'Olympic Destroyer (2018)', 'UEFI rootkit campaigns'], 'motivations': ['Espionage', 'Political influence', 'Military intelligence'], 'sophistication': 'Very High' }, 'APT29': { 'country': 'Russia', 'flag': '🇷🇺', 'aliases': ['Cozy Bear', 'The Dukes', 'NOBELIUM', 'Midnight Blizzard', 'UNC2452'], 'description': 'Russian foreign intelligence service (SVR) cyber unit. Extremely sophisticated group known for stealth, persistence, and advanced techniques in espionage operations.', 'first_observed': '2008', 'attribution_confidence': 'High', 'sponsor': 'State-sponsored (SVR)', 'malware': ['HAMMERTOSS', 'COZYCAR', 'SeaDuke', 'SUNBURST', 'TEARDROP', 'BEACON'], 'tools': ['PowerShell', 'WMI', 'Cobalt Strike', 'AdFind', 'BloodHound'], 'targets': ['Government agencies', 'Think tanks', 'Healthcare organizations', 'Technology companies'], 'sectors': ['Government', 'Healthcare', 'Technology', 'Research', 'NGOs'], 'regions': ['United States', 'Europe', 'Global'], 'ttps': ['T1566.002', 'T1071.001', 'T1055', 'T1027', 'T1078', 'T1490'], 'mitre_groups': ['G0016'], 'notable_campaigns': ['SolarWinds supply chain attack (2020)', 'COVID-19 research targeting', 'Azure/M365 attacks'], 'motivations': ['Espionage', 'Intelligence gathering', 'Political influence'], 'sophistication': 'Very High' }, 'Lazarus': { 'country': 'North Korea', 'flag': '🇰🇵', 'aliases': ['Lazarus Group', 'Hidden Cobra', 'ZINC', 'TEMP.Hermit', 'Labyrinth Chollima'], 'description': 'North Korean state-sponsored hacking group known for financially motivated attacks, cryptocurrency theft, and destructive operations. Connected to RGB (Reconnaissance General Bureau).', 'first_observed': '2009', 'attribution_confidence': 'High', 'sponsor': 'State-sponsored (RGB)', 'malware': ['WannaCry', 'HOPLIGHT', 'TYPEFRAME', 'BADCALL', 'FALLCHILL', 'ELECTRICFISH'], 'tools': ['PowerShell', 'Mimikatz', 'PsExec', 'Living-off-the-land binaries'], 'targets': ['Financial institutions', 'Cryptocurrency exchanges', 'Entertainment companies', 'Defense contractors'], 'sectors': ['Financial Services', 'Entertainment', 'Cryptocurrency', 'Defense', 'Healthcare'], 'regions': ['Global', 'South Korea', 'United States', 'Europe'], 'ttps': ['T1566.001', 'T1059.003', 'T1055', 'T1027', 'T1486', 'T1490'], 'mitre_groups': ['G0032'], 'notable_campaigns': ['Sony Pictures attack (2014)', 'WannaCry ransomware (2017)', 'SWIFT banking attacks'], 'motivations': ['Financial gain', 'Espionage', 'Destruction', 'Sanctions evasion'], 'sophistication': 'High' }, 'Equation': { 'country': 'United States (suspected)', 'flag': '🇺🇸', 'aliases': ['Equation Group', 'EQGRP', 'Tilded Team'], 'description': 'Highly sophisticated cyber espionage group suspected to be linked to the NSA. Known for advanced persistent threats, zero-day exploits, and firmware-level implants.', 'first_observed': '2001', 'attribution_confidence': 'Medium', 'sponsor': 'State-sponsored (suspected NSA)', 'malware': ['DOUBLEFANTASY', 'EQUATIONDRUG', 'GRAYFISH', 'FANNY', 'STUXNET'], 'tools': ['EternalBlue', 'EternalRomance', 'DoublePulsar', 'FuzzBunch'], 'targets': ['High-value targets', 'Government agencies', 'Telecommunications', 'Research institutions'], 'sectors': ['Government', 'Telecommunications', 'Research', 'Technology', 'Energy'], 'regions': ['Middle East', 'Asia', 'Europe', 'Global'], 'ttps': ['T1055', 'T1027', 'T1083', 'T1068', 'T1542.009', 'T1014'], 'mitre_groups': ['G0020'], 'notable_campaigns': ['Operation Equation (2008-2015)', 'STUXNET collaboration', 'Flame malware'], 'motivations': ['Espionage', 'Intelligence gathering', 'Sabotage'], 'sophistication': 'Extremely High' }, 'Carbanak': { 'country': 'International', 'flag': '🌍', 'aliases': ['FIN7', 'Carbanak Group', 'Anunak', 'Carbon Spider'], 'description': 'Financially motivated cybercriminal organization responsible for stealing over $1 billion from financial institutions worldwide through ATM and point-of-sale attacks.', 'first_observed': '2013', 'attribution_confidence': 'High', 'sponsor': 'Cybercriminal', 'malware': ['Carbanak', 'CARBANAK', 'HALFBAKED', 'BABYMETAL', 'GRIFFON'], 'tools': ['Cobalt Strike', 'Mimikatz', 'PowerShell Empire', 'Metasploit'], 'targets': ['Financial institutions', 'Banks', 'Payment processors', 'Hospitality', 'Retail'], 'sectors': ['Financial Services', 'Hospitality', 'Retail', 'Restaurant'], 'regions': ['Global', 'United States', 'Europe', 'Asia'], 'ttps': ['T1566.001', 'T1059.003', 'T1055', 'T1027', 'T1021.001', 'T1083'], 'mitre_groups': ['G0008', 'G0046'], 'notable_campaigns': ['Carbanak banking attacks', 'FIN7 point-of-sale attacks', 'Restaurant POS campaigns'], 'motivations': ['Financial gain'], 'sophistication': 'High' }, 'APT40': { 'country': 'China', 'flag': '🇨🇳', 'aliases': ['Leviathan', 'TEMP.Periscope', 'TEMP.Jumper', 'Kryptonite Panda'], 'description': 'Chinese state-sponsored cyber espionage group focused on maritime industries, engineering companies, and research organizations to support China\'s Belt and Road Initiative.', 'first_observed': '2013', 'attribution_confidence': 'High', 'sponsor': 'State-sponsored (MSS Hainan)', 'malware': ['BADFLICK', 'PHOTO', 'HOMEFRY', 'MURKYTOP', 'LUNCHMONEY'], 'tools': ['China Chopper', 'Mimikatz', 'PowerShell', 'WMI'], 'targets': ['Maritime industries', 'Engineering companies', 'Research organizations', 'Government agencies'], 'sectors': ['Maritime', 'Engineering', 'Research', 'Government', 'Healthcare'], 'regions': ['United States', 'Europe', 'Asia-Pacific'], 'ttps': ['T1566.001', 'T1190', 'T1059.003', 'T1055', 'T1027'], 'mitre_groups': ['G0065'], 'notable_campaigns': ['Maritime industry targeting', 'COVID-19 research theft', 'Belt and Road surveillance'], 'motivations': ['Espionage', 'Economic advantage', 'Strategic intelligence'], 'sophistication': 'High' } } # Cybersecurity indicators (identiques à Streamlit) self.security_indicators = { 'malware': r'\b(trojan|virus|worm|ransomware|backdoor|rootkit|spyware|adware|botnet|rat|loader)\b', 'techniques': r'\bT\d{4}(\.\d{3})?\b', 'domains': r'\b[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b', 'ips': r'\b(?:\d{1,3}\.){3}\d{1,3}\b', 'hashes': r'\b[a-fA-F0-9]{32,64}\b', 'cve': r'\bCVE-\d{4}-\d{4,}\b', 'tools': r'\b(cobalt strike|metasploit|mimikatz|powershell|psexec|wmi|bloodhound)\b' } self.load_model() def download_model_from_hf(self): """Téléchargement robuste avec vérification du checksum""" try: model_url = "https://huggingface.co/melissachall/cysecbert-apt-classifier/resolve/main/best_cysecbert_max_performance.pt" logger.info(f"Downloading model from: {model_url}") response = requests.get(model_url, timeout=300, stream=True) if response.status_code == 200: model_path = "downloaded_model.pt" total_size = int(response.headers.get('content-length', 0)) with open(model_path, "wb") as f: downloaded = 0 for chunk in response.iter_content(chunk_size=8192): if chunk: f.write(chunk) downloaded += len(chunk) if total_size > 0: percent = (downloaded / total_size) * 100 if downloaded % 1000000 == 0: # Log every MB logger.info(f"Download progress: {percent:.1f}%") logger.info(f"✅ Model downloaded: {downloaded} bytes") # ✅ VALIDATION CRITIQUE DU MODÈLE TÉLÉCHARGÉ try: test_checkpoint = torch.load(model_path, map_location='cpu', weights_only=False) # Vérifier que les champs critiques existent required_fields = ['model_state_dict', 'class_names'] missing_fields = [field for field in required_fields if field not in test_checkpoint] if missing_fields: raise ValueError(f"Missing critical fields in checkpoint: {missing_fields}") class_names = test_checkpoint.get('class_names', []) if len(class_names) == 0: raise ValueError("Checkpoint has empty class_names") logger.info(f"✅ Model validation passed. Classes: {class_names}") return model_path except Exception as e: logger.error(f"❌ Downloaded model validation failed: {e}") if os.path.exists(model_path): os.remove(model_path) return None else: logger.error(f"❌ HTTP {response.status_code} for {model_url}") return None except Exception as e: logger.error(f"❌ Download error: {e}") return None def load_model(self): """Chargement EXACTEMENT identique à Streamlit""" try: # Étape 1: Télécharger le modèle model_path = self.download_model_from_hf() if not model_path or not os.path.exists(model_path): raise RuntimeError("❌ Cannot download model from HuggingFace") # Étape 2: Charger le checkpoint logger.info(f"Loading checkpoint from {model_path}") checkpoint = torch.load(model_path, map_location=self.device, weights_only=False) # ✅ ÉTAPE CRITIQUE: Récupérer les métadonnées EXACTEMENT comme Streamlit self.class_names = checkpoint.get('class_names', []) if not self.class_names: raise RuntimeError("❌ Checkpoint missing class_names. Upload complete .pt file with metadata.") self.label_encoder = checkpoint.get('label_encoder') num_classes = len(self.class_names) logger.info(f"✅ Class names from checkpoint: {self.class_names}") logger.info(f"✅ Number of classes: {num_classes}") # ✅ VÉRIFICATION: Les class_names doivent correspondre aux profils profile_classes = set(self.apt_profiles.keys()) checkpoint_classes = set(self.class_names) if profile_classes != checkpoint_classes: logger.warning(f"⚠️ MISMATCH DETECTED!") logger.warning(f" Profile classes: {profile_classes}") logger.warning(f" Checkpoint classes: {checkpoint_classes}") logger.warning(f" Missing in profiles: {checkpoint_classes - profile_classes}") logger.warning(f" Extra in profiles: {profile_classes - checkpoint_classes}") # Étape 3: Créer le modèle avec les bonnes dimensions self.model = CySecBERTMaxPerformance( num_classes=num_classes, dropout_rate=checkpoint.get('config', {}).get('dropout_rate', 0.15) ).to(self.device) # Étape 4: Charger les poids if 'model_state_dict' not in checkpoint: raise RuntimeError("❌ Checkpoint missing model_state_dict") self.model.load_state_dict(checkpoint['model_state_dict']) self.model.eval() logger.info("✅ MODEL LOADED SUCCESSFULLY - IDENTICAL TO STREAMLIT!") # Nettoyage if os.path.exists(model_path): os.remove(model_path) return True except Exception as e: logger.error(f"❌ Model loading error: {e}") raise RuntimeError(f"Cannot load model: {e}") def extract_features(self, text: str) -> Dict[str, List[str]]: """IDENTIQUE à Streamlit""" features = {} text_lower = text.lower() for feature_type, pattern in self.security_indicators.items(): matches = re.findall(pattern, text_lower, re.IGNORECASE) features[feature_type] = list(set(matches))[:10] # Limit to 10 items return features def get_attribution_factors(self, text: str, predicted_class: str) -> List[str]: """IDENTIQUE à Streamlit""" factors = [] text_lower = text.lower() if predicted_class in self.apt_profiles: profile = self.apt_profiles[predicted_class] # Check for group mentions if predicted_class.lower() in text_lower: factors.append(f"Direct mention of {predicted_class}") # Check for aliases for alias in profile.get('aliases', []): if alias.lower() in text_lower: factors.append(f"Alias detected: {alias}") # Check for known malware for malware in profile.get('malware', []): if malware.lower() in text_lower: factors.append(f"Known malware: {malware}") # Check for tools for tool in profile.get('tools', []): if tool.lower() in text_lower: factors.append(f"Known tool: {tool}") # Check for target sectors for target in profile.get('targets', []): if target.lower() in text_lower: factors.append(f"Target sector match: {target}") # Check for TTPs for ttp in profile.get('ttps', []): if ttp in text: factors.append(f"MITRE technique: {ttp}") return factors def classify(self, text: str, confidence_threshold: float = 0.5) -> ClassificationResult: """Classification EXACTEMENT identique à Streamlit""" start_time = time.time() # Vérifications strictes if self.model is None: raise RuntimeError("❌ Model not loaded") if not hasattr(self.model, 'tokenizer') or self.model.tokenizer is None: raise RuntimeError("❌ Tokenizer not loaded") logger.info("🚀 Using CySecBERTMaxPerformance (identical to Streamlit)") # Tokenisation IDENTIQUE encoding = self.model.tokenizer( text, max_length=self.model.max_length, padding='max_length', truncation=True, return_tensors='pt' ) input_ids = encoding['input_ids'].to(self.device) attention_mask = encoding['attention_mask'].to(self.device) # Prédiction with torch.no_grad(): outputs = self.model(input_ids, attention_mask) probabilities = outputs['probabilities'].cpu().numpy()[0] # Top 5 IDENTIQUE top5_indices = np.argsort(probabilities)[::-1][:5] predicted_class = self.class_names[top5_indices[0]] confidence = float(probabilities[top5_indices[0]]) # Distribution top 5 top5_probabilities = { self.class_names[idx]: float(probabilities[idx]) for idx in top5_indices } logger.info(f"✅ Prediction: {predicted_class} ({confidence:.1%})") logger.info(f"✅ Top 5: {top5_probabilities}") # Features et attribution extracted_features = self.extract_features(text) attribution_factors = self.get_attribution_factors(text, predicted_class) processing_time = time.time() - start_time return ClassificationResult( predicted_class=predicted_class, confidence=confidence, top5_probabilities=top5_probabilities, processing_time=processing_time, extracted_features=extracted_features, attribution_factors=attribution_factors, timestamp=datetime.now().isoformat() ) # ===== FONCTIONS UTILITAIRES IDENTIQUES ===== def process_uploaded_file(uploaded_file): """Process uploaded file and extract text content""" if uploaded_file is None: return "" try: file_name = uploaded_file.name.lower() if file_name.endswith('.txt'): content = uploaded_file.read() if isinstance(content, bytes): return content.decode('utf-8', errors='ignore') return str(content) elif file_name.endswith('.json'): content = uploaded_file.read() if isinstance(content, bytes): content = content.decode('utf-8') try: json_data = json.loads(content) text_fields = [] def extract_text_from_json(obj, depth=0): if depth > 3: return if isinstance(obj, dict): for key, value in obj.items(): if isinstance(value, str) and len(value) > 10: text_fields.append(f"{key}: {value}") elif isinstance(value, (dict, list)): extract_text_from_json(value, depth + 1) elif isinstance(obj, list): for item in obj: extract_text_from_json(item, depth + 1) extract_text_from_json(json_data) return "\n".join(text_fields) except: return content else: # Generic text extraction content = uploaded_file.read() if isinstance(content, bytes): return content.decode('utf-8', errors='ignore') return str(content) except Exception as e: logger.error(f"File processing error: {e}") return f"Error processing file: {str(e)}" def create_prediction_plot(top5_probs): """Créer le graphique des top 5 prédictions""" fig = go.Figure(go.Bar( x=list(top5_probs.values()), y=list(top5_probs.keys()), orientation='h', marker=dict( color=['#667eea', '#764ba2', '#f093fb', '#f5576c', '#4facfe'][:len(top5_probs)], line=dict(color='rgba(50,50,50,0.8)', width=1) ), text=[f"{prob:.2%}" for prob in top5_probs.values()], textposition='auto', textfont=dict(size=12, color='white') )) fig.update_layout( title=dict( text="🎯 Top 5 APT Group Predictions", font=dict(size=18, color='#2c3e50'), x=0.5 ), xaxis=dict( title=dict(text="Confidence Score", font=dict(size=14)), tickfont=dict(size=12), range=[0, max(top5_probs.values()) * 1.1] ), yaxis=dict( title=dict(text="APT Groups", font=dict(size=14)), tickfont=dict(size=12) ), height=400, margin=dict(l=100, r=50, t=80, b=50), plot_bgcolor='rgba(248,249,250,0.8)', paper_bgcolor='white' ) return fig def format_apt_profile(predicted_class, classifier): """Formater le profil APT (IDENTIQUE à Streamlit)""" if predicted_class not in classifier.apt_profiles: return f"
⚠️ No profile available for '{predicted_class}'. This might indicate a class mapping issue.
" profile = classifier.apt_profiles[predicted_class] html = f"""

{profile.get('flag', '🌍')} {predicted_class} - Complete Profile

📋 Basic Information

Origin: {profile.get('country', 'Unknown')}

First Observed: {profile.get('first_observed', 'Unknown')}

Attribution Confidence: {profile.get('attribution_confidence', 'Unknown')}

Sponsor: {profile.get('sponsor', 'Unknown')}

Sophistication: {profile.get('sophistication', 'Unknown')}

🎭 Known Aliases

    {''.join([f"
  • {alias}
    • " for alias in profile.get('aliases', [])])}

🦠 Associated Malware

{''.join([f"{malware}" for malware in profile.get('malware', [])])}

🎯 Typical Targets

    {''.join([f"
  • {target}
    • " for target in profile.get('targets', [])])}

🛠️ Known Tools

{''.join([f"{tool}" for tool in profile.get('tools', [])])}

⚙️ MITRE ATT&CK TTPs

{''.join([f"{ttp}" for ttp in profile.get('ttps', [])])}

📖 Description

{profile.get('description', 'No description available')}

🚨 Notable Campaigns

""" return html def classify_text(text, uploaded_file, confidence_threshold, show_features, show_attribution, show_profile): """Fonction principale de classification - IDENTIQUE à Streamlit""" # Traitement de l'input input_text = "" file_processed = False if uploaded_file is not None: file_content = process_uploaded_file(uploaded_file) if file_content.strip(): input_text = file_content file_processed = True if not input_text and text.strip(): input_text = text if not input_text: return ( "Please enter text or upload a file", "No confidence", None, "0.000s", "No input provided", "No features extracted" if show_features else "", "No attribution factors" if show_attribution else "", "No profile available" if show_profile else "", {}, "File processed successfully" if file_processed else "" ) # Classification try: classifier = APTClassifier() result = classifier.classify(input_text, confidence_threshold) # Formatage des résultats confidence_color = "#27ae60" if result.confidence > 0.8 else "#f39c12" if result.confidence > 0.6 else "#e74c3c" main_result = f"""

🛡️ APT Classification Result

✅ Model synchronized with Streamlit version

Predicted Group

{result.predicted_class}

Confidence

{result.confidence:.2%}

Processing Time

{result.processing_time:.3f}s

Analysis completed at: {datetime.fromisoformat(result.timestamp).strftime('%H:%M:%S UTC')}

""" # Graphique plot = create_prediction_plot(result.top5_probabilities) # Features features_html = "" if show_features and any(result.extracted_features.values()): features_html = "
" features_html += "

🔍 Extracted Cybersecurity Features

" icon_map = {'malware': '🦠', 'techniques': '⚙️', 'domains': '🌐', 'ips': '🔢', 'hashes': '#️⃣', 'cve': '🚨', 'tools': '🛠️'} for feature_type, feature_list in result.extracted_features.items(): if feature_list: icon = icon_map.get(feature_type, '📌') features_html += f"

{icon} {feature_type.title()}:

" for feature in feature_list[:5]: features_html += f"{feature}" features_html += "
" # Attribution attribution_html = "" if show_attribution and result.attribution_factors: attribution_html = "
" attribution_html += "

🎯 Attribution Factors

" for factor in result.attribution_factors: attribution_html += f"
{factor}
" attribution_html += "
" # Profil profile_html = "" if show_profile: profile_html = format_apt_profile(result.predicted_class, classifier) # Export data export_data = { 'predicted_class': result.predicted_class, 'confidence': result.confidence, 'processing_time': result.processing_time, 'top5_probabilities': result.top5_probabilities, 'extracted_features': result.extracted_features, 'attribution_factors': result.attribution_factors, 'timestamp': result.timestamp, 'model_info': 'CySecBERTMaxPerformance - Synchronized with Streamlit' } file_status = f"✅ File '{uploaded_file.name}' processed successfully ({len(input_text)} characters)" if file_processed else "" return ( result.predicted_class, f"{result.confidence:.2%}", plot, f"{result.processing_time:.3f}s", main_result, features_html, attribution_html, profile_html, export_data, file_status ) except Exception as e: error_msg = f"❌ Classification error: {str(e)}" logger.error(error_msg) return ( "Error", "0%", None, "0.000s", f"
{error_msg}
", "", "", "", {}, "" ) # ===== INTERFACE GRADIO ===== # CSS optimisé css = """ @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap'); .gradio-container { font-family: 'Inter', sans-serif !important; max-width: 1200px !important; margin: 0 auto !important; } .header-container { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 3rem 2rem; border-radius: 20px; text-align: center; margin-bottom: 2rem; box-shadow: 0 15px 35px rgba(102, 126, 234, 0.3); } .header-title { color: white; font-size: 3.5rem; font-weight: 700; margin: 0; text-shadow: 2px 2px 4px rgba(0,0,0,0.3); } .header-subtitle { color: rgba(255,255,255,0.95); font-size: 1.3rem; font-weight: 400; margin: 1rem 0 0 0; } .status-success { background: linear-gradient(135deg, #28a745 0%, #20c997 100%); color: white; padding: 1rem 2rem; border-radius: 10px; margin: 1rem 0; text-align: center; font-weight: 500; box-shadow: 0 4px 15px rgba(40, 167, 69, 0.3); } """ # Examples optimisés pour les 7 groupes de Streamlit example_texts = [ "Advanced persistent threat campaign attributed to APT28 (Fancy Bear) targeting government entities. Spear-phishing emails deliver X-Agent payload with T1566.001 techniques. Network analysis reveals C2 communications consistent with Sofacy operations and GRU Unit 26165 characteristics.", "Financial institutions targeted by Lazarus Group operations. Watering hole attacks deploy custom malware for cryptocurrency theft. TTPs include T1566.001 and T1059.003, consistent with Hidden Cobra methodologies. HOPLIGHT and TYPEFRAME malware observed.", "Government agencies report sophisticated malware attributed to APT29 (Cozy Bear). Advanced T1566.002 techniques with HAMMERTOSS C2. Campaign characteristics consistent with Russian SVR operations. NOBELIUM techniques with SUNBURST and TEARDROP malware detected.", "Chinese cyber espionage operations attributed to APT1 targeting intellectual property. Comment Crew techniques observed with WEBC2 and BACKDOOR.BARKIOFORK malware. Campaign consistent with PLA Unit 61398 operations targeting industrial companies.", "Advanced threat operations attributed to Equation Group. Zero-day exploits and firmware-level implants detected. DOUBLEFANTASY and EQUATIONDRUG malware with characteristics consistent with NSA-linked operations." ] # Interface Gradio principale with gr.Blocks(theme=gr.themes.Soft(), css=css, title="APT Classification System - Fixed") as demo: # Header gr.HTML("""

🛡️ APT Classification System

CySecBERTMaxPerformance - Synchronized with Streamlit

🔄 Fixed version - Identical behavior to Streamlit interface

""") # Status gr.HTML("""
✅ SYNCHRONIZED MODEL • Fixed class mapping • Identical to Streamlit version
""") with gr.Row(): # Main input column with gr.Column(scale=2): gr.Markdown("### 📝 Threat Intelligence Input") with gr.Tab("Text Input"): text_input = gr.Textbox( lines=6, placeholder="Describe the cybersecurity incident, including TTPs, malware, targets, and attribution indicators...", label="Incident Description", show_label=False ) gr.Examples( examples=[[text] for text in example_texts], inputs=text_input, label="📚 Test Cases for APT Groups" ) with gr.Tab("File Upload"): file_input = gr.File( file_types=[".txt", ".log", ".json", ".csv",".pdf"], label="Upload Threat Intelligence Report", file_count="single" ) file_status = gr.HTML("") # Configuration column with gr.Column(scale=1): gr.Markdown("### ⚙️ Analysis Configuration") confidence_threshold = gr.Slider( 0.0, 1.0, value=0.3, step=0.05, label="🎯 Confidence Threshold", info="Minimum confidence for predictions" ) gr.Markdown("### 📊 Display Options") show_features = gr.Checkbox(value=True, label="🔍 Extract Features") show_attribution = gr.Checkbox(value=True, label="🎯 Show Attribution") show_profile = gr.Checkbox(value=True, label="📋 Complete Profile") analyze_button = gr.Button( "🔍 ANALYZE THREAT", variant="primary", size="lg" ) gr.Markdown("### 📈 Model Information") gr.HTML("""
Status: ✅ Fixed & Synchronized
Architecture: CySecBERT + Custom Layers
Classes: 7 APT groups (same as Streamlit)
Source: HuggingFace Hub download
Validation: Class mapping verified
""") # Results section gr.Markdown("## 📊 Analysis Results") with gr.Row(): predicted_class = gr.Textbox(label="🎯 Predicted APT Group", interactive=False) confidence_display = gr.Textbox(label="📊 Confidence Score", interactive=False) processing_time = gr.Textbox(label="⚡ Processing Time", interactive=False) # Main results main_result = gr.HTML(label="Main Results") # Visualization prediction_plot = gr.Plot(label="Top 5 Predictions Visualization") # Additional results with gr.Row(): features_output = gr.HTML(label="Extracted Features") attribution_output = gr.HTML(label="Attribution Factors") # APT Profile profile_output = gr.HTML(label="Complete APT Profile") # Export gr.Markdown("### 💾 Export Results") export_data = gr.JSON(label="Analysis Data", visible=False) # Event handler analyze_button.click( classify_text, inputs=[text_input, file_input, confidence_threshold, show_features, show_attribution, show_profile], outputs=[predicted_class, confidence_display, prediction_plot, processing_time, main_result, features_output, attribution_output, profile_output, export_data, file_status] ) if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, share=False, show_error=True )