-- Migration 017: Multi-Factor Authentication (TOTP) Implementation
-- CRITICAL SECURITY: TOTP-based MFA for financial system compliance
-- CVSS 9.1 vulnerability remediation: Implement mandatory MFA for all users

-- Enable required extensions
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- ============================================================================
-- MFA Configuration Table
-- ============================================================================
CREATE TABLE IF NOT EXISTS mfa_config (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL UNIQUE REFERENCES users(id) ON DELETE CASCADE,
    
    -- TOTP Configuration (RFC 6238)
    totp_secret_encrypted BYTEA NOT NULL, -- Encrypted TOTP secret (Base32)
    totp_algorithm VARCHAR(10) NOT NULL DEFAULT 'SHA1', -- SHA1, SHA256, SHA512
    totp_digits INTEGER NOT NULL DEFAULT 6, -- 6 or 8 digits
    totp_period INTEGER NOT NULL DEFAULT 30, -- Time step in seconds (usually 30)
    
    -- MFA Status
    is_enabled BOOLEAN NOT NULL DEFAULT false,
    is_verified BOOLEAN NOT NULL DEFAULT false, -- True after first successful verification
    
    -- Enrollment
    enrolled_at TIMESTAMP,
    verified_at TIMESTAMP,
    last_used_at TIMESTAMP,
    
    -- Recovery and Security
    backup_codes_remaining INTEGER NOT NULL DEFAULT 10,
    failed_verification_attempts INTEGER NOT NULL DEFAULT 0,
    last_failed_attempt_at TIMESTAMP,
    locked_until TIMESTAMP,
    
    -- Device Trust (optional future enhancement)
    trusted_device_ids JSONB DEFAULT '[]'::jsonb,
    
    -- Audit fields
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP NOT NULL DEFAULT NOW(),
    created_by VARCHAR(255),
    updated_by VARCHAR(255),
    
    -- Constraints
    CONSTRAINT mfa_config_totp_algorithm_check CHECK (totp_algorithm IN ('SHA1', 'SHA256', 'SHA512')),
    CONSTRAINT mfa_config_totp_digits_check CHECK (totp_digits IN (6, 8)),
    CONSTRAINT mfa_config_totp_period_check CHECK (totp_period > 0 AND totp_period <= 60),
    CONSTRAINT mfa_config_backup_codes_check CHECK (backup_codes_remaining >= 0 AND backup_codes_remaining <= 20)
);

-- ============================================================================
-- MFA Backup Codes Table
-- ============================================================================
CREATE TABLE IF NOT EXISTS mfa_backup_codes (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    
    -- Code storage (hashed for security)
    code_hash VARCHAR(64) NOT NULL UNIQUE, -- SHA-256 hash of the backup code
    code_hint VARCHAR(10) NOT NULL, -- First 4 chars for user reference (e.g., "A3F2...")
    
    -- Usage tracking
    is_used BOOLEAN NOT NULL DEFAULT false,
    used_at TIMESTAMP,
    used_from_ip INET,
    
    -- Expiration
    expires_at TIMESTAMP NOT NULL,
    
    -- Audit
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    
    -- Constraints
    CONSTRAINT mfa_backup_codes_expiration_check CHECK (expires_at > created_at)
);

-- ============================================================================
-- MFA Verification Attempts Log
-- ============================================================================
CREATE TABLE IF NOT EXISTS mfa_verification_log (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    
    -- Verification details
    verification_method VARCHAR(50) NOT NULL, -- 'totp', 'backup_code', 'trusted_device'
    success BOOLEAN NOT NULL,
    
    -- Context
    ip_address INET,
    user_agent TEXT,
    location_data JSONB, -- Optional geolocation data
    
    -- TOTP-specific
    totp_code VARCHAR(8), -- Not stored in plaintext for failed attempts
    totp_drift INTEGER, -- Time drift in periods (for analysis)
    
    -- Error information
    error_code VARCHAR(100),
    error_message TEXT,
    
    -- Timestamp
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    
    -- Constraints
    CONSTRAINT mfa_verification_log_method_check CHECK (verification_method IN 
        ('totp', 'backup_code', 'trusted_device', 'recovery'))
);

