fetch_ml/internal/crypto/tenant_keys_test.go

package crypto_test

import (
	"testing"

	"github.com/jfraeys/fetch_ml/internal/crypto"
	"github.com/jfraeys/fetch_ml/internal/crypto/kms"
	kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// setupTenantKeyManager creates a test TenantKeyManager with memory provider
func setupTenantKeyManager(t *testing.T) *crypto.TenantKeyManager {
	t.Helper()
	return crypto.NewTestTenantKeyManager(nil)
}

// TestNewTenantKeyManager tests the constructor
func TestNewTenantKeyManager(t *testing.T) {
	t.Parallel()

	provider := kms.NewMemoryProvider()
	cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
	config := kmsconfig.Config{Provider: kmsconfig.ProviderTypeMemory}

	km := crypto.NewTenantKeyManager(provider, cache, config, nil)
	require.NotNil(t, km)
}

// TestNewTestTenantKeyManager tests the test constructor
func TestNewTestTenantKeyManager(t *testing.T) {
	t.Parallel()

	km := crypto.NewTestTenantKeyManager(nil)
	require.NotNil(t, km)
}

// TestProvisionTenant tests creating tenant root keys
func TestProvisionTenant(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	cases := []struct {
		name     string
		tenantID string
		wantErr  bool
	}{
		{"valid tenant", "tenant-1", false},
		{"another tenant", "tenant-2", false},
		{"empty tenant", "", true},
		{"whitespace tenant", "   ", true},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			hierarchy, err := km.ProvisionTenant(tc.tenantID)
			if tc.wantErr {
				require.Error(t, err)
				assert.Nil(t, hierarchy)
				return
			}
			require.NoError(t, err)
			assert.NotNil(t, hierarchy)
			assert.Equal(t, tc.tenantID, hierarchy.TenantID)
			assert.NotEmpty(t, hierarchy.RootKeyID)
			assert.NotEmpty(t, hierarchy.KMSKeyID)
			assert.Equal(t, "AES-256-GCM", hierarchy.Algorithm)
			assert.False(t, hierarchy.CreatedAt.IsZero())
		})
	}
}

// TestRotateTenantKey tests key rotation
func TestRotateTenantKey(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision initial key
	initial, err := km.ProvisionTenant("rotate-tenant")
	require.NoError(t, err)
	initialKeyID := initial.KMSKeyID

	// Rotate key
	rotated, err := km.RotateTenantKey("rotate-tenant", initial)
	require.NoError(t, err)
	assert.NotNil(t, rotated)
	assert.Equal(t, "rotate-tenant", rotated.TenantID)
	assert.NotEqual(t, initialKeyID, rotated.KMSKeyID)
	assert.NotEqual(t, initial.RootKeyID, rotated.RootKeyID)
}

// TestRevokeTenant tests tenant key revocation
func TestRevokeTenant(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision a tenant
	hierarchy, err := km.ProvisionTenant("revoke-tenant")
	require.NoError(t, err)

	// Revoke tenant
	err = km.RevokeTenant(hierarchy)
	require.NoError(t, err)
}

// TestGenerateDataEncryptionKey tests DEK generation
func TestGenerateDataEncryptionKey(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision tenant first
	hierarchy, err := km.ProvisionTenant("dek-tenant")
	require.NoError(t, err)

	// Generate DEK
	wrappedDEK, err := km.GenerateDataEncryptionKey("dek-tenant", "artifact-1", hierarchy.KMSKeyID)
	require.NoError(t, err)
	assert.NotNil(t, wrappedDEK)
	assert.Equal(t, "dek-tenant", wrappedDEK.TenantID)
	assert.Equal(t, "artifact-1", wrappedDEK.ArtifactID)
	assert.NotEmpty(t, wrappedDEK.WrappedKey)
	assert.Equal(t, "AES-256-GCM", wrappedDEK.Algorithm)
	assert.False(t, wrappedDEK.CreatedAt.IsZero())
}

// TestUnwrapDataEncryptionKey tests DEK unwrapping
func TestUnwrapDataEncryptionKey(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision tenant and generate DEK
	hierarchy, err := km.ProvisionTenant("unwrap-tenant")
	require.NoError(t, err)

	wrappedDEK, err := km.GenerateDataEncryptionKey("unwrap-tenant", "artifact-unwrap", hierarchy.KMSKeyID)
	require.NoError(t, err)

	// Unwrap DEK
	dek, err := km.UnwrapDataEncryptionKey(wrappedDEK, hierarchy.KMSKeyID)
	require.NoError(t, err)
	assert.NotNil(t, dek)
	assert.Len(t, dek, 32) // AES-256 key
}

