/**
 * Copyright © 2012-2025 by Data Virtuality GmbH
 */
package org.teiid.jdbc;

import org.teiid.core.crypto.FipsUtil;
import org.teiid.core.crypto.FipsUtil.FipsMode;
import org.teiid.net.TeiidURL;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.Provider;
import java.security.PublicKey;
import java.security.Security;
import java.security.interfaces.DSAPublicKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.ECParameterSpec;
import java.sql.SQLException;
import java.util.Properties;
import java.util.logging.Logger;

import static org.teiid.core.crypto.FipsUtil.BCFIPS_PROVIDER;
import static org.teiid.core.crypto.FipsUtil.BC_KEYSTORE_TYPE;

/**
 * Manager class for FIPS mode functionality.
 * Handles initialization and registration of FIPS-compliant security providers.
 *
 * @author Roman Pedchenko
 * @date 2025.08.07
 */
public class FipsManager {
    private static final Logger logger = Logger.getLogger(FipsManager.class.getName());

    private static boolean fipsInitialized = false;
    private static Provider fipsProvider = null;
    private static Provider fipsJsseProvider = null;

    /**
     * Initialize FIPS mode if enabled in the properties.
     *
     * @param props Connection properties
     * @throws SQLException if FIPS mode initialization fails
     */
    public static synchronized void initializeFipsMode(Properties props) throws SQLException {
        if (props == null) {
            removeFipsProvider();
            return;
        }

        FipsMode fipsMode = FipsMode.of(props.getProperty(TeiidURL.CONNECTION.FIPS));
        if (!fipsMode.isEnabled()) {
            removeFipsProvider();
            return;
        }

        String providerName = props.getProperty(TeiidURL.CONNECTION.FIPS_PROVIDER, BCFIPS_PROVIDER);

        try {
            registerFipsProvider(fipsMode, providerName);
            FipsUtil.setFipsMode(true, providerName);
            logger.fine(String.format(JDBCPlugin.Util.getString("fips_mode_initialized"), providerName));
        } catch (Exception e) {
            throw new SQLException(String.format(JDBCPlugin.Util.getString("fips_mode_init_failed"), providerName, e.getMessage()), e);
        }
    }

    /**
     * Register the specified FIPS provider.
     *
     * @param fipsMode     fips mode
     * @param providerName Name of the FIPS provider to register
     * @throws Exception if provider registration fails
     */
    private static void registerFipsProvider(FipsMode fipsMode, String providerName) throws Exception {
        if (fipsInitialized && fipsProvider != null && fipsProvider.getName().equals(providerName)) {
            return;
        }

        logger.fine(String.format(JDBCPlugin.Util.getString("fips_initializing_provider"), providerName));

        Provider provider;
        Provider jsseProvider = null;
        if (BCFIPS_PROVIDER.equals(providerName)) {
            try {
                Class<?> providerClass = Class.forName("org.bouncycastle.jcajce.provider.BouncyCastleFipsProvider");
                provider = (Provider) providerClass.getDeclaredConstructor().newInstance();
                try {
                    Class<?> fipsStatusClass = Class.forName("org.bouncycastle.crypto.fips.FipsStatus");
                    Object isReadyObj = fipsStatusClass.getMethod("isReady").invoke(null);
                    if (isReadyObj instanceof Boolean isReady && !isReady) {
                        throw new Exception(JDBCPlugin.Util.getString("fips_bcfips_not_ready"));
                    }
                } catch (ClassNotFoundException e) {
                    logger.warning(String.format(JDBCPlugin.Util.getString("fips_could_not_verify_status"), e.getMessage()));
                }
                try {
                    providerClass = Class.forName("org.bouncycastle.jsse.provider.BouncyCastleJsseProvider");
                    jsseProvider = (Provider) providerClass.getDeclaredConstructor(boolean.class).newInstance(true);
                } catch (ClassNotFoundException e) {
                    throw new Exception(String.format(JDBCPlugin.Util.getString("fips_bctls_fips_not_found"), e.getMessage()));
                }
            } catch (ClassNotFoundException e) {
                throw new Exception(JDBCPlugin.Util.getString("fips_bcfips_not_found"), e);
            }
        } else if (providerName.startsWith("SunPKCS11-")) {
            provider = Security.getProvider(providerName);
            if (provider == null) {
                throw new Exception(String.format(JDBCPlugin.Util.getString("fips_sunpkcs11_not_found"), providerName));
            }
        } else {
            provider = Security.getProvider(providerName);
            if (provider == null) {
                throw new Exception(String.format(JDBCPlugin.Util.getString("fips_provider_not_found"), providerName));
            }
        }

        Security.insertProviderAt(provider, 1);
        fipsProvider = provider;
        if (jsseProvider != null) {
            Security.insertProviderAt(jsseProvider, 2);
            fipsJsseProvider = jsseProvider;
        }
        fipsInitialized = true;

        logger.fine(String.format(JDBCPlugin.Util.getString("fips_initialized_success"), provider.getName()));
    }

