fetch_ml/internal/crypto/tenant_keys.go

// Package crypto provides tenant-scoped encryption key management for multi-tenant deployments.
// This implements Phase 9.4: Per-Tenant Encryption Keys with KMS integration per ADR-012 through ADR-015.
package crypto

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/jfraeys/fetch_ml/internal/audit"
	"github.com/jfraeys/fetch_ml/internal/crypto/kms"
)

// KeyHierarchy defines the tenant key structure
// Root Key (per tenant in KMS) -> Data Encryption Keys (per artifact, cached per ADR-012)
type KeyHierarchy struct {
	TenantID  string    `json:"tenant_id"`
	RootKeyID string    `json:"root_key_id"`
	KMSKeyID  string    `json:"kms_key_id"` // External KMS key identifier per ADR-014
	CreatedAt time.Time `json:"created_at"`
	Algorithm string    `json:"algorithm"` // Always "AES-256-GCM"
}

// TenantKeyManager manages per-tenant encryption keys using external KMS per ADR-012 through ADR-015.
// Root keys are stored in the KMS; DEKs are generated locally and cached.
type TenantKeyManager struct {
	kms    kms.KMSProvider // External KMS for root key operations
	cache  *kms.DEKCache   // In-process DEK cache per ADR-012
	config kms.Config      // KMS configuration
	ctx    context.Context // Background context for operations
	audit  *audit.Logger   // Audit logger for key operations per ADR-012
}

// NewTenantKeyManager creates a new tenant key manager with KMS integration.
func NewTenantKeyManager(provider kms.KMSProvider, cache *kms.DEKCache, config kms.Config, auditLogger *audit.Logger) *TenantKeyManager {
	return &TenantKeyManager{
		kms:    provider,
		cache:  cache,
		config: config,
		ctx:    context.Background(),
		audit:  auditLogger,
	}
}

// ProvisionTenant creates a new root key for a tenant in the KMS.
func (km *TenantKeyManager) ProvisionTenant(tenantID string) (*KeyHierarchy, error) {
	if strings.TrimSpace(tenantID) == "" {
		return nil, fmt.Errorf("tenant ID cannot be empty")
	}

	// Create KMS key for tenant
	kmsKeyID, err := km.kms.CreateKey(km.ctx, tenantID)
	if err != nil {
		return nil, fmt.Errorf("failed to create KMS key: %w", err)
	}

	// Create key ID from hash of tenant ID + timestamp
	h := sha256.Sum256([]byte(tenantID + time.Now().String()))
	rootKeyID := hex.EncodeToString(h[:8]) // First 8 bytes as ID

	// Log key creation per ADR-012
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSKeyCreate, tenantID, "", kmsKeyID, true, "")
	}

	return &KeyHierarchy{
		TenantID:  tenantID,
		RootKeyID: rootKeyID,
		KMSKeyID:  kmsKeyID,
		CreatedAt: time.Now().UTC(),
		Algorithm: "AES-256-GCM",
	}, nil
}

// RotateTenantKey rotates the root key for a tenant.
// Creates new KMS key and schedules deletion of old key per ADR-015.
func (km *TenantKeyManager) RotateTenantKey(tenantID string, hierarchy *KeyHierarchy) (*KeyHierarchy, error) {
	// Schedule deletion of old key (90 day window per ADR-015)
	_, err := km.kms.ScheduleKeyDeletion(km.ctx, hierarchy.KMSKeyID, 90)
	if err != nil {
		return nil, fmt.Errorf("failed to schedule old key deletion: %w", err)
	}

	// Flush DEK cache for this tenant
	km.cache.Flush(tenantID)

	// Provision new key
	newHierarchy, err := km.ProvisionTenant(tenantID)
	if err != nil {
		return nil, err
	}

	// Log key rotation per ADR-012
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSKeyRotate, tenantID, "", newHierarchy.KMSKeyID, true, "")
	}

	return newHierarchy, nil
}