// TestEncryptArtifact tests artifact encryption
func TestEncryptArtifact(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision tenant
	hierarchy, err := km.ProvisionTenant("encrypt-tenant")
	require.NoError(t, err)

	// Encrypt data
	plaintext := []byte("sensitive data for encryption test")
	encrypted, err := km.EncryptArtifact("encrypt-tenant", "artifact-enc", hierarchy.KMSKeyID, plaintext)
	require.NoError(t, err)
	assert.NotNil(t, encrypted)
	assert.NotEmpty(t, encrypted.Ciphertext)
	assert.NotNil(t, encrypted.DEK)
	assert.Equal(t, hierarchy.KMSKeyID, encrypted.KMSKeyID)
	assert.Equal(t, "AES-256-GCM", encrypted.Algorithm)
}

// TestDecryptArtifact tests artifact decryption
func TestDecryptArtifact(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision tenant
	hierarchy, err := km.ProvisionTenant("decrypt-tenant")
	require.NoError(t, err)

	// Encrypt data
	plaintext := []byte("secret message for decryption test")
	encrypted, err := km.EncryptArtifact("decrypt-tenant", "artifact-dec", hierarchy.KMSKeyID, plaintext)
	require.NoError(t, err)

	// Decrypt data
	decrypted, err := km.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
	require.NoError(t, err)
	assert.Equal(t, plaintext, decrypted)
}

// TestEncryptDecryptRoundtrip tests full encryption/decryption cycle
func TestEncryptDecryptRoundtrip(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision tenant
	hierarchy, err := km.ProvisionTenant("roundtrip-tenant")
	require.NoError(t, err)

	// Test various payloads
	payloads := [][]byte{
		[]byte("short"),
		[]byte("this is a longer message with more content"),
		[]byte{},
		[]byte{0x00, 0x01, 0x02, 0x03, 0xFF},
	}

	for i, payload := range payloads {
		artifactID := "artifact-rt-" + string(rune('a'+i))
		encrypted, err := km.EncryptArtifact("roundtrip-tenant", artifactID, hierarchy.KMSKeyID, payload)
		require.NoError(t, err)

		decrypted, err := km.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
		require.NoError(t, err)
		// Handle nil vs empty slice comparison
		if len(payload) == 0 {
			assert.Empty(t, decrypted, "Payload %d should decrypt to empty", i)
		} else {
			assert.Equal(t, payload, decrypted, "Payload %d should decrypt correctly", i)
		}
	}
}

// TestKeyHierarchyStructure tests KeyHierarchy fields
func TestKeyHierarchyStructure(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	hierarchy, err := km.ProvisionTenant("struct-tenant")
	require.NoError(t, err)

	// Verify structure
	assert.NotEmpty(t, hierarchy.TenantID)
	assert.NotEmpty(t, hierarchy.RootKeyID)
	assert.NotEmpty(t, hierarchy.KMSKeyID)
	assert.NotEmpty(t, hierarchy.Algorithm)
	assert.False(t, hierarchy.CreatedAt.IsZero())
}

// TestWrappedDEKStructure tests WrappedDEK fields
func TestWrappedDEKStructure(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	hierarchy, err := km.ProvisionTenant("wrapped-dek-tenant")
	require.NoError(t, err)

	wrappedDEK, err := km.GenerateDataEncryptionKey("wrapped-dek-tenant", "artifact-wrapped", hierarchy.KMSKeyID)
	require.NoError(t, err)

	assert.NotEmpty(t, wrappedDEK.TenantID)
	assert.NotEmpty(t, wrappedDEK.ArtifactID)
	assert.NotEmpty(t, wrappedDEK.WrappedKey)
	assert.NotEmpty(t, wrappedDEK.Algorithm)
	assert.False(t, wrappedDEK.CreatedAt.IsZero())
}

// TestEncryptedArtifactStructure tests EncryptedArtifact fields
func TestEncryptedArtifactStructure(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	hierarchy, err := km.ProvisionTenant("enc-art-tenant")
	require.NoError(t, err)

	plaintext := []byte("test data")
	encrypted, err := km.EncryptArtifact("enc-art-tenant", "artifact-struct", hierarchy.KMSKeyID, plaintext)
	require.NoError(t, err)

	assert.NotEmpty(t, encrypted.Ciphertext)
	assert.NotNil(t, encrypted.DEK)
	assert.NotEmpty(t, encrypted.KMSKeyID)
	assert.NotEmpty(t, encrypted.Algorithm)
}

// TestMultipleTenantsIsolation tests that tenants are properly isolated
func TestMultipleTenantsIsolation(t *testing.T) {
	t.Parallel()

	km := setupTenantKeyManager(t)

	// Provision two tenants
	tenant1, err := km.ProvisionTenant("isolation-tenant-1")
	require.NoError(t, err)

	tenant2, err := km.ProvisionTenant("isolation-tenant-2")
	require.NoError(t, err)

	// Encrypt with tenant 1
	plaintext := []byte("cross-tenant data")
	encrypted, err := km.EncryptArtifact("isolation-tenant-1", "cross-artifact", tenant1.KMSKeyID, plaintext)
	require.NoError(t, err)

	// Try to decrypt with tenant 2's key (should fail)
	_, err = km.DecryptArtifact(encrypted, tenant2.KMSKeyID)
	assert.Error(t, err)
}