Add new SSE, KMS encryptions

tests/test_kms_api.py

@@ -0,0 +1,506 @@
+"""Tests for KMS API endpoints."""
+from __future__ import annotations
+
+import base64
+import json
+import secrets
+
+import pytest
+
+
+@pytest.fixture
+def kms_client(tmp_path):
+    """Create a test client with KMS enabled."""
+    from app import create_app
+    
+    app = create_app({
+        "TESTING": True,
+        "STORAGE_ROOT": str(tmp_path / "storage"),
+        "IAM_CONFIG": str(tmp_path / "iam.json"),
+        "BUCKET_POLICY_PATH": str(tmp_path / "policies.json"),
+        "ENCRYPTION_ENABLED": True,
+        "KMS_ENABLED": True,
+        "ENCRYPTION_MASTER_KEY_PATH": str(tmp_path / "master.key"),
+        "KMS_KEYS_PATH": str(tmp_path / "kms_keys.json"),
+    })
+    
+    # Create default IAM config with admin user
+    iam_config = {
+        "users": [
+            {
+                "access_key": "test-access-key",
+                "secret_key": "test-secret-key",
+                "display_name": "Test User",
+                "permissions": ["*"]
+            }
+        ]
+    }
+    (tmp_path / "iam.json").write_text(json.dumps(iam_config))
+    
+    return app.test_client()
+
+
+@pytest.fixture
+def auth_headers():
+    """Get authentication headers."""
+    return {
+        "X-Access-Key": "test-access-key",
+        "X-Secret-Key": "test-secret-key",
+    }
+
+
+class TestKMSKeyManagement:
+    """Tests for KMS key management endpoints."""
+    
+    def test_create_key(self, kms_client, auth_headers):
+        """Test creating a KMS key."""
+        response = kms_client.post(
+            "/kms/keys",
+            json={"Description": "Test encryption key"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "KeyMetadata" in data
+        assert data["KeyMetadata"]["Description"] == "Test encryption key"
+        assert data["KeyMetadata"]["Enabled"] is True
+        assert "KeyId" in data["KeyMetadata"]
+    
+    def test_create_key_with_custom_id(self, kms_client, auth_headers):
+        """Test creating a key with a custom ID."""
+        response = kms_client.post(
+            "/kms/keys",
+            json={"KeyId": "my-custom-key", "Description": "Custom key"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert data["KeyMetadata"]["KeyId"] == "my-custom-key"
+    
+    def test_list_keys(self, kms_client, auth_headers):
+        """Test listing KMS keys."""
+        # Create some keys
+        kms_client.post("/kms/keys", json={"Description": "Key 1"}, headers=auth_headers)
+        kms_client.post("/kms/keys", json={"Description": "Key 2"}, headers=auth_headers)
+        
+        response = kms_client.get("/kms/keys", headers=auth_headers)
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "Keys" in data
+        assert len(data["Keys"]) == 2
+    
+    def test_get_key(self, kms_client, auth_headers):
+        """Test getting a specific key."""
+        # Create a key
+        create_response = kms_client.post(
+            "/kms/keys",
+            json={"KeyId": "test-key", "Description": "Test key"},
+            headers=auth_headers,
+        )
+        
+        response = kms_client.get("/kms/keys/test-key", headers=auth_headers)
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert data["KeyMetadata"]["KeyId"] == "test-key"
+        assert data["KeyMetadata"]["Description"] == "Test key"
+    
+    def test_get_nonexistent_key(self, kms_client, auth_headers):
+        """Test getting a key that doesn't exist."""
+        response = kms_client.get("/kms/keys/nonexistent", headers=auth_headers)
+        
+        assert response.status_code == 404
+    
+    def test_delete_key(self, kms_client, auth_headers):
+        """Test deleting a key."""
+        # Create a key
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        # Delete it
+        response = kms_client.delete("/kms/keys/test-key", headers=auth_headers)
+        
+        assert response.status_code == 204
+        
+        # Verify it's gone
+        get_response = kms_client.get("/kms/keys/test-key", headers=auth_headers)
+        assert get_response.status_code == 404
+    
+    def test_enable_disable_key(self, kms_client, auth_headers):
+        """Test enabling and disabling a key."""
+        # Create a key
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        # Disable
+        response = kms_client.post("/kms/keys/test-key/disable", headers=auth_headers)
+        assert response.status_code == 200
+        
+        # Verify disabled
+        get_response = kms_client.get("/kms/keys/test-key", headers=auth_headers)
+        assert get_response.get_json()["KeyMetadata"]["Enabled"] is False
+        
+        # Enable
+        response = kms_client.post("/kms/keys/test-key/enable", headers=auth_headers)
+        assert response.status_code == 200
+        
+        # Verify enabled
+        get_response = kms_client.get("/kms/keys/test-key", headers=auth_headers)
+        assert get_response.get_json()["KeyMetadata"]["Enabled"] is True
+
+
+class TestKMSEncryption:
+    """Tests for KMS encryption operations."""
+    
+    def test_encrypt_decrypt(self, kms_client, auth_headers):
+        """Test encrypting and decrypting data."""
+        # Create a key
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        plaintext = b"Hello, World!"
+        plaintext_b64 = base64.b64encode(plaintext).decode()
+        
+        # Encrypt
+        encrypt_response = kms_client.post(
+            "/kms/encrypt",
+            json={"KeyId": "test-key", "Plaintext": plaintext_b64},
+            headers=auth_headers,
+        )
+        
+        assert encrypt_response.status_code == 200
+        encrypt_data = encrypt_response.get_json()
+        
+        assert "CiphertextBlob" in encrypt_data
+        assert encrypt_data["KeyId"] == "test-key"
+        
+        # Decrypt
+        decrypt_response = kms_client.post(
+            "/kms/decrypt",
+            json={"CiphertextBlob": encrypt_data["CiphertextBlob"]},
+            headers=auth_headers,
+        )
+        
+        assert decrypt_response.status_code == 200
+        decrypt_data = decrypt_response.get_json()
+        
+        decrypted = base64.b64decode(decrypt_data["Plaintext"])
+        assert decrypted == plaintext
+    
+    def test_encrypt_with_context(self, kms_client, auth_headers):
+        """Test encryption with encryption context."""
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        plaintext = b"Contextualized data"
+        plaintext_b64 = base64.b64encode(plaintext).decode()
+        context = {"purpose": "testing", "bucket": "my-bucket"}
+        
+        # Encrypt with context
+        encrypt_response = kms_client.post(
+            "/kms/encrypt",
+            json={
+                "KeyId": "test-key",
+                "Plaintext": plaintext_b64,
+                "EncryptionContext": context,
+            },
+            headers=auth_headers,
+        )
+        
+        assert encrypt_response.status_code == 200
+        ciphertext = encrypt_response.get_json()["CiphertextBlob"]
+        
+        # Decrypt with same context succeeds
+        decrypt_response = kms_client.post(
+            "/kms/decrypt",
+            json={
+                "CiphertextBlob": ciphertext,
+                "EncryptionContext": context,
+            },
+            headers=auth_headers,
+        )
+        
+        assert decrypt_response.status_code == 200
+        
+        # Decrypt with wrong context fails
+        wrong_context_response = kms_client.post(
+            "/kms/decrypt",
+            json={
+                "CiphertextBlob": ciphertext,
+                "EncryptionContext": {"wrong": "context"},
+            },
+            headers=auth_headers,
+        )
+        
+        assert wrong_context_response.status_code == 400
+    
+    def test_encrypt_missing_key_id(self, kms_client, auth_headers):
+        """Test encryption without KeyId."""
+        response = kms_client.post(
+            "/kms/encrypt",
+            json={"Plaintext": base64.b64encode(b"data").decode()},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 400
+        assert "KeyId is required" in response.get_json()["message"]
+    
+    def test_encrypt_missing_plaintext(self, kms_client, auth_headers):
+        """Test encryption without Plaintext."""
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        response = kms_client.post(
+            "/kms/encrypt",
+            json={"KeyId": "test-key"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 400
+        assert "Plaintext is required" in response.get_json()["message"]
+
+
+class TestKMSDataKey:
+    """Tests for KMS data key generation."""
+    
+    def test_generate_data_key(self, kms_client, auth_headers):
+        """Test generating a data key."""
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        response = kms_client.post(
+            "/kms/generate-data-key",
+            json={"KeyId": "test-key"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "Plaintext" in data
+        assert "CiphertextBlob" in data
+        assert data["KeyId"] == "test-key"
+        
+        # Verify plaintext key is 256 bits (32 bytes)
+        plaintext_key = base64.b64decode(data["Plaintext"])
+        assert len(plaintext_key) == 32
+    
+    def test_generate_data_key_aes_128(self, kms_client, auth_headers):
+        """Test generating an AES-128 data key."""
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        response = kms_client.post(
+            "/kms/generate-data-key",
+            json={"KeyId": "test-key", "KeySpec": "AES_128"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        # Verify plaintext key is 128 bits (16 bytes)
+        plaintext_key = base64.b64decode(data["Plaintext"])
+        assert len(plaintext_key) == 16
+    
+    def test_generate_data_key_without_plaintext(self, kms_client, auth_headers):
+        """Test generating a data key without plaintext."""
+        kms_client.post("/kms/keys", json={"KeyId": "test-key"}, headers=auth_headers)
+        
+        response = kms_client.post(
+            "/kms/generate-data-key-without-plaintext",
+            json={"KeyId": "test-key"},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "CiphertextBlob" in data
+        assert "Plaintext" not in data
+
+
+class TestKMSReEncrypt:
+    """Tests for KMS re-encryption."""
+    
+    def test_re_encrypt(self, kms_client, auth_headers):
+        """Test re-encrypting data with a different key."""
+        # Create two keys
+        kms_client.post("/kms/keys", json={"KeyId": "key-1"}, headers=auth_headers)
+        kms_client.post("/kms/keys", json={"KeyId": "key-2"}, headers=auth_headers)
+        
+        # Encrypt with key-1
+        plaintext = b"Data to re-encrypt"
+        encrypt_response = kms_client.post(
+            "/kms/encrypt",
+            json={
+                "KeyId": "key-1",
+                "Plaintext": base64.b64encode(plaintext).decode(),
+            },
+            headers=auth_headers,
+        )
+        
+        ciphertext = encrypt_response.get_json()["CiphertextBlob"]
+        
+        # Re-encrypt with key-2
+        re_encrypt_response = kms_client.post(
+            "/kms/re-encrypt",
+            json={
+                "CiphertextBlob": ciphertext,
+                "DestinationKeyId": "key-2",
+            },
+            headers=auth_headers,
+        )
+        
+        assert re_encrypt_response.status_code == 200
+        data = re_encrypt_response.get_json()
+        
+        assert data["SourceKeyId"] == "key-1"
+        assert data["KeyId"] == "key-2"
+        
+        # Verify new ciphertext can be decrypted
+        decrypt_response = kms_client.post(
+            "/kms/decrypt",
+            json={"CiphertextBlob": data["CiphertextBlob"]},
+            headers=auth_headers,
+        )
+        
+        decrypted = base64.b64decode(decrypt_response.get_json()["Plaintext"])
+        assert decrypted == plaintext
+
+
+class TestKMSRandom:
+    """Tests for random number generation."""
+    
+    def test_generate_random(self, kms_client, auth_headers):
+        """Test generating random bytes."""
+        response = kms_client.post(
+            "/kms/generate-random",
+            json={"NumberOfBytes": 64},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        random_bytes = base64.b64decode(data["Plaintext"])
+        assert len(random_bytes) == 64
+    
+    def test_generate_random_default_size(self, kms_client, auth_headers):
+        """Test generating random bytes with default size."""
+        response = kms_client.post(
+            "/kms/generate-random",
+            json={},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        random_bytes = base64.b64decode(data["Plaintext"])
+        assert len(random_bytes) == 32  # Default is 32 bytes
+
+
+class TestClientSideEncryption:
+    """Tests for client-side encryption helpers."""
+    
+    def test_generate_client_key(self, kms_client, auth_headers):
+        """Test generating a client encryption key."""
+        response = kms_client.post(
+            "/kms/client/generate-key",
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "key" in data
+        assert data["algorithm"] == "AES-256-GCM"
+        
+        key = base64.b64decode(data["key"])
+        assert len(key) == 32
+    
+    def test_client_encrypt_decrypt(self, kms_client, auth_headers):
+        """Test client-side encryption and decryption."""
+        # Generate a key
+        key_response = kms_client.post("/kms/client/generate-key", headers=auth_headers)
+        key = key_response.get_json()["key"]
+        
+        # Encrypt
+        plaintext = b"Client-side encrypted data"
+        encrypt_response = kms_client.post(
+            "/kms/client/encrypt",
+            json={
+                "Plaintext": base64.b64encode(plaintext).decode(),
+                "Key": key,
+            },
+            headers=auth_headers,
+        )
+        
+        assert encrypt_response.status_code == 200
+        encrypted = encrypt_response.get_json()
+        
+        # Decrypt
+        decrypt_response = kms_client.post(
+            "/kms/client/decrypt",
+            json={
+                "Ciphertext": encrypted["ciphertext"],
+                "Nonce": encrypted["nonce"],
+                "Key": key,
+            },
+            headers=auth_headers,
+        )
+        
+        assert decrypt_response.status_code == 200
+        decrypted = base64.b64decode(decrypt_response.get_json()["Plaintext"])
+        assert decrypted == plaintext
+
+
+class TestEncryptionMaterials:
+    """Tests for S3 encryption materials endpoint."""
+    
+    def test_get_encryption_materials(self, kms_client, auth_headers):
+        """Test getting encryption materials for client-side S3 encryption."""
+        # Create a key
+        kms_client.post("/kms/keys", json={"KeyId": "s3-key"}, headers=auth_headers)
+        
+        response = kms_client.post(
+            "/kms/materials/s3-key",
+            json={},
+            headers=auth_headers,
+        )
+        
+        assert response.status_code == 200
+        data = response.get_json()
+        
+        assert "PlaintextKey" in data
+        assert "EncryptedKey" in data
+        assert data["KeyId"] == "s3-key"
+        assert data["Algorithm"] == "AES-256-GCM"
+        
+        # Verify key is 256 bits
+        key = base64.b64decode(data["PlaintextKey"])
+        assert len(key) == 32
+
+
+class TestKMSAuthentication:
+    """Tests for KMS authentication requirements."""
+    
+    def test_unauthenticated_request_fails(self, kms_client):
+        """Test that unauthenticated requests are rejected."""
+        response = kms_client.get("/kms/keys")
+        
+        # Should fail with 403 (no credentials)
+        assert response.status_code == 403
+    
+    def test_invalid_credentials_fail(self, kms_client):
+        """Test that invalid credentials are rejected."""
+        response = kms_client.get(
+            "/kms/keys",
+            headers={
+                "X-Access-Key": "wrong-key",
+                "X-Secret-Key": "wrong-secret",
+            },
+        )
+        
+        assert response.status_code == 403