-- ============================================================================
-- MFA Enrollment Sessions (temporary during setup)
-- ============================================================================
CREATE TABLE IF NOT EXISTS mfa_enrollment_sessions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    
    -- Temporary TOTP secret (during enrollment)
    temp_totp_secret_encrypted BYTEA NOT NULL,
    qr_code_data TEXT NOT NULL, -- QR code URI (otpauth://)
    
    -- Session management
    is_active BOOLEAN NOT NULL DEFAULT true,
    expires_at TIMESTAMP NOT NULL,
    completed_at TIMESTAMP,
    
    -- Verification
    verification_attempts INTEGER NOT NULL DEFAULT 0,
    max_verification_attempts INTEGER NOT NULL DEFAULT 3,
    
    -- Created timestamp
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    
    -- Constraints
    CONSTRAINT mfa_enrollment_expiration_check CHECK (expires_at > created_at),
    CONSTRAINT mfa_enrollment_attempts_check CHECK (verification_attempts >= 0)
);

-- ============================================================================
-- Indexes for Performance
-- ============================================================================
CREATE INDEX IF NOT EXISTS idx_mfa_config_user_id ON mfa_config(user_id);
CREATE INDEX IF NOT EXISTS idx_mfa_config_is_enabled ON mfa_config(is_enabled);
CREATE INDEX IF NOT EXISTS idx_mfa_config_last_used_at ON mfa_config(last_used_at);

CREATE INDEX IF NOT EXISTS idx_mfa_backup_codes_user_id ON mfa_backup_codes(user_id);
CREATE INDEX IF NOT EXISTS idx_mfa_backup_codes_code_hash ON mfa_backup_codes(code_hash);
CREATE INDEX IF NOT EXISTS idx_mfa_backup_codes_is_used ON mfa_backup_codes(is_used);
CREATE INDEX IF NOT EXISTS idx_mfa_backup_codes_expires_at ON mfa_backup_codes(expires_at);

CREATE INDEX IF NOT EXISTS idx_mfa_verification_log_user_id ON mfa_verification_log(user_id);
CREATE INDEX IF NOT EXISTS idx_mfa_verification_log_created_at ON mfa_verification_log(created_at);
CREATE INDEX IF NOT EXISTS idx_mfa_verification_log_success ON mfa_verification_log(success);
CREATE INDEX IF NOT EXISTS idx_mfa_verification_log_ip_address ON mfa_verification_log(ip_address);

CREATE INDEX IF NOT EXISTS idx_mfa_enrollment_sessions_user_id ON mfa_enrollment_sessions(user_id);
CREATE INDEX IF NOT EXISTS idx_mfa_enrollment_sessions_is_active ON mfa_enrollment_sessions(is_active);
CREATE INDEX IF NOT EXISTS idx_mfa_enrollment_sessions_expires_at ON mfa_enrollment_sessions(expires_at);

-- ============================================================================
-- Encryption Key Management Functions
-- ============================================================================

-- Function to encrypt TOTP secrets using application encryption key
-- In production, this should use a proper key management service (KMS)
CREATE OR REPLACE FUNCTION encrypt_totp_secret(
    plaintext_secret TEXT,
    encryption_key TEXT DEFAULT current_setting('app.encryption_key', true)
) RETURNS BYTEA AS $$
BEGIN
    -- Use pgcrypto for AES-256 encryption
    -- In production, encryption_key should be from Vault/KMS
    RETURN pgp_sym_encrypt(plaintext_secret, encryption_key, 'cipher-algo=aes256');
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Function to decrypt TOTP secrets
CREATE OR REPLACE FUNCTION decrypt_totp_secret(
    encrypted_secret BYTEA,
    encryption_key TEXT DEFAULT current_setting('app.encryption_key', true)
) RETURNS TEXT AS $$
BEGIN
    RETURN pgp_sym_decrypt(encrypted_secret, encryption_key);
EXCEPTION
    WHEN OTHERS THEN
        -- Log decryption failures
        RAISE EXCEPTION 'TOTP secret decryption failed: %', SQLERRM;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- ============================================================================
-- MFA Management Functions
-- ============================================================================

-- Function to hash backup codes securely
CREATE OR REPLACE FUNCTION hash_backup_code(backup_code TEXT)
RETURNS VARCHAR(64) AS $$
BEGIN
    RETURN encode(digest(backup_code, 'sha256'), 'hex');
END;
$$ LANGUAGE plpgsql IMMUTABLE;

