Files
MyFSIO/tests/test_kms_api.py
2025-12-01 00:46:12 +08:00

507 lines
17 KiB
Python

"""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