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.
package crypto

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

// KeyHierarchy defines the tenant key structure
// Root Key (per tenant) -> Data Encryption Keys (per artifact)
type KeyHierarchy struct {
	TenantID   string    `json:"tenant_id"`
	RootKeyID  string    `json:"root_key_id"`
	CreatedAt  time.Time `json:"created_at"`
	Algorithm  string    `json:"algorithm"` // Always "AES-256-GCM"
}

// TenantKeyManager manages per-tenant encryption keys
// In production, root keys should be stored in a KMS (HashiCorp Vault, AWS KMS, etc.)
type TenantKeyManager struct {
	// In-memory store for development; use external KMS in production
	rootKeys map[string][]byte // tenantID -> root key
}

// NewTenantKeyManager creates a new tenant key manager
func NewTenantKeyManager() *TenantKeyManager {
	return &TenantKeyManager{
		rootKeys: make(map[string][]byte),
	}
}

// ProvisionTenant creates a new root key for a tenant
// In production, this would call out to a KMS to create a key
func (km *TenantKeyManager) ProvisionTenant(tenantID string) (*KeyHierarchy, error) {
	if strings.TrimSpace(tenantID) == "" {
		return nil, fmt.Errorf("tenant ID cannot be empty")
	}

	// Generate root key (32 bytes for AES-256)
	rootKey := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, rootKey); err != nil {
		return nil, fmt.Errorf("failed to generate root key: %w", err)
	}

	// Create key ID from hash of key (for reference, not for key derivation)
	h := sha256.Sum256(rootKey)
	rootKeyID := hex.EncodeToString(h[:8]) // First 8 bytes as ID

	// Store root key
	km.rootKeys[tenantID] = rootKey

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

// RotateTenantKey rotates the root key for a tenant
// Existing data must be re-encrypted with the new key
func (km *TenantKeyManager) RotateTenantKey(tenantID string) (*KeyHierarchy, error) {
	// Delete old key
	delete(km.rootKeys, tenantID)

	// Provision new key
	return km.ProvisionTenant(tenantID)
}

// RevokeTenant removes all keys for a tenant
// This effectively makes all encrypted data inaccessible
func (km *TenantKeyManager) RevokeTenant(tenantID string) error {
	if _, exists := km.rootKeys[tenantID]; !exists {
		return fmt.Errorf("tenant %s not found", tenantID)
	}

	// Overwrite key before deleting (best effort)
	key := km.rootKeys[tenantID]
	for i := range key {
		key[i] = 0
	}
	delete(km.rootKeys, tenantID)

	return nil
}

// GenerateDataEncryptionKey creates a unique DEK for an artifact
// The DEK is wrapped (encrypted) under the tenant's root key
func (km *TenantKeyManager) GenerateDataEncryptionKey(tenantID string, artifactID string) (*WrappedDEK, error) {
	rootKey, exists := km.rootKeys[tenantID]
	if !exists {
		return nil, fmt.Errorf("no root key found for tenant %s", tenantID)
	}

	// 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 root key
	wrappedKey, err := km.wrapKey(rootKey, dek)
	if err != nil {
		return nil, fmt.Errorf("failed to wrap DEK: %w", 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
}

// UnwrapDataEncryptionKey decrypts a wrapped DEK using the tenant's root key
func (km *TenantKeyManager) UnwrapDataEncryptionKey(wrappedDEK *WrappedDEK) ([]byte, error) {
	rootKey, exists := km.rootKeys[wrappedDEK.TenantID]
	if !exists {
		return nil, fmt.Errorf("no root key found for tenant %s", wrappedDEK.TenantID)
	}

	return km.unwrapKey(rootKey, wrappedDEK.WrappedKey)
}

// 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"`
}

// wrapKey encrypts a key using AES-256-GCM with the provided root key
func (km *TenantKeyManager) wrapKey(rootKey, keyToWrap []byte) (string, error) {
	block, err := aes.NewCipher(rootKey)
	if err != nil {
		return "", fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("failed to create GCM: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", fmt.Errorf("failed to generate nonce: %w", err)
	}

	ciphertext := gcm.Seal(nonce, nonce, keyToWrap, nil)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// unwrapKey decrypts a wrapped key using AES-256-GCM
func (km *TenantKeyManager) unwrapKey(rootKey []byte, wrappedKey string) ([]byte, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(wrappedKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decode wrapped key: %w", err)
	}

	block, err := aes.NewCipher(rootKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	return gcm.Open(nil, nonce, ciphertext, nil)
}

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

	// Unwrap the DEK for use
	dek, err := km.UnwrapDataEncryptionKey(wrappedDEK)
	if err != nil {
		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 {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

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

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

// DecryptArtifact decrypts artifact data using its wrapped DEK
func (km *TenantKeyManager) DecryptArtifact(encrypted *EncryptedArtifact) ([]byte, error) {
	// Unwrap the DEK
	dek, err := km.UnwrapDataEncryptionKey(encrypted.DEK)
	if err != nil {
		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 {
		return nil, fmt.Errorf("failed to decode ciphertext: %w", err)
	}

	block, err := aes.NewCipher(dek)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	return gcm.Open(nil, nonce, ciphertext, nil)
}

// EncryptedArtifact represents an encrypted artifact with its wrapped DEK
type EncryptedArtifact struct {
	Ciphertext string       `json:"ciphertext"` // base64 encoded
	DEK        *WrappedDEK  `json:"dek"`
	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"`
}