-- Function to validate backup code
CREATE OR REPLACE FUNCTION validate_backup_code(
    p_user_id UUID,
    p_backup_code TEXT,
    p_ip_address INET DEFAULT NULL
) RETURNS BOOLEAN AS $$
DECLARE
    v_code_hash VARCHAR(64);
    v_code_id UUID;
    v_is_valid BOOLEAN := false;
BEGIN
    v_code_hash := hash_backup_code(p_backup_code);
    
    -- Check if backup code exists and is valid
    SELECT id INTO v_code_id
    FROM mfa_backup_codes
    WHERE user_id = p_user_id
        AND code_hash = v_code_hash
        AND is_used = false
        AND expires_at > NOW();
    
    IF v_code_id IS NOT NULL THEN
        -- Mark code as used
        UPDATE mfa_backup_codes
        SET is_used = true,
            used_at = NOW(),
            used_from_ip = p_ip_address
        WHERE id = v_code_id;
        
        -- Decrement backup codes remaining
        UPDATE mfa_config
        SET backup_codes_remaining = backup_codes_remaining - 1,
            last_used_at = NOW()
        WHERE user_id = p_user_id;
        
        v_is_valid := true;
        
        -- Log successful verification
        INSERT INTO mfa_verification_log (
            user_id, verification_method, success, ip_address
        ) VALUES (
            p_user_id, 'backup_code', true, p_ip_address
        );
    ELSE
        -- Log failed verification
        INSERT INTO mfa_verification_log (
            user_id, verification_method, success, ip_address, error_code
        ) VALUES (
            p_user_id, 'backup_code', false, p_ip_address, 'INVALID_BACKUP_CODE'
        );
    END IF;
    
    RETURN v_is_valid;
END;
$$ LANGUAGE plpgsql;

-- Function to check if MFA is required for user
CREATE OR REPLACE FUNCTION is_mfa_required(p_user_id UUID)
RETURNS BOOLEAN AS $$
DECLARE
    v_user_role VARCHAR(50);
    v_mfa_enabled BOOLEAN := false;
BEGIN
    -- Get user role
    SELECT role INTO v_user_role FROM users WHERE id = p_user_id;
    
    -- Check if MFA is enabled for user
    SELECT is_enabled INTO v_mfa_enabled
    FROM mfa_config
    WHERE user_id = p_user_id;
    
    -- MFA is required for all privileged roles in financial systems
    -- Also required if user has MFA enabled (cannot disable once enabled)
    RETURN (v_user_role IN ('admin', 'trader', 'risk_manager', 'compliance_officer'))
           OR COALESCE(v_mfa_enabled, false);
END;
$$ LANGUAGE plpgsql;

-- Function to record MFA verification attempt
CREATE OR REPLACE FUNCTION record_mfa_attempt(
    p_user_id UUID,
    p_method VARCHAR(50),
    p_success BOOLEAN,
    p_ip_address INET DEFAULT NULL,
    p_user_agent TEXT DEFAULT NULL,
    p_totp_drift INTEGER DEFAULT NULL,
    p_error_code VARCHAR(100) DEFAULT NULL
) RETURNS UUID AS $$
DECLARE
    v_log_id UUID;
BEGIN
    INSERT INTO mfa_verification_log (
        user_id, verification_method, success, ip_address,
        user_agent, totp_drift, error_code
    ) VALUES (
        p_user_id, p_method, p_success, p_ip_address,
        p_user_agent, p_totp_drift, p_error_code
    ) RETURNING id INTO v_log_id;
    
    -- Update MFA config based on success/failure
    IF p_success THEN
        UPDATE mfa_config
        SET last_used_at = NOW(),
            failed_verification_attempts = 0,
            locked_until = NULL
        WHERE user_id = p_user_id;
    ELSE
        -- Increment failed attempts
        UPDATE mfa_config
        SET failed_verification_attempts = failed_verification_attempts + 1,
            last_failed_attempt_at = NOW()
        WHERE user_id = p_user_id;
        
        -- Lock account after 5 failed attempts (15 minutes lockout)
        UPDATE mfa_config
        SET locked_until = NOW() + INTERVAL '15 minutes'
        WHERE user_id = p_user_id
            AND failed_verification_attempts >= 5
            AND locked_until IS NULL;
    END IF;
    
    RETURN v_log_id;
END;
$$ LANGUAGE plpgsql;

