ai-agent/storage/adapters/sqlite_adapter.py

from security.encrypt import encrypt_data, decrypt_data, generate_key
import sqlite3
import threading
import hashlib
import logging
import base64
from typing import Optional, Dict, Any, Union
from security.encrypt import encrypt_data, decrypt_data
from security.rbac_engine import RBACEngine

class AccessDenied(Exception):
    """Raised when RBAC validation fails"""

class NotFound(Exception):
    """Raised when requested key doesn't exist"""

class EncryptionError(Exception):
    """Raised when encryption operation fails"""

class SQLiteAdapter:
    """SQLite storage adapter with RBAC and encryption support.
    Integrates with TaskDispatcher for secure storage operations.
    
    Attributes:
        db_path: Path to SQLite database file
        encryption_key: Key used for data encryption
        rbac: RBAC engine instance (must be set before use)
        _lock: Thread lock for concurrent access
        logger: Logger instance for operations
    """
    
    rbac: RBACEngine = None  # Must be set by application
    logger = logging.getLogger('SQLiteAdapter')
    
    def __init__(self, db_path: str, encryption_key: str):
        """Initialize SQLite adapter with dispatcher integration.
        
        Args:
            db_path: Path to SQLite database file
            encryption_key: Encryption key for data protection (32+ chars or 32 bytes)
            
        Raises:
            RuntimeError: If encryption key is invalid
        """
        if not encryption_key or (isinstance(encryption_key, str) and len(encryption_key) < 32):
            raise RuntimeError("Encryption key must be at least 32 characters or 32 bytes")
            
        self.db_path = db_path
        self.encryption_key = self._convert_key(encryption_key)
        self._lock = threading.Lock()
        self._init_db()
        self.logger.info(f"Initialized SQLite adapter for {db_path}")
    
    def _init_db(self):
        """Initialize database tables."""
        with self._lock, sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS storage (
                    key_hash TEXT PRIMARY KEY,
                    encrypted_value BLOB,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    created_by TEXT
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS access_log (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    key_hash TEXT,
                    operation TEXT,
                    user_id TEXT,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY(key_hash) REFERENCES storage(key_hash)
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS performance_metrics (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    operation TEXT NOT NULL,
                    execution_time_ms INTEGER NOT NULL,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    user_id TEXT,
                    key_hash TEXT,
                    FOREIGN KEY(key_hash) REFERENCES storage(key_hash)
                )
            """)
    
    def _hash_key(self, key):
        """Generate SHA-256 hash of key."""
        return hashlib.sha256(key.encode()).hexdigest()
        
    def _convert_key(self, key: Union[str, bytes]) -> bytes:
        """Convert encryption key to bytes format required by AES-256-GCM.
        Supports:
        - 32-byte raw keys
        - 44-byte base64 encoded keys
        - String keys (hashed to 32 bytes)
        
        Args:
            key: Original key (str or bytes)
            
        Returns:
            bytes: 32-byte key
            
        Raises:
            RuntimeError: If key cannot be converted to 32 bytes
        """
        if isinstance(key, bytes):
            if len(key) == 32:
                return key
            if len(key) == 44:  # Base64 encoded
                try:
                    decoded = base64.urlsafe_b64decode(key)
                    if len(decoded) == 32:
                        return decoded
                    raise ValueError("Decoded key must be 32 bytes")
                except Exception as e:
                    raise RuntimeError(f"Invalid base64 key: {str(e)}")
            raise RuntimeError("Bytes key must be 32 raw bytes or 44 base64 bytes")
            
        # Convert string key to bytes using SHA-256 for consistent length
        return hashlib.sha256(key.encode()).digest()[:32]
    
    def create(self, key: str, value: Any, user_id: str, client_cert_info: Optional[Dict] = None) -> bool:
        """Create new storage entry with RBAC check.
        Integrates with dispatcher's RBAC system.
        
        Args:
            key: Storage key
            value: Value to store (will be encrypted)
            user_id: User ID for RBAC validation
            client_cert_info: Optional client certificate info for enhanced auth
        
        Returns:
            bool: True if successful, False if unauthorized
        
        Raises:
            ValueError: If key or value is invalid
        """
        if not key or not isinstance(key, str):
            raise ValueError("Key must be non-empty string")
            
        if not self.rbac or not self.rbac.validate_permission(user_id, "storage", "create"):
            self.logger.warning(f"Unauthorized create attempt by {user_id}")
            return False
            
        key_hash = self._hash_key(key)
        try:
            encrypted_value = encrypt_data(value, self.encryption_key)
        except Exception as e:
            self.logger.error(f"Encryption failed: {str(e)}")
            raise RuntimeError("Data encryption failed") from e
        
        with self._lock, sqlite3.connect(self.db_path) as conn:
            try:
                conn.execute(
                    "INSERT INTO storage (key_hash, encrypted_value, created_by) VALUES (?, ?, ?)",
                    (key_hash, encrypted_value, user_id)
                )
                conn.execute(
                    "INSERT INTO access_log (key_hash, operation, user_id) VALUES (?, ?, ?)",
                    (key_hash, "create", user_id)
                )
                return True
            except sqlite3.IntegrityError:
                return False
    
    def read(self, key: str, user_id: str) -> Optional[Any]:
        """Read storage entry with RBAC check.
        Integrates with dispatcher's RBAC system.
        
        Args:
            key: Storage key
            user_id: User ID for RBAC validation
        
        Returns:
            Decrypted value or None if unauthorized/not found
        
        Raises:
            ValueError: If key is invalid
            RuntimeError: If decryption fails
        """
        if not key or not isinstance(key, str):
            raise ValueError("Key must be non-empty string")
            
        if not self.rbac or not self.rbac.validate_permission(user_id, "storage", "read"):
            self.logger.warning(f"Unauthorized read attempt by {user_id}")
            return None
            
        key_hash = self._hash_key(key)
        
        with self._lock, sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "SELECT encrypted_value FROM storage WHERE key_hash = ?",
                (key_hash,))
            row = cursor.fetchone()
            
            if row:
                conn.execute(
                    "INSERT INTO access_log (key_hash, operation, user_id) VALUES (?, ?, ?)",
                    (key_hash, "read", user_id)
                )
                try:
                    return decrypt_data(row[0], self.encryption_key)
                except Exception as e:
                    self.logger.error(f"Decryption failed: {str(e)}")
                    raise RuntimeError("Data decryption failed") from e
            return None

    def delete(self, key: str, user_id: str) -> bool:
        """Delete storage entry with RBAC check.
        Integrates with dispatcher's RBAC system.
        
        Args:
            key: Storage key
            user_id: User ID for RBAC validation
        
        Returns:
            bool: True if successful, False if unauthorized or not found
        
        Raises:
            ValueError: If key is invalid
        """
        if not key or not isinstance(key, str):
            raise ValueError("Key must be non-empty string")
            
        if not self.rbac or not self.rbac.validate_permission(user_id, "storage", "delete"):
            self.logger.warning(f"Unauthorized delete attempt by {user_id}")
            return False
            
        key_hash = self._hash_key(key)
        
        with self._lock, sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "DELETE FROM storage WHERE key_hash = ?",
                (key_hash,))
            
            if cursor.rowcount > 0:
                conn.execute(
                    "INSERT INTO access_log (key_hash, operation, user_id) VALUES (?, ?, ?)",
                    (key_hash, "delete", user_id)
                )
                return True
            return False
    
    def close(self):
        """Close database connections."""
        pass  # SQLite connections are closed automatically
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def update(self, key: str, value: bytes, user_id: str) -> None:
        """Update storage entry with RBAC check.
        Integrates with dispatcher's RBAC system.
        
        Args:
            key: Storage key (non-empty string)
            value: New value to store as bytes (will be encrypted)
            user_id: User ID for RBAC validation
        
        Returns:
            None
        
        Raises:
            ValueError: If key is invalid
            NotFound: If key doesn't exist
            AccessDenied: If RBAC validation fails
            EncryptionError: If encryption fails
        """
        if not key or not isinstance(key, str):
            raise ValueError("Key must be non-empty string")
            
        if not self.rbac or not self.rbac.validate_permission(user_id, "storage", "update"):
            self.logger.warning(f"Unauthorized update attempt by {user_id}")
            raise AccessDenied("Update permission denied")
            
        key_hash = self._hash_key(key)
        try:
            encrypted_value = encrypt_data(value, self.encryption_key)
        except Exception as e:
            self.logger.error(f"Encryption failed: {str(e)}")
            raise EncryptionError("Data encryption failed") from e
        
        with self._lock, sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "UPDATE storage SET encrypted_value = ?, updated_at = CURRENT_TIMESTAMP WHERE key_hash = ?",
                (encrypted_value, key_hash))
            
            if cursor.rowcount == 0:
                raise NotFound(f"Key {key} not found")
                
            conn.execute(
                "INSERT INTO access_log (key_hash, operation, user_id) VALUES (?, ?, ?)",
                (key_hash, "update", user_id)
            )

    def begin_transaction(self):
        """Begin a new transaction."""
        self._lock.acquire()
        self._transaction_conn = sqlite3.connect(self.db_path)
        self._transaction_conn.execute("BEGIN")

    def commit_transaction(self):
        """Commit the current transaction."""
        if hasattr(self, '_transaction_conn'):
            self._transaction_conn.commit()
            self._transaction_conn.close()
            self._lock.release()

    def rollback_transaction(self):
        """Rollback the current transaction."""
        if hasattr(self, '_transaction_conn'):
            self._transaction_conn.rollback()
            self._transaction_conn.close()
            self._lock.release()