sameersyed/Defence_FrameWork

1import json2import re3from typing import Dict, List, Tuple4 5class SecurityAnalysisEngine:6    def __init__(self, knowledge_base_path: str):7        with open(knowledge_base_path, 'r') as f:8            self.kb = json.load(f)9    10    def analyze_intent(self, user_query: str) -> Dict:11        """Extract intent from user query using keyword matching"""12        query_lower = user_query.lower()13        14        # Detect model type15        model_type = self._detect_model_type(query_lower)16        17        # Detect application domain18        domain = self._detect_domain(query_lower)19        20        # Detect security concerns21        threats = self._detect_threats(query_lower)22        23        # Calculate confidence24        confidence = self._calculate_confidence(query_lower, threats)25        26        return {27            "model_type": model_type,28            "domain": domain,29            "threats": threats,30            "query": user_query,31            "confidence": confidence32        }33    34    def _detect_model_type(self, query: str) -> str:35        """Identify ML model type from query"""36        model_keywords = {37            "cnn": ["cnn", "convolutional", "resnet", "vgg", "image classifier"],38            "rnn": ["rnn", "lstm", "gru", "recurrent"],39            "transformer": ["transformer", "bert", "gpt", "attention", "llm"],40            "gan": ["gan", "generative", "generator", "discriminator"],41            "reinforcement": ["reinforcement", "rl", "q-learning", "policy"]42        }43        44        for model, keywords in model_keywords.items():45            if any(kw in query for kw in keywords):46                return model47        return "general"48    49    def _detect_domain(self, query: str) -> str:50        """Identify application domain"""51        domains = {52            "computer_vision": ["image", "vision", "object detection", "facial", "autonomous vehicle"],53            "nlp": ["text", "nlp", "language", "sentiment", "chatbot"],54            "healthcare": ["medical", "diagnosis", "patient", "healthcare"],55            "finance": ["financial", "fraud", "trading", "banking"],56            "iot": ["iot", "sensor", "embedded", "edge device"]57        }58        59        for domain, keywords in domains.items():60            if any(kw in query for kw in keywords):61                return domain62        return "general"63    64    def _detect_threats(self, query: str) -> List[str]:65        """Identify potential security threats from query"""66        detected = []67        68        for threat_id, threat_data in self.kb["attack_types"].items():69            keywords = threat_data["keywords"]70            scenarios = threat_data["scenarios"]71            72            # Check if any keyword or scenario matches73            if any(kw in query for kw in keywords) or any(sc in query for sc in scenarios):74                detected.append(threat_id)75        76        # If no specific threat detected, infer from context77        if not detected:78            if any(word in query for word in ["secure", "protect", "safe", "vulnerability", "risk"]):79                detected = ["adversarial_attacks", "data_poisoning"]  # Default common threats80        81        return detected82    83    def _calculate_confidence(self, query: str, threats: List[str]) -> int:84        """Calculate confidence score based on keyword matches"""85        if not threats:86            return 50  # Low confidence if no threats detected87        88        total_keywords = 089        matched_keywords = 090        91        for threat_id in threats:92            if threat_id in self.kb["attack_types"]:93                keywords = self.kb["attack_types"][threat_id]["keywords"]94                total_keywords += len(keywords)95                matched_keywords += sum(1 for kw in keywords if kw in query)96        97        if total_keywords == 0:98            return 5099        100        confidence = int((matched_keywords / total_keywords) * 100)101        return min(95, max(50, confidence))  # Clamp between 50-95%102    103    def classify_threats(self, threats: List[str]) -> List[Dict]:104        """Get detailed information about identified threats"""105        threat_details = []106        107        for threat_id in threats:108            if threat_id in self.kb["attack_types"]:109                threat = self.kb["attack_types"][threat_id]110                threat_details.append({111                    "id": threat_id,112                    "name": threat["name"],113                    "description": threat["description"],114                    "risk_level": threat["risk_level"],115                    "scenarios": threat["scenarios"]116                })117        118        return threat_details119    120    def recommend_defenses(self, threats: List[str]) -> List[Dict]:121        """Map threats to appropriate defense mechanisms"""122        defense_ids = set()123        124        for threat_id in threats:125            if threat_id in self.kb["attack_types"]:126                threat_defenses = self.kb["attack_types"][threat_id]["defenses"]127                defense_ids.update(threat_defenses)128        129        defenses = []130        for defense_id in defense_ids:131            if defense_id in self.kb["defenses"]:132                defense = self.kb["defenses"][defense_id]133                defenses.append({134                    "id": defense_id,135                    "name": defense["name"],136                    "description": defense["description"],137                    "effectiveness": defense["effectiveness"],138                    "complexity": defense["complexity"]139                })140        141        # Sort by effectiveness142        effectiveness_order = {"high": 3, "medium": 2, "low": 1}143        defenses.sort(key=lambda x: effectiveness_order.get(x["effectiveness"], 0), reverse=True)144        145        return defenses146    147    def generate_response(self, intent: Dict, threats: List[Dict], defenses: List[Dict]) -> str:148        """Generate clean, structured response without markdown symbols"""149        response_parts = []150        151        if not threats:152            return """🔒 ShieldML Security Advisor153 154I can help you analyze security risks in ML/DL systems.155 156Ask me about:157• Adversarial attacks on your model158• Data poisoning and backdoor detection159• Privacy attacks (membership inference)160• Model extraction/stealing161• Physical-world attacks162• Defense strategies and implementations163 164Example: "How to protect my CNN from adversarial attacks?"165"""166        167        # Header168        response_parts.append("🔒 Real-Time Security Analysis\n\n")169        response_parts.append(f"Analysis Time: {self._get_timestamp()}\n")170        response_parts.append(f"Domain: {intent['domain'].replace('_', ' ').title()}\n")171        response_parts.append(f"Model Type: {intent['model_type'].upper()}\n")172        response_parts.append(f"Confidence: {intent.get('confidence', 75)}%\n\n")173        response_parts.append("─" * 60 + "\n\n")174        175        # Threat summary176        critical_count = sum(1 for t in threats if t['risk_level'] == 'critical')177        high_count = sum(1 for t in threats if t['risk_level'] == 'high')178        179        if critical_count > 0:180            response_parts.append(f"🚨 ALERT: {critical_count} CRITICAL threat(s) detected\n\n")181        if high_count > 0:182            response_parts.append(f"⚠️ WARNING: {high_count} HIGH-risk threat(s) detected\n\n")183        184        # Detailed threats185        response_parts.append("⚠️ Detected Threats:\n\n")186        187        for i, threat in enumerate(threats, 1):188            risk_emoji = "🚨" if threat['risk_level'] == 'critical' else "⚠️" if threat['risk_level'] == 'high' else "⚡"189            190            response_parts.append(f"{risk_emoji} {i}. {threat['name']}\n")191            response_parts.append(f"   Risk Level: {threat['risk_level'].upper()}\n")192            response_parts.append(f"   Attack Phase: {threat.get('attack_phase', 'N/A').upper()}\n\n")193            response_parts.append(f"   Description: {threat['description']}\n\n")194            195            if threat.get('real_world_impact'):196                response_parts.append(f"   Real-World Impact: {threat['real_world_impact']}\n\n")197            198            if threat['scenarios']:199                scenarios = ', '.join(threat['scenarios'][:4])200                response_parts.append(f"   Common in: {scenarios}\n\n")201        202        response_parts.append("─" * 60 + "\n\n")203        204        # Defense recommendations205        if defenses:206            response_parts.append("🛡️ Defense Recommendations (Priority Order):\n\n")207            208            for i, defense in enumerate(defenses[:5], 1):209                eff_emoji = "🟢" if defense['effectiveness'] == 'high' else "🟡" if defense['effectiveness'] == 'medium' else "🔴"210                comp_emoji = "⚡" if defense['complexity'] == 'low' else "⚙️" if defense['complexity'] == 'medium' else "🔧"211                212                response_parts.append(f"{i}. {defense['name']} {eff_emoji}\n")213                response_parts.append(f"   Effectiveness: {defense['effectiveness'].capitalize()} | Complexity: {defense['complexity'].capitalize()} {comp_emoji}\n\n")214                response_parts.append(f"   {defense['description']}\n\n")215                216                if defense.get('implementation'):217                    response_parts.append(f"   Implementation: {defense['implementation']}\n\n")218        219        response_parts.append("─" * 60 + "\n\n")220        221        # Action items222        response_parts.append("✅ Immediate Action Items:\n\n")223        224        if critical_count > 0:225            response_parts.append("URGENT (Do Now):\n")226            response_parts.append("• Halt production deployment immediately\n")227            response_parts.append("• Audit training data for poisoning\n")228            response_parts.append("• Implement critical defenses before release\n\n")229        230        response_parts.append("Short-term (1-2 weeks):\n")231        response_parts.append("• Implement top 3 recommended defenses\n")232        response_parts.append("• Set up monitoring and logging\n")233        response_parts.append("• Test model robustness with adversarial examples\n\n")234        235        response_parts.append("Long-term (1-3 months):\n")236        response_parts.append("• Establish security testing pipeline\n")237        response_parts.append("• Regular security audits\n")238        response_parts.append("• Train team on ML security best practices\n\n")239        240        response_parts.append("─" * 60 + "\n\n")241        242        # Additional help243        response_parts.append("📚 Need More Help?\n\n")244        response_parts.append("• Upload CSV dataset for automated risk analysis\n")245        response_parts.append("• Upload images to detect spam/adversarial content\n")246        response_parts.append("• Ask follow-up questions for specific scenarios\n")247        248        return "".join(response_parts)249    250    def _get_timestamp(self):251        """Get current timestamp"""252        from datetime import datetime253        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")254