-- Enable pgcrypto for MFA secret encryption
-- Wave 121 Agent 1: Fix CRITICAL security blocker (MFA plaintext secrets)
-- CVSS Impact: Prevents unauthorized access to TOTP secrets

-- Enable pgcrypto extension for encryption functions
CREATE EXTENSION IF NOT EXISTS pgcrypto;

-- Create encryption key table (single row, versioned keys for rotation)
CREATE TABLE IF NOT EXISTS mfa_encryption_keys (
    id SERIAL PRIMARY KEY,
    key_version INTEGER UNIQUE NOT NULL DEFAULT 1,
    encryption_key BYTEA NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() NOT NULL,
    rotated_at TIMESTAMP WITH TIME ZONE,
    is_active BOOLEAN DEFAULT TRUE NOT NULL,
    CONSTRAINT ensure_single_active_key CHECK (
        is_active = TRUE OR rotated_at IS NOT NULL
    )
);

-- Create index for active key lookup
CREATE UNIQUE INDEX IF NOT EXISTS idx_mfa_encryption_keys_active
ON mfa_encryption_keys(key_version) WHERE is_active = TRUE;

-- Generate initial encryption key (256-bit AES key)
INSERT INTO mfa_encryption_keys (key_version, encryption_key, is_active)
VALUES (
    1,
    gen_random_bytes(32),  -- 256-bit key
    TRUE
)
ON CONFLICT (key_version) DO NOTHING;

-- Function to get active encryption key
CREATE OR REPLACE FUNCTION get_active_mfa_key()
RETURNS BYTEA AS $$
DECLARE
    v_key BYTEA;
BEGIN
    SELECT encryption_key INTO v_key
    FROM mfa_encryption_keys
    WHERE is_active = TRUE
    ORDER BY key_version DESC
    LIMIT 1;

    IF v_key IS NULL THEN
        RAISE EXCEPTION 'No active MFA encryption key found';
    END IF;

    RETURN v_key;
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;

-- Function to encrypt TOTP secret using pgcrypto
-- Uses AES-256-CBC with the active encryption key
CREATE OR REPLACE FUNCTION encrypt_mfa_secret(p_secret TEXT)
RETURNS BYTEA AS $$
DECLARE
    v_key BYTEA;
BEGIN
    v_key := get_active_mfa_key();

    -- Encrypt using AES-256-CBC
    -- Format: pgp_sym_encrypt(data, password, options)
    -- We use the raw encryption key directly for maximum security
    RETURN encrypt(p_secret::BYTEA, v_key, 'aes');
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;

-- Function to decrypt TOTP secret using pgcrypto
CREATE OR REPLACE FUNCTION decrypt_mfa_secret(p_encrypted BYTEA)
RETURNS TEXT AS $$
DECLARE
    v_key BYTEA;
    v_decrypted BYTEA;
BEGIN
    v_key := get_active_mfa_key();

    -- Decrypt using AES-256-CBC
    v_decrypted := decrypt(p_encrypted, v_key, 'aes');

    RETURN convert_from(v_decrypted, 'UTF8');
EXCEPTION
    WHEN OTHERS THEN
        RAISE EXCEPTION 'Failed to decrypt MFA secret: %', SQLERRM;
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;

-- Function to rotate encryption key (for key management)
CREATE OR REPLACE FUNCTION rotate_mfa_encryption_key()
RETURNS INTEGER AS $$
DECLARE
    v_new_key_version INTEGER;
    v_new_key BYTEA;
BEGIN
    -- Get next key version
    SELECT COALESCE(MAX(key_version), 0) + 1 INTO v_new_key_version
    FROM mfa_encryption_keys;

    -- Generate new 256-bit key
    v_new_key := gen_random_bytes(32);

    -- Mark old keys as inactive
    UPDATE mfa_encryption_keys
    SET is_active = FALSE, rotated_at = NOW()
    WHERE is_active = TRUE;

    -- Insert new key
    INSERT INTO mfa_encryption_keys (key_version, encryption_key, is_active)
    VALUES (v_new_key_version, v_new_key, TRUE);

    RAISE NOTICE 'MFA encryption key rotated to version %', v_new_key_version;

    RETURN v_new_key_version;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Grant execute permissions to the application role
GRANT EXECUTE ON FUNCTION get_active_mfa_key() TO foxhunt;
GRANT EXECUTE ON FUNCTION encrypt_mfa_secret(TEXT) TO foxhunt;
GRANT EXECUTE ON FUNCTION decrypt_mfa_secret(BYTEA) TO foxhunt;
GRANT EXECUTE ON FUNCTION rotate_mfa_encryption_key() TO foxhunt;

-- Prevent direct access to encryption keys table
REVOKE ALL ON TABLE mfa_encryption_keys FROM PUBLIC;
GRANT SELECT ON TABLE mfa_encryption_keys TO foxhunt;

-- Add comment for documentation
COMMENT ON FUNCTION encrypt_mfa_secret(TEXT) IS 'Encrypts TOTP secret using AES-256-CBC with pgcrypto. Uses active encryption key from mfa_encryption_keys table.';
COMMENT ON FUNCTION decrypt_mfa_secret(BYTEA) IS 'Decrypts TOTP secret using AES-256-CBC with pgcrypto. Requires active encryption key.';
COMMENT ON FUNCTION rotate_mfa_encryption_key() IS 'Rotates MFA encryption key. New key becomes active, old keys are marked inactive. Use during key rotation ceremonies.';

-- Verification query to ensure encryption works
DO $$
DECLARE
    v_test_encrypted BYTEA;
    v_test_decrypted TEXT;
BEGIN
    -- Test encryption
    v_test_encrypted := encrypt_mfa_secret('TEST_SECRET_12345');

    -- Test decryption
    v_test_decrypted := decrypt_mfa_secret(v_test_encrypted);

    IF v_test_decrypted = 'TEST_SECRET_12345' THEN
        RAISE NOTICE '✅ MFA encryption verification PASSED';
    ELSE
        RAISE EXCEPTION '❌ MFA encryption verification FAILED: expected "TEST_SECRET_12345", got "%"', v_test_decrypted;
    END IF;
END;
$$;