    /**
     * Check if FIPS mode is currently enabled.
     *
     * @return true if FIPS mode is enabled, false otherwise
     */
    public static boolean isFipsModeEnabled() {
        return fipsInitialized && fipsProvider != null;
    }

    public static boolean isFipsJsseModeEnabled() {
        return isFipsModeEnabled() && fipsJsseProvider != null;
    }

    /**
     * Get the name of the current FIPS provider.
     *
     * @return Name of the current FIPS provider, or null if FIPS mode is not enabled
     */
    public static String getCurrentFipsProvider() {
        return fipsProvider != null ? fipsProvider.getName() : null;
    }

    public static String getCurrentJsseProvider() {
        return fipsJsseProvider != null ? fipsJsseProvider.getName() : null;
    }

    public static void removeFipsProvider() {
        if (isFipsModeEnabled()) {
            Security.removeProvider(fipsProvider.getName());
            if (fipsJsseProvider != null)
                Security.removeProvider(fipsJsseProvider.getName());
        }
        fipsInitialized = false;
        fipsProvider = null;
        fipsJsseProvider = null;
        FipsUtil.setFipsMode(false, null);
    }

    public static void validateContext(SSLContext context) throws GeneralSecurityException {
        if (context == null)
            return;
        if (getCurrentJsseProvider() != null && !context.getProvider().getName().equals(getCurrentJsseProvider()))
            throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_context_mismatch"));
    }

    public static void validateKeystoreType(boolean checkKeystoreType, String keystoreType) throws GeneralSecurityException {
        if (checkKeystoreType && BCFIPS_PROVIDER.equalsIgnoreCase(getCurrentFipsProvider()) && !BC_KEYSTORE_TYPE.equalsIgnoreCase(keystoreType))
            throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_strict_requires_bcfks"));
    }

    public static void configureSocket(FipsMode fipsMode, SSLSocket socket) {
        if (fipsMode.isEnabled()) {
            socket.addHandshakeCompletedListener(event -> {
                logger.info(String.format(JDBCPlugin.Util.getString("fips_handshake_completed"),
                        event.getSession().getProtocol(),
                        event.getSession().getCipherSuite()));
                try {
                    java.security.cert.Certificate[] peer;
                    try {
                        peer = event.getSession().getPeerCertificates();
                    } catch (javax.net.ssl.SSLPeerUnverifiedException e) {
                        // no peer authentication, nothing to validate
                        return;
                    }
                    if (peer != null && peer.length > 0 && peer[0] instanceof java.security.cert.X509Certificate cert)
                        validatePublicKeyStrength(cert.getPublicKey());
                } catch (GeneralSecurityException gse) {
                    try {
                        socket.close();
                    } catch (IOException ignored) { }
                    throw new RuntimeException(gse.getMessage(), gse);
                }
            });
        }
    }

    /**
     * Validate that the given public key meets minimal FIPS strength requirements (>= 112 bits of security).
     * - RSA: modulus length >= 2048
     * - DSA: (L,N) in {(2048,224), (2048,256), (3072,256)}
     * - ECDSA/EdDSA: order bit length >= 224
     */
    public static void validatePublicKeyStrength(PublicKey key) throws GeneralSecurityException {
        if (key == null) return;

        String algo = key.getAlgorithm();
        if (key instanceof RSAPublicKey rsa) {
            int bits = rsa.getModulus().bitLength();
            if (bits < 2048)
                throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_rsa_wrong_size"));
            return;
        }

        if (key instanceof DSAPublicKey dsa) {
            int L = dsa.getParams().getP().bitLength();
            int N = dsa.getParams().getQ().bitLength();
            boolean ok = (L == 2048 && (N == 224 || N == 256)) || (L == 3072 && N == 256);
            if (!ok)
                throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_dsa_wrong_size"));
            return;
        }

        // Accept EdDSA (Ed25519/Ed448). If EdDSA types are unavailable, check by algorithm name.
        if ("Ed25519".equalsIgnoreCase(algo) || "Ed448".equalsIgnoreCase(algo) || algo.toUpperCase().startsWith("ED"))
            // Ed25519/Ed448 both meet >=224-bit order length
            return;

        if (key instanceof ECPublicKey ec) {
            ECParameterSpec spec = ec.getParams();
            if (spec != null) {
                int nBits = spec.getOrder() != null ? spec.getOrder().bitLength() : 0;
                if (nBits == 0 && spec.getCurve() != null && spec.getCurve().getField() != null) {
                    // Fallback to field size
                    try { nBits = spec.getCurve().getField().getFieldSize(); } catch (Exception ignored) {}
                }
                if (nBits < 224)
                    throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_ec_wrong_size"));
                return;
            }
            // If no params, fail conservatively
            throw new GeneralSecurityException(JDBCPlugin.Util.getString("fips_ec_missing"));
        }
        // Unknown algorithm - do not fail hard, but be conservative and allow (could be handled elsewhere)
    }
}