// RevokeTenant disables and schedules deletion of all keys for a tenant.
// This effectively makes all encrypted data inaccessible per ADR-015.
func (km *TenantKeyManager) RevokeTenant(hierarchy *KeyHierarchy) error {
	// Immediately disable the key per ADR-015
	if err := km.kms.DisableKey(km.ctx, hierarchy.KMSKeyID); err != nil {
		return fmt.Errorf("failed to disable key: %w", err)
	}

	// Log key disable per ADR-015
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSKeyDisable, hierarchy.TenantID, "", hierarchy.KMSKeyID, true, "")
	}

	// Schedule hard deletion after 90 days per ADR-015
	_, err := km.kms.ScheduleKeyDeletion(km.ctx, hierarchy.KMSKeyID, 90)
	if err != nil {
		return fmt.Errorf("failed to schedule key deletion: %w", err)
	}

	// Log key deletion scheduled per ADR-015
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSKeyDelete, hierarchy.TenantID, "", hierarchy.KMSKeyID, true, "")
	}

	// Flush DEK cache for this tenant
	km.cache.Flush(hierarchy.TenantID)

	return nil
}

// GenerateDataEncryptionKey creates a unique DEK for an artifact.
// The DEK is wrapped (encrypted) under the tenant's KMS root key.
func (km *TenantKeyManager) GenerateDataEncryptionKey(tenantID, artifactID, kmsKeyID string) (*WrappedDEK, error) {
	// Generate unique DEK (32 bytes for AES-256)
	dek := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, dek); err != nil {
		return nil, fmt.Errorf("failed to generate DEK: %w", err)
	}

	// Wrap DEK with KMS root key
	wrappedKey, err := km.wrapKeyWithKMS(km.ctx, kmsKeyID, dek)
	if err != nil {
		return nil, fmt.Errorf("failed to wrap DEK: %w", err)
	}

	// Store DEK in cache for future use per ADR-012
	if err := km.cache.Put(tenantID, artifactID, kmsKeyID, dek); err != nil {
		// Log but don't fail - caching is optimization
		_ = err
	}

	// Clear plaintext DEK from memory
	for i := range dek {
		dek[i] = 0
	}

	return &WrappedDEK{
		TenantID:   tenantID,
		ArtifactID: artifactID,
		WrappedKey: wrappedKey,
		Algorithm:  "AES-256-GCM",
		CreatedAt:  time.Now().UTC(),
	}, nil
}

