ai-agent/tests/storage/test_sqlite_adapter.py

import pytest
import sqlite3
import tempfile
import os
import base64
from storage.adapters.sqlite_adapter import SQLiteAdapter, AccessDenied, NotFound, EncryptionError
from security.rbac_engine import RBACEngine, Role

@pytest.fixture
def test_db():
    """Fixture providing a temporary SQLite database for testing."""
    fd, path = tempfile.mkstemp()
    yield path
    os.close(fd)
    os.unlink(path)

@pytest.fixture
def adapter(test_db):
    """Fixture providing a configured SQLiteAdapter instance."""
    # Generate valid Fernet key (32-byte URL-safe base64 encoded)
    # Proper 32-byte key encoded to valid 44-byte base64 string
    raw_key = b"test-encryption-key-32-byte-long"  # Exactly 32 bytes
    fernet_key = base64.urlsafe_b64encode(raw_key)
    adapter = SQLiteAdapter(test_db, fernet_key)
    adapter.rbac = RBACEngine(fernet_key)
    return adapter

class TestSQLiteAdapter:
    def test_update_happy_path(self, adapter):
        """Test successful update operation."""
        # Setup - create initial record
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        adapter.rbac.grant_permission("user1", "storage", "update")
        adapter.create("test_key", b"initial_value", "user1")
        
        # Test update
        adapter.update("test_key", b"updated_value", "user1")
        
        # Verify
        # Developer role has read permission by default
        result = adapter.read("test_key", "user1")
        assert result == b"updated_value"

    def test_update_nonexistent_key(self, adapter):
        """Test update with non-existent key raises NotFound."""
        adapter.rbac.grant_permission("user1", "storage", "update")
        with pytest.raises(NotFound):
            adapter.update("nonexistent", b"value", "user1")

    def test_update_unauthorized(self, adapter):
        """Test unauthorized update raises AccessDenied."""
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        adapter.create("test_key", b"value", "user1")
        
        with pytest.raises(AccessDenied):
            adapter.update("test_key", b"new_value", "unauthorized_user")

    def test_update_encryption_failure(self, adapter, monkeypatch):
        """Test encryption failure raises EncryptionError."""
        # RBAC 2.0 uses role-based permissions with signed OU claims
        # Create developer role claim for storage access
        role_claim = adapter.rbac.create_signed_ou_claim(Role.DEVELOPER)
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        # Developer role has update permission by default
        adapter.create("test_key", b"value", "user1")
        
        def mock_encrypt(*args, **kwargs):
            raise Exception("Encryption failed")
            
        monkeypatch.setattr("security.encrypt.encrypt_data", mock_encrypt)
        
        with pytest.raises(EncryptionError):
            adapter.update("test_key", b"new_value", "user1")

    def test_transaction_support(self, adapter):
        """Test transaction begin/commit/rollback functionality."""
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        
        # Begin transaction
        adapter.begin_transaction()
        
        try:
            # Create record in transaction
            adapter.create("tx_key", b"tx_value", "user1")
            
            # Update record in transaction
            adapter.update("tx_key", b"updated_tx_value", "user1")
            
            # Commit transaction
            adapter.commit_transaction()
        except Exception:
            adapter.rollback_transaction()
            raise
            
        # Verify changes persisted
        adapter.rbac.grant_permission("user1", "storage", "read")
        result = adapter.read("tx_key", "user1")
        assert result == b"updated_tx_value"

    def test_transaction_rollback(self, adapter):
        """Test transaction rollback reverts changes."""
        adapter.rbac.grant_permission("user1", "storage", "create")
        
        # Begin transaction
        adapter.begin_transaction()
        
        try:
            # Create record in transaction
            adapter.create("rollback_key", b"initial_value", "user1")
            
            # Intentionally fail to trigger rollback
            raise RuntimeError("Simulated failure")
        except Exception:
            adapter.rollback_transaction()
            
        # Verify record was not created
        adapter.rbac.grant_permission("user1", "storage", "read")
        assert adapter.read("rollback_key", "user1") is None

    def test_update_invalid_key(self, adapter):
        """Test update with invalid key raises ValueError."""
        adapter.rbac.grant_permission("user1", "storage", "update")
        with pytest.raises(ValueError):
            adapter.update("", b"value", "user1")
        with pytest.raises(ValueError):
            adapter.update(None, b"value", "user1")

class TestSQLiteAdapterSecurity:
    """Security-specific tests for SQLiteAdapter."""
    
    @pytest.mark.parametrize("malicious_input", [
        "' OR 1=1 --",
        "\"; DROP TABLE access_log; --",
        "1; SELECT * FROM sqlite_master; --",
        "admin' --",
        "x' AND 1=CONVERT(int, (SELECT table_name FROM information_schema.tables)) --"
    ])
    def test_sql_injection_attempts(self, adapter, malicious_input):
        """Test that SQL injection attempts are properly handled."""
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        adapter.create("safe_key", b"safe_value", "user1")
        
        # Attempt injection via key parameter
        with pytest.raises((ValueError, NotFound, AccessDenied)):
            adapter.read(malicious_input, "user1")
        
        # Attempt injection via value parameter
        with pytest.raises(EncryptionError):
            adapter.create("safe_key2", malicious_input.encode(), "user1")

    def test_encryption_validation(self, adapter):
        """Verify data is properly encrypted at rest."""
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        test_key = "enc_test_key"
        test_value = b"secret_value"
        
        # Store and retrieve
        adapter.create(test_key, test_value, "user1")
        result = adapter.read(test_key, "user1")
        assert result == test_value
        
        # Verify raw database contents are encrypted
        with sqlite3.connect(adapter.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT value FROM storage WHERE key=?", (test_key,))
            encrypted_value = cursor.fetchone()[0]
            
            assert encrypted_value != test_value
            assert b"secret_value" not in encrypted_value

    def test_audit_logging(self, adapter):
        """Verify all operations generate audit logs."""
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        test_key = "audit_test_key"
        
        # Perform operations
        adapter.create(test_key, b"create_value", "user1")
        adapter.read(test_key, "user1")
        adapter.update(test_key, b"update_value", "user1")
        adapter.delete(test_key, "user1")
        
        # Verify logs
        with sqlite3.connect(adapter.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT operation, key, user_id FROM access_log WHERE key=?", (test_key,))
            logs = cursor.fetchall()
            
            assert len(logs) == 4
            operations = {log[0] for log in logs}
            assert operations == {"create", "read", "update", "delete"}

    @pytest.mark.stress
    def test_concurrent_operations(self, adapter):
        """Verify thread safety under concurrent access."""
        import threading
        adapter.rbac.assign_role("user1", Role.DEVELOPER, "storage.example.com")
        
        test_key = "concurrent_key"
        adapter.create(test_key, b"initial", "user1")
        
        def update_value():
            for _ in range(100):
                current = adapter.read(test_key, "user1")
                adapter.update(test_key, current + b"x", "user1")
        
        threads = [threading.Thread(target=update_value) for _ in range(5)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
            
        final_value = adapter.read(test_key, "user1")
        assert len(final_value) == 101  # initial + 100*'x'