CHKIM79/scalable-ai-agent-system

1"""2Security and Privacy Module3Implements encryption, authentication, compliance, audit logging, and privacy protection4"""5import asyncio6import logging7import json8import hashlib9import hmac10import secrets11import base6412from typing import Dict, List, Any, Optional, Union, Callable13from dataclasses import dataclass, field14from enum import Enum15from datetime import datetime, timedelta16from cryptography.fernet import Fernet17import sqlite318import jwt19from passlib.context import CryptContext20import re21 22 23class SecurityLevel(Enum):24    LOW = "low"25    MEDIUM = "medium"26    HIGH = "high"27    CRITICAL = "critical"28 29 30class AuthenticationMethod(Enum):31    PASSWORD = "password"32    API_KEY = "api_key"33    JWT_TOKEN = "jwt_token"34    OAUTH2 = "oauth2"35    MULTI_FACTOR = "multi_factor"36 37 38class ComplianceStandard(Enum):39    GDPR = "gdpr"40    HIPAA = "hipaa"41    SOC2 = "soc2"42    PCI_DSS = "pci_dss"43    ISO27001 = "iso27001"44 45 46class AuditEventType(Enum):47    LOGIN = "login"48    LOGOUT = "logout"49    DATA_ACCESS = "data_access"50    DATA_MODIFICATION = "data_modification"51    PERMISSION_CHANGE = "permission_change"52    SECURITY_VIOLATION = "security_violation"53    SYSTEM_ERROR = "system_error"54 55 56@dataclass57class User:58    id: str59    username: str60    email: str61    password_hash: str62    roles: List[str] = field(default_factory=list)63    permissions: List[str] = field(default_factory=list)64    is_active: bool = True65    created_at: datetime = field(default_factory=datetime.now)66    last_login: Optional[datetime] = None67    failed_login_attempts: int = 068    metadata: Dict[str, Any] = field(default_factory=dict)69 70 71@dataclass72class APIKey:73    id: str74    key_hash: str75    name: str76    user_id: str77    permissions: List[str] = field(default_factory=list)78    expires_at: Optional[datetime] = None79    is_active: bool = True80    created_at: datetime = field(default_factory=datetime.now)81    last_used: Optional[datetime] = None82    usage_count: int = 083 84 85@dataclass86class AuditEvent:87    id: str88    event_type: AuditEventType89    user_id: Optional[str]90    resource: str91    action: str92    details: Dict[str, Any]93    ip_address: Optional[str] = None94    user_agent: Optional[str] = None95    timestamp: datetime = field(default_factory=datetime.now)96    risk_score: float = 0.097 98 99class EncryptionManager:100    """Handles encryption and decryption operations"""101    102    def __init__(self, master_key: Optional[bytes] = None):103        self.master_key = master_key or Fernet.generate_key()104        self.cipher_suite = Fernet(self.master_key)105        self.password_context = CryptContext(schemes=["bcrypt"], deprecated="auto")106        self.logger = logging.getLogger(__name__)107    108    def encrypt_data(self, data: Union[str, bytes]) -> str:109        """Encrypt data using symmetric encryption"""110        if isinstance(data, str):111            data = data.encode()112        113        encrypted_data = self.cipher_suite.encrypt(data)114        return base64.b64encode(encrypted_data).decode()115    116    def decrypt_data(self, encrypted_data: str) -> str:117        """Decrypt data using symmetric encryption"""118        try:119            encrypted_bytes = base64.b64decode(encrypted_data.encode())120            decrypted_data = self.cipher_suite.decrypt(encrypted_bytes)121            return decrypted_data.decode()122        except Exception as e:123            self.logger.error(f"Decryption failed: {e}")124            raise ValueError("Invalid encrypted data")125    126    def hash_password(self, password: str) -> str:127        """Hash password using bcrypt"""128        return self.password_context.hash(password)129    130    def verify_password(self, password: str, password_hash: str) -> bool:131        """Verify password against hash"""132        return self.password_context.verify(password, password_hash)133    134    def generate_api_key(self) -> str:135        """Generate secure API key"""136        return secrets.token_urlsafe(32)137    138    def hash_api_key(self, api_key: str) -> str:139        """Hash API key for storage"""140        return hashlib.sha256(api_key.encode()).hexdigest()141    142    def generate_jwt_secret(self) -> str:143        """Generate JWT secret key"""144        return secrets.token_urlsafe(64)145 146 147class AuthenticationManager:148    """Manages user authentication and authorization"""149    150    def __init__(self, db_path: str = "security.db", jwt_secret: Optional[str] = None):151        self.db_path = db_path152        self.jwt_secret = jwt_secret or secrets.token_urlsafe(64)153        self.encryption_manager = EncryptionManager()154        155        self.users: Dict[str, User] = {}156        self.api_keys: Dict[str, APIKey] = {}157        self.active_sessions: Dict[str, Dict[str, Any]] = {}158        159        self.db_connection = None160        self.logger = logging.getLogger(__name__)161        162        # Security settings163        self.max_failed_attempts = 5164        self.lockout_duration = timedelta(minutes=30)165        self.jwt_expiration = timedelta(hours=24)166        self.api_key_expiration = timedelta(days=90)167    168    async def initialize(self):169        """Initialize authentication manager"""170        self.db_connection = sqlite3.connect(self.db_path)171        await self._create_tables()172        await self._load_users()173        await self._load_api_keys()174        175        self.logger.info("Authentication manager initialized")176    177    async def _create_tables(self):178        """Create authentication tables"""179        cursor = self.db_connection.cursor()180        181        cursor.execute('''182            CREATE TABLE IF NOT EXISTS users (183                id TEXT PRIMARY KEY,184                username TEXT UNIQUE,185                email TEXT UNIQUE,186                password_hash TEXT,187                roles TEXT,188                permissions TEXT,189                is_active BOOLEAN,190                created_at DATETIME,191                last_login DATETIME,192                failed_login_attempts INTEGER,193                metadata TEXT194            )195        ''')196        197        cursor.execute('''198            CREATE TABLE IF NOT EXISTS api_keys (199                id TEXT PRIMARY KEY,200                key_hash TEXT UNIQUE,201                name TEXT,202                user_id TEXT,203                permissions TEXT,204                expires_at DATETIME,205                is_active BOOLEAN,206                created_at DATETIME,207                last_used DATETIME,208                usage_count INTEGER,209                FOREIGN KEY (user_id) REFERENCES users (id)210            )211        ''')212        213        self.db_connection.commit()214    215    async def _load_users(self):216        """Load users from database"""217        cursor = self.db_connection.cursor()218        cursor.execute('SELECT * FROM users WHERE is_active = 1')219        220        for row in cursor.fetchall():221            user = User(222                id=row[0],223                username=row[1],224                email=row[2],225                password_hash=row[3],226                roles=json.loads(row[4]) if row[4] else [],227                permissions=json.loads(row[5]) if row[5] else [],228                is_active=bool(row[6]),229                created_at=datetime.fromisoformat(row[7]),230                last_login=datetime.fromisoformat(row[8]) if row[8] else None,231                failed_login_attempts=row[9],232                metadata=json.loads(row[10]) if row[10] else {}233            )234            self.users[user.id] = user235    236    async def _load_api_keys(self):237        """Load API keys from database"""238        cursor = self.db_connection.cursor()239        cursor.execute('SELECT * FROM api_keys WHERE is_active = 1')240        241        for row in cursor.fetchall():242            api_key = APIKey(243                id=row[0],244                key_hash=row[1],245                name=row[2],246                user_id=row[3],247                permissions=json.loads(row[4]) if row[4] else [],248                expires_at=datetime.fromisoformat(row[5]) if row[5] else None,249                is_active=bool(row[6]),250                created_at=datetime.fromisoformat(row[7]),251                last_used=datetime.fromisoformat(row[8]) if row[8] else None,252                usage_count=row[9]253            )254            self.api_keys[api_key.key_hash] = api_key255    256    async def create_user(self, username: str, email: str, password: str, 257                         roles: List[str] = None, permissions: List[str] = None) -> User:258        """Create a new user"""259        260        # Validate input261        if not self._validate_email(email):262            raise ValueError("Invalid email format")263        264        if not self._validate_password(password):265            raise ValueError("Password does not meet security requirements")266        267        # Check for existing user268        cursor = self.db_connection.cursor()269        cursor.execute('SELECT id FROM users WHERE username = ? OR email = ?', (username, email))270        if cursor.fetchone():271            raise ValueError("User already exists")272        273        # Create user274        user = User(275            id=f"user_{secrets.token_hex(8)}",276            username=username,277            email=email,278            password_hash=self.encryption_manager.hash_password(password),279            roles=roles or [],280            permissions=permissions or []281        )282        283        # Store in database284        cursor.execute('''285            INSERT INTO users 286            (id, username, email, password_hash, roles, permissions, is_active, created_at, failed_login_attempts, metadata)287            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)288        ''', (289            user.id, user.username, user.email, user.password_hash,290            json.dumps(user.roles), json.dumps(user.permissions),291            user.is_active, user.created_at, user.failed_login_attempts,292            json.dumps(user.metadata)293        ))294        self.db_connection.commit()295        296        self.users[user.id] = user297        self.logger.info(f"Created user: {username}")298        299        return user300    301    async def authenticate_user(self, username: str, password: str, 302                              ip_address: str = None, user_agent: str = None) -> Optional[str]:303        """Authenticate user and return JWT token"""304        305        # Find user306        user = None307        for u in self.users.values():308            if u.username == username or u.email == username:309                user = u310                break311        312        if not user or not user.is_active:313            return None314        315        # Check lockout316        if user.failed_login_attempts >= self.max_failed_attempts:317            lockout_end = user.last_login + self.lockout_duration if user.last_login else datetime.now()318            if datetime.now() < lockout_end:319                self.logger.warning(f"User {username} is locked out")320                return None321        322        # Verify password323        if not self.encryption_manager.verify_password(password, user.password_hash):324            user.failed_login_attempts += 1325            await self._update_user(user)326            self.logger.warning(f"Failed login attempt for user: {username}")327            return None328        329        # Successful login330        user.failed_login_attempts = 0331        user.last_login = datetime.now()332        await self._update_user(user)333        334        # Generate JWT token335        token = self._generate_jwt_token(user)336        337        self.logger.info(f"User authenticated: {username}")338        return token339    340    async def verify_jwt_token(self, token: str) -> Optional[User]:341        """Verify JWT token and return user"""342        343        try:344            payload = jwt.decode(token, self.jwt_secret, algorithms=["HS256"])345            user_id = payload.get('user_id')346            347            if user_id and user_id in self.users:348                user = self.users[user_id]349                if user.is_active:350                    return user351            352        except jwt.InvalidTokenError:353            self.logger.warning("Invalid JWT token")354        355        return None356    357    def _generate_jwt_token(self, user: User) -> str:358        """Generate JWT token for user"""359        360        payload = {361            'user_id': user.id,362            'username': user.username,363            'roles': user.roles,364            'permissions': user.permissions,365            'exp': datetime.utcnow() + self.jwt_expiration,366            'iat': datetime.utcnow()367        }368        369        return jwt.encode(payload, self.jwt_secret, algorithm="HS256")370    371    def check_permission(self, user: User, permission: str) -> bool:372        """Check if user has specific permission"""373        374        if not user or not user.is_active:375            return False376        377        # Check direct permissions378        if permission in user.permissions:379            return True380        381        # Check role-based permissions (simplified)382        role_permissions = {383            'admin': ['*'],  # Admin has all permissions384            'user': ['read', 'write'],385            'readonly': ['read']386        }387        388        for role in user.roles:389            if role in role_permissions:390                role_perms = role_permissions[role]391                if '*' in role_perms or permission in role_perms:392                    return True393        394        return False395    396    def _validate_email(self, email: str) -> bool:397        """Validate email format"""398        pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'399        return re.match(pattern, email) is not None400    401    def _validate_password(self, password: str) -> bool:402        """Validate password strength"""403        # At least 8 characters, one uppercase, one lowercase, one digit404        if len(password) < 8:405            return False406        407        has_upper = any(c.isupper() for c in password)408        has_lower = any(c.islower() for c in password)409        has_digit = any(c.isdigit() for c in password)410        411        return has_upper and has_lower and has_digit412    413    async def _update_user(self, user: User):414        """Update user in database"""415        cursor = self.db_connection.cursor()416        cursor.execute('''417            UPDATE users SET last_login = ?, failed_login_attempts = ?418            WHERE id = ?419        ''', (user.last_login, user.failed_login_attempts, user.id))420        self.db_connection.commit()421 422 423class AuditLogger:424    """Handles security audit logging"""425    426    def __init__(self, db_path: str = "audit.db"):427        self.db_path = db_path428        self.db_connection = None429        self.logger = logging.getLogger(__name__)430        431        # Risk scoring rules432        self.risk_rules = {433            AuditEventType.LOGIN: 1.0,434            AuditEventType.LOGOUT: 0.5,435            AuditEventType.DATA_ACCESS: 2.0,436            AuditEventType.DATA_MODIFICATION: 3.0,437            AuditEventType.PERMISSION_CHANGE: 4.0,438            AuditEventType.SECURITY_VIOLATION: 5.0,439            AuditEventType.SYSTEM_ERROR: 2.5440        }441    442    async def initialize(self):443        """Initialize audit logger"""444        self.db_connection = sqlite3.connect(self.db_path)445        await self._create_tables()446        447        self.logger.info("Audit logger initialized")448    449    async def _create_tables(self):450        """Create audit tables"""451        cursor = self.db_connection.cursor()452        453        cursor.execute('''454            CREATE TABLE IF NOT EXISTS audit_events (455                id TEXT PRIMARY KEY,456                event_type TEXT,457                user_id TEXT,458                resource TEXT,459                action TEXT,460                details TEXT,461                ip_address TEXT,462                user_agent TEXT,463                timestamp DATETIME,464                risk_score REAL465            )466        ''')467        468        cursor.execute('''469            CREATE TABLE IF NOT EXISTS security_alerts (470                id INTEGER PRIMARY KEY AUTOINCREMENT,471                alert_type TEXT,472                severity TEXT,473                message TEXT,474                details TEXT,475                timestamp DATETIME,476                resolved BOOLEAN DEFAULT FALSE477            )478        ''')479        480        self.db_connection.commit()481    482    async def log_event(self, event_type: AuditEventType, user_id: Optional[str],483                       resource: str, action: str, details: Dict[str, Any],484                       ip_address: str = None, user_agent: str = None) -> AuditEvent:485        """Log audit event"""486        487        event = AuditEvent(488            id=f"audit_{secrets.token_hex(8)}",489            event_type=event_type,490            user_id=user_id,491            resource=resource,492            action=action,493            details=details,494            ip_address=ip_address,495            user_agent=user_agent,496            risk_score=self._calculate_risk_score(event_type, details)497        )498        499        # Store in database500        cursor = self.db_connection.cursor()501        cursor.execute('''502            INSERT INTO audit_events 503            (id, event_type, user_id, resource, action, details, ip_address, user_agent, timestamp, risk_score)504            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)505        ''', (506            event.id, event.event_type.value, event.user_id, event.resource,507            event.action, json.dumps(event.details), event.ip_address,508            event.user_agent, event.timestamp, event.risk_score509        ))510        self.db_connection.commit()511        512        # Check for security alerts513        await self._check_security_alerts(event)514        515        return event516    517    def _calculate_risk_score(self, event_type: AuditEventType, details: Dict[str, Any]) -> float:518        """Calculate risk score for event"""519        520        base_score = self.risk_rules.get(event_type, 1.0)521        522        # Adjust based on details523        multiplier = 1.0524        525        if details.get('failed_attempt'):526            multiplier += 0.5527        528        if details.get('suspicious_ip'):529            multiplier += 1.0530        531        if details.get('privilege_escalation'):532            multiplier += 2.0533        534        return min(base_score * multiplier, 10.0)  # Cap at 10535    536    async def _check_security_alerts(self, event: AuditEvent):537        """Check if event should trigger security alerts"""538        539        # High risk score alert540        if event.risk_score >= 4.0:541            await self._create_alert(542                "HIGH_RISK_EVENT",543                "high",544                f"High risk event detected: {event.event_type.value}",545                event.details546            )547    548    async def _create_alert(self, alert_type: str, severity: str, message: str, details: Dict[str, Any]):549        """Create security alert"""550        cursor = self.db_connection.cursor()551        552        cursor.execute('''553            INSERT INTO security_alerts (alert_type, severity, message, details, timestamp)554            VALUES (?, ?, ?, ?, ?)555        ''', (alert_type, severity, message, json.dumps(details), datetime.now()))556        557        self.db_connection.commit()558        559        self.logger.warning(f"Security alert: {alert_type} - {message}")560 561 562class SecurityManager:563    """Main security manager coordinating all security components"""564    565    def __init__(self, db_path: str = "security.db"):566        self.db_path = db_path567        self.encryption_manager = EncryptionManager()568        self.auth_manager = AuthenticationManager(db_path)569        self.audit_logger = AuditLogger(db_path.replace('.db', '_audit.db'))570        571        self.logger = logging.getLogger(__name__)572    573    async def initialize(self):574        """Initialize security manager"""575        await self.auth_manager.initialize()576        await self.audit_logger.initialize()577        578        self.logger.info("Security manager initialized")579    580    async def authenticate_request(self, token: str = None, api_key: str = None) -> Optional[User]:581        """Authenticate request using token or API key"""582        583        if token:584            user = await self.auth_manager.verify_jwt_token(token)585            if user:586                await self.audit_logger.log_event(587                    AuditEventType.LOGIN,588                    user.id,589                    "authentication",590                    "jwt_verification",591                    {"method": "jwt", "success": True}592                )593                return user594        595        if api_key:596            user = await self.auth_manager.authenticate_api_key(api_key)597            if user:598                await self.audit_logger.log_event(599                    AuditEventType.LOGIN,600                    user.id,601                    "authentication",602                    "api_key_verification",603                    {"method": "api_key", "success": True}604                )605                return user606        607        # Log failed authentication608        await self.audit_logger.log_event(609            AuditEventType.SECURITY_VIOLATION,610            None,611            "authentication",612            "failed_authentication",613            {"method": "unknown", "success": False}614        )615        616        return None617    618    async def authorize_action(self, user: User, resource: str, action: str) -> bool:619        """Authorize user action on resource"""620        621        # Check permissions622        permission_required = f"{resource}:{action}"623        has_permission = self.auth_manager.check_permission(user, permission_required)624        625        # Log access attempt626        await self.audit_logger.log_event(627            AuditEventType.DATA_ACCESS if has_permission else AuditEventType.SECURITY_VIOLATION,628            user.id,629            resource,630            action,631            {"authorized": has_permission, "permission_required": permission_required}632        )633        634        return has_permission635    636    def encrypt_sensitive_data(self, data: str) -> str:637        """Encrypt sensitive data"""638        return self.encryption_manager.encrypt_data(data)639    640    def decrypt_sensitive_data(self, encrypted_data: str) -> str:641        """Decrypt sensitive data"""642        return self.encryption_manager.decrypt_data(encrypted_data)643    644    async def get_security_metrics(self) -> Dict[str, Any]:645        """Get security metrics"""646        647        # Get recent audit events648        recent_events = await self.audit_logger.get_audit_events(649            start_time=datetime.now() - timedelta(hours=24),650            limit=1000651        )652        653        # Calculate metrics654        total_events = len(recent_events)655        high_risk_events = sum(1 for e in recent_events if e.risk_score >= 4.0)656        failed_logins = sum(1 for e in recent_events 657                           if e.event_type == AuditEventType.LOGIN and 658                           e.details.get('failed_attempt', False))659        660        active_users = len([u for u in self.auth_manager.users.values() if u.is_active])661        active_api_keys = len([k for k in self.auth_manager.api_keys.values() if k.is_active])662        663        return {664            'total_events_24h': total_events,665            'high_risk_events_24h': high_risk_events,666            'failed_logins_24h': failed_logins,667            'active_users': active_users,668            'active_api_keys': active_api_keys,669            'security_score': max(0, 100 - (high_risk_events * 10) - (failed_logins * 2))670        }671    672    async def shutdown(self):673        """Shutdown security manager"""674        if self.auth_manager.db_connection:675            self.auth_manager.db_connection.close()676        677        if self.audit_logger.db_connection:678            self.audit_logger.db_connection.close()679        680        self.logger.info("Security manager shutdown")681