// wrapKeyWithKMS encrypts a key using the KMS.
func (km *TenantKeyManager) wrapKeyWithKMS(ctx context.Context, kmsKeyID string, keyToWrap []byte) (string, error) {
	ciphertext, err := km.kms.Encrypt(ctx, kmsKeyID, keyToWrap)
	if err != nil {
		return "", fmt.Errorf("KMS encrypt failed: %w", err)
	}
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// UnwrapDataEncryptionKey decrypts a wrapped DEK using the tenant's KMS root key.
// Per ADR-012/013: Checks cache first, falls back to KMS with fail-closed grace window.
func (km *TenantKeyManager) UnwrapDataEncryptionKey(wrappedDEK *WrappedDEK, kmsKeyID string) ([]byte, error) {
	// Try cache first per ADR-012 - include KMSKeyID in cache key for isolation
	if dek, ok := km.cache.Get(wrappedDEK.TenantID, wrappedDEK.ArtifactID, kmsKeyID, false); ok {
		return dek, nil
	}

	// Check KMS health for grace window determination per ADR-013
	kmsHealthy := km.kms.HealthCheck(km.ctx) == nil

	// If KMS is unavailable and we have a cached entry in grace window, use it per ADR-013
	if !kmsHealthy {
		if dek, ok := km.cache.Get(wrappedDEK.TenantID, wrappedDEK.ArtifactID, kmsKeyID, true); ok {
			// Grace window DEK returned - logged by caller
			return dek, nil
		}
		// No cached DEK and KMS unavailable - fail closed per ADR-013
		return nil, fmt.Errorf("KMS unavailable and no cached DEK (fail-closed per ADR-013)")
	}

	// Unwrap via KMS
	ciphertext, err := base64.StdEncoding.DecodeString(wrappedDEK.WrappedKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decode wrapped key: %w", err)
	}

	dek, err := km.kms.Decrypt(km.ctx, kmsKeyID, ciphertext)
	if err != nil {
		return nil, fmt.Errorf("KMS decrypt failed: %w", err)
	}

	// Store in cache for future use per ADR-012 - include KMSKeyID
	if err := km.cache.Put(wrappedDEK.TenantID, wrappedDEK.ArtifactID, kmsKeyID, dek); err != nil {
		// Log but don't fail - caching is optimization
		_ = err
	}

	return dek, nil
}

// WrappedDEK represents a data encryption key wrapped under a tenant root key
type WrappedDEK struct {
	TenantID   string    `json:"tenant_id"`
	ArtifactID string    `json:"artifact_id"`
	WrappedKey string    `json:"wrapped_key"` // base64 encoded
	Algorithm  string    `json:"algorithm"`
	CreatedAt  time.Time `json:"created_at"`
}

// NewTestTenantKeyManager creates a tenant key manager with memory provider for testing.
// This provides backward compatibility for existing tests.
func NewTestTenantKeyManager(auditLogger *audit.Logger) *TenantKeyManager {
	provider := kms.NewMemoryProvider()
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	config := kms.Config{Provider: kms.ProviderTypeMemory}
	return NewTenantKeyManager(provider, cache, config, auditLogger)
}

// EncryptArtifact encrypts artifact data using a tenant-specific DEK.
func (km *TenantKeyManager) EncryptArtifact(tenantID, artifactID, kmsKeyID string, plaintext []byte) (*EncryptedArtifact, error) {
	// Generate a new DEK for this artifact
	wrappedDEK, err := km.GenerateDataEncryptionKey(tenantID, artifactID, kmsKeyID)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, err
	}

	// Get the DEK (from cache or unwrap)
	dek, err := km.UnwrapDataEncryptionKey(wrappedDEK, kmsKeyID)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, err
	}
	defer func() {
		// Clear DEK from memory after use
		for i := range dek {
			dek[i] = 0
		}
	}()

	// Encrypt the data with the DEK
	block, err := aes.NewCipher(dek)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)

	// Log successful encryption per ADR-012
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSEncrypt, tenantID, artifactID, kmsKeyID, true, "")
	}

	return &EncryptedArtifact{
		Ciphertext: base64.StdEncoding.EncodeToString(ciphertext),
		DEK:        wrappedDEK,
		KMSKeyID:   kmsKeyID,
		Algorithm:  "AES-256-GCM",
	}, nil
}

// DecryptArtifact decrypts artifact data using its wrapped DEK.
func (km *TenantKeyManager) DecryptArtifact(encrypted *EncryptedArtifact, kmsKeyID string) ([]byte, error) {
	tenantID := encrypted.DEK.TenantID
	artifactID := encrypted.DEK.ArtifactID

	// Unwrap the DEK
	dek, err := km.UnwrapDataEncryptionKey(encrypted.DEK, kmsKeyID)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to unwrap DEK: %w", err)
	}
	defer func() {
		for i := range dek {
			dek[i] = 0
		}
	}()

	// Decrypt the data
	ciphertext, err := base64.StdEncoding.DecodeString(encrypted.Ciphertext)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to decode ciphertext: %w", err)
	}

	block, err := aes.NewCipher(dek)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		err := fmt.Errorf("ciphertext too short")
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, err
	}

	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		if km.audit != nil {
			km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, false, err.Error())
		}
		return nil, err
	}

	// Log successful decryption per ADR-012
	if km.audit != nil {
		km.audit.LogKMSOperation(audit.EventKMSDecrypt, tenantID, artifactID, kmsKeyID, true, "")
	}

	return plaintext, nil
}

// EncryptedArtifact represents an encrypted artifact with its wrapped DEK
type EncryptedArtifact struct {
	Ciphertext string      `json:"ciphertext"` // base64 encoded
	DEK        *WrappedDEK `json:"dek"`
	KMSKeyID   string      `json:"kms_key_id"` // Per ADR-014
	Algorithm  string      `json:"algorithm"`
}

// AuditLogEntry represents an audit log entry for encryption/decryption operations
type AuditLogEntry struct {
	Timestamp  time.Time `json:"timestamp"`
	Operation  string    `json:"operation"` // "encrypt", "decrypt", "key_rotation"
	TenantID   string    `json:"tenant_id"`
	ArtifactID string    `json:"artifact_id,omitempty"`
	KeyID      string    `json:"key_id"`
	Success    bool      `json:"success"`
	Error      string    `json:"error,omitempty"`
}