-- Function to clean up expired MFA data
CREATE OR REPLACE FUNCTION cleanup_expired_mfa_data()
RETURNS INTEGER AS $$
DECLARE
    cleanup_count INTEGER := 0;
BEGIN
    -- Delete expired backup codes
    DELETE FROM mfa_backup_codes
    WHERE expires_at < NOW() OR (is_used = true AND used_at < NOW() - INTERVAL '90 days');
    GET DIAGNOSTICS cleanup_count = ROW_COUNT;
    
    -- Delete old verification logs (keep 1 year)
    DELETE FROM mfa_verification_log
    WHERE created_at < NOW() - INTERVAL '1 year';
    
    -- Delete expired enrollment sessions
    DELETE FROM mfa_enrollment_sessions
    WHERE expires_at < NOW() OR (completed_at IS NOT NULL AND completed_at < NOW() - INTERVAL '7 days');
    
    -- Clear expired account locks
    UPDATE mfa_config
    SET locked_until = NULL
    WHERE locked_until < NOW();
    
    RETURN cleanup_count;
END;
$$ LANGUAGE plpgsql;

-- ============================================================================
-- Triggers
-- ============================================================================

CREATE TRIGGER update_mfa_config_updated_at BEFORE UPDATE ON mfa_config
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

-- ============================================================================
-- Row Level Security (RLS)
-- ============================================================================

ALTER TABLE mfa_config ENABLE ROW LEVEL SECURITY;
ALTER TABLE mfa_backup_codes ENABLE ROW LEVEL SECURITY;
ALTER TABLE mfa_verification_log ENABLE ROW LEVEL SECURITY;
ALTER TABLE mfa_enrollment_sessions ENABLE ROW LEVEL SECURITY;

-- Users can only access their own MFA data
CREATE POLICY mfa_config_self_access ON mfa_config
    FOR ALL TO foxhunt_user
    USING (user_id = current_setting('app.current_user_id')::uuid OR
           current_setting('app.current_user_role') = 'admin');

CREATE POLICY mfa_backup_codes_self_access ON mfa_backup_codes
    FOR ALL TO foxhunt_user
    USING (user_id = current_setting('app.current_user_id')::uuid OR
           current_setting('app.current_user_role') = 'admin');

CREATE POLICY mfa_verification_log_self_access ON mfa_verification_log
    FOR SELECT TO foxhunt_user
    USING (user_id = current_setting('app.current_user_id')::uuid OR
           current_setting('app.current_user_role') = 'admin');

CREATE POLICY mfa_enrollment_sessions_self_access ON mfa_enrollment_sessions
    FOR ALL TO foxhunt_user
    USING (user_id = current_setting('app.current_user_id')::uuid);

-- ============================================================================
-- Grants
-- ============================================================================

GRANT SELECT, INSERT, UPDATE ON mfa_config TO foxhunt_user;
GRANT SELECT, INSERT, UPDATE ON mfa_backup_codes TO foxhunt_user;
GRANT SELECT, INSERT ON mfa_verification_log TO foxhunt_user;
GRANT SELECT, INSERT, UPDATE, DELETE ON mfa_enrollment_sessions TO foxhunt_user;

-- ============================================================================
-- Comments
-- ============================================================================

COMMENT ON TABLE mfa_config IS 'Multi-factor authentication configuration (TOTP) for users';
COMMENT ON TABLE mfa_backup_codes IS 'One-time backup codes for MFA recovery';
COMMENT ON TABLE mfa_verification_log IS 'Audit log for MFA verification attempts';
COMMENT ON TABLE mfa_enrollment_sessions IS 'Temporary sessions during MFA enrollment';

COMMENT ON FUNCTION encrypt_totp_secret IS 'Encrypt TOTP secrets using AES-256';
COMMENT ON FUNCTION decrypt_totp_secret IS 'Decrypt TOTP secrets for verification';
COMMENT ON FUNCTION hash_backup_code IS 'Hash backup codes using SHA-256';
COMMENT ON FUNCTION validate_backup_code IS 'Validate and consume a backup code';
COMMENT ON FUNCTION is_mfa_required IS 'Check if MFA is required for a user';
COMMENT ON FUNCTION record_mfa_attempt IS 'Record MFA verification attempt with security tracking';
COMMENT ON FUNCTION cleanup_expired_mfa_data IS 'Clean up expired MFA sessions and codes';