ai-agent/security/memory/core.py

import logging
from dataclasses import dataclass
from typing import Dict, Optional, Any
import os
import json
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from security.rbac_engine import RBACEngine, ClientCertInfo

logger = logging.getLogger('MemoryCore')

class EncryptionError(Exception):
    pass

class DecryptionError(Exception):
    pass

class AccessDenied(Exception):
    pass

class NotFound(Exception):
    pass

@dataclass
class AuditEntry:
    timestamp: str
    operation: str
    key_hash: str
    status: bool
    caller: str
    details: Optional[str] = None

class MemoryCore:
    def __init__(self, encryption_key: bytes, rbac_engine: RBACEngine):
        # Initialize encryption
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        self.aes_key = kdf.derive(encryption_key)
        self.salt = salt

        # Data storage
        self.data: Dict[str, bytes] = {}

        # RBAC integration
        self.rbac = rbac_engine

        # Audit log
        self.audit_log: list[AuditEntry] = []

    def _encrypt(self, plaintext: bytes) -> bytes:
        """Encrypt data using AES-256-GCM"""
        nonce = os.urandom(12)
        aesgcm = AESGCM(self.aes_key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)
        return nonce + ciphertext

    def _decrypt(self, ciphertext: bytes) -> bytes:
        """Decrypt data using AES-256-GCM"""
        nonce = ciphertext[:12]
        ciphertext = ciphertext[12:]
        aesgcm = AESGCM(self.aes_key)
        try:
            return aesgcm.decrypt(nonce, ciphertext, None)
        except Exception as e:
            raise DecryptionError(f"Decryption failed: {str(e)}")

    def _hash_key(self, key: str) -> str:
        """Create secure hash of key for audit logging"""
        return hashlib.sha256(key.encode()).hexdigest()

    def _audit(self, operation: str, key: str, status: bool, 
              caller: Optional[str] = None, details: Optional[str] = None):
        """Record audit entry"""
        entry = AuditEntry(
            timestamp=datetime.now().isoformat(),
            operation=operation,
            key_hash=self._hash_key(key),
            status=status,
            caller=caller or "system",
            details=details
        )
        self.audit_log.append(entry)
        logger.info(f"Audit: {entry}")

    def create(self, key: str, value: bytes, 
              user: Optional[str] = None,
              cert_info: Optional[ClientCertInfo] = None) -> bool:
        """Create new encrypted entry with RBAC check"""
        if not self.rbac.validate_permission("memory", "create", user=user, client_cert_info=cert_info):
            self._audit("create", key, False, user or cert_info.subject.get('CN'), "RBAC check failed")
            raise AccessDenied("Create permission denied")

        try:
            encrypted = self._encrypt(value)
            self.data[key] = encrypted
            self._audit("create", key, True, user or cert_info.subject.get('CN'))
            return True
        except Exception as e:
            self._audit("create", key, False, user or cert_info.subject.get('CN'), str(e))
            raise EncryptionError(f"Encryption failed: {str(e)}")

    def read(self, key: str, 
            user: Optional[str] = None,
            cert_info: Optional[ClientCertInfo] = None) -> bytes:
        """Read and decrypt entry with RBAC check"""
        if not self.rbac.validate_permission("memory", "read", user=user, client_cert_info=cert_info):
            self._audit("read", key, False, user or cert_info.subject.get('CN'), "RBAC check failed")
            raise AccessDenied("Read permission denied")

        if key not in self.data:
            self._audit("read", key, False, user or cert_info.subject.get('CN'), "Key not found")
            raise NotFound(f"Key {key} not found")

        try:
            decrypted = self._decrypt(self.data[key])
            self._audit("read", key, True, user or cert_info.subject.get('CN'))
            return decrypted
        except Exception as e:
            self._audit("read", key, False, user or cert_info.subject.get('CN'), str(e))
            raise DecryptionError(f"Decryption failed: {str(e)}")

    def update(self, key: str, value: bytes, 
              user: Optional[str] = None,
              cert_info: Optional[ClientCertInfo] = None) -> bool:
        """Update encrypted entry with RBAC check"""
        if not self.rbac.validate_permission("memory", "update", user=user, client_cert_info=cert_info):
            self._audit("update", key, False, user or cert_info.subject.get('CN'), "RBAC check failed")
            raise AccessDenied("Update permission denied")

        if key not in self.data:
            self._audit("update", key, False, user or cert_info.subject.get('CN'), "Key not found")
            raise NotFound(f"Key {key} not found")

        try:
            encrypted = self._encrypt(value)
            self.data[key] = encrypted
            self._audit("update", key, True, user or cert_info.subject.get('CN'))
            return True
        except Exception as e:
            self._audit("update", key, False, user or cert_info.subject.get('CN'), str(e))
            raise EncryptionError(f"Encryption failed: {str(e)}")

    def delete(self, key: str, 
              user: Optional[str] = None,
              cert_info: Optional[ClientCertInfo] = None) -> bool:
        """Delete entry with RBAC check"""
        if not self.rbac.validate_permission("memory", "delete", user=user, client_cert_info=cert_info):
            self._audit("delete", key, False, user or cert_info.subject.get('CN'), "RBAC check failed")
            raise AccessDenied("Delete permission denied")

        if key not in self.data:
            self._audit("delete", key, False, user or cert_info.subject.get('CN'), "Key not found")
            raise NotFound(f"Key {key} not found")

        del self.data[key]
        self._audit("delete", key, True, user or cert_info.subject.get('CN'))
        return True