Create security_system.py

src/security_system.py

@@ -0,0 +1,87 @@
+import re
+import time
+from datetime import datetime, timedelta
+from typing import Dict, List, Optional
+import logging
+
+class SecuritySystem:
+    """Security system for input validation and rate limiting"""
+    
+    def __init__(self):
+        self.rate_limits = {}
+        self.suspicious_patterns = [
+            r"(?i)(password|token|key|secret)",
+            r"(?i)(delete|drop|alter|update).*table",
+            r"(?i)(script|javascript|onload|onerror)",
+            r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b",
+            r"(?i)(admin|root|sudo)"
+        ]
+        self.setup_logging()
+    
+    def setup_logging(self):
+        """Setup logging"""
+        self.logger = logging.getLogger(__name__)
+    
+    def check_request(self, query: str, user_id: str) -> Dict:
+        """Check request for security issues"""
+        result = {
+            "is_suspicious": False,
+            "alerts": [],
+            "risk_score": 0
+        }
+        
+        # Rate limiting
+        if not self.check_rate_limit(user_id):
+            result["is_suspicious"] = True
+            result["alerts"].append("Rate limit exceeded")
+            result["risk_score"] = 100
+            return result
+        
+        # Pattern matching
+        for pattern in self.suspicious_patterns:
+            if re.search(pattern, query):
+                result["alerts"].append(f"Suspicious pattern: {pattern}")
+                result["risk_score"] += 20
+        
+        # Query length analysis
+        if len(query) > 10000:
+            result["alerts"].append("Excessively long query")
+            result["risk_score"] += 30
+        
+        # Special character analysis
+        special_chars = len(re.findall(r'[^\w\s]', query))
+        if special_chars > len(query) * 0.3:
+            result["alerts"].append("High percentage of special characters")
+            result["risk_score"] += 25
+        
+        # Determine if suspicious
+        if result["risk_score"] >= 50:
+            result["is_suspicious"] = True
+        
+        return result
+    
+    def check_rate_limit(self, user_id: str, requests_per_minute: int = 60) -> bool:
+        """Check rate limit for user"""
+        current_time = datetime.now()
+        
+        # Clean old entries
+        self.rate_limits[user_id] = [
+            t for t in self.rate_limits.get(user_id, [])
+            if current_time - t < timedelta(minutes=1)
+        ]
+        
+        # Check rate limit
+        if len(self.rate_limits[user_id]) >= requests_per_minute:
+            return False
+        
+        # Add current request
+        self.rate_limits[user_id].append(current_time)
+        return True
+    
+    def sanitize_input(self, text: str) -> str:
+        """Sanitize user input"""
+        # Remove potentially dangerous characters
+        sanitized = re.sub(r'[<>"\'&]', '', text)
+        sanitized = re.sub(r'(\b)(DROP|DELETE|INSERT|UPDATE)(\b)', '', sanitized, flags=re.IGNORECASE)
+        sanitized = re.sub(r';\s*\w+', '', sanitized)
+        return sanitized.strip()