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

import javax.net.ssl.*;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;

/**
 *
 *
 * @author Roman Pedchenko
 * @date 2025.08.07
 */
public class FipsConnectionValidator {
    private static final List<String> FIPS_APPROVED_CIPHERS = List.of(
            "TLS_AES_256_GCM_SHA384",
            "TLS_AES_128_GCM_SHA256",
            "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
            "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
            "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
            "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
            "TLS_DHE_RSA_WITH_AES_256_GCM_SHA384",
            "TLS_DHE_RSA_WITH_AES_128_GCM_SHA256"
    );
    private static final List<String> FIPS_APPROVED_PROTOCOLS = List.of("TLSv1.3", "TLSv1.2");

    public static void validateFipsConnection(SSLSession session) throws SQLException {
        if (session == null) {
            throw new SQLException("FIPS VALIDATION FAILED: no SSL session");
        }

        String protocol = session.getProtocol();
        if (!FIPS_APPROVED_PROTOCOLS.contains(protocol)) {
            throw new SQLException("FIPS VALIDATION FAILED: wrong protocol: " + protocol +
                                   ". Permitted protocols: " + FIPS_APPROVED_PROTOCOLS);
        }

        String cipherSuite = session.getCipherSuite();
        if (!FIPS_APPROVED_CIPHERS.contains(cipherSuite)) {
            throw new SQLException("FIPS VALIDATION FAILED: wrong cipher: " + cipherSuite +
                                   ". FIPS-approved ciphers are: " + FIPS_APPROVED_CIPHERS);
        }

        System.out.println("- FIPS VALIDATION PASSED: " + protocol + " / " + cipherSuite);
    }

    public static SSLContext createStrictFipsSSLContext() throws Exception {
        if (!FipsManager.isFipsModeEnabled()) {
            throw new SQLException("FIPS mode is not enabled!");
        }

        String fipsProvider = FipsManager.getCurrentFipsProvider();

        if ("BCFIPS".equals(fipsProvider)) {
            return createBCFipsSSLContext();
        } else if (fipsProvider != null) {
            SSLContext sslContext = SSLContext.getInstance("TLS", fipsProvider);
            sslContext.init(null, new TrustManager[]{new StrictFipsTrustManager()}, null);
            return sslContext;
        } else
            throw new SQLException("No FIPS providers found!");
    }

    private static SSLContext createBCFipsSSLContext() throws Exception {
        SSLContext sslContext;

        sslContext = SSLContext.getInstance("TLS");
        sslContext.init(null, new TrustManager[]{new StrictFipsTrustManager()},
                org.teiid.core.crypto.FipsUtil.getSecureRandom());

        System.out.println("- Created SSL-context for  BCFIPS over default TLS");

        return wrapWithFipsValidation(sslContext);
    }

    private static SSLContext wrapWithFipsValidation(SSLContext originalContext) {
        return new SSLContext(new FipsValidatingSSLContextSpi(originalContext),
                originalContext.getProvider(),
                originalContext.getProtocol()) {
        };
    }

    private static class FipsValidatingSSLContextSpi extends SSLContextSpi {
        private final SSLContext delegate;

        public FipsValidatingSSLContextSpi(SSLContext delegate) {
            this.delegate = delegate;
        }

        @Override
        protected SSLSocketFactory engineGetSocketFactory() {
            return new FipsValidatingSSLSocketFactory(delegate.getSocketFactory());
        }

        @Override
        protected SSLServerSocketFactory engineGetServerSocketFactory() {
            return delegate.getServerSocketFactory();
        }

        @Override
        protected SSLEngine engineCreateSSLEngine() {
            return delegate.createSSLEngine();
        }

        @Override
        protected SSLEngine engineCreateSSLEngine(String host, int port) {
            return delegate.createSSLEngine(host, port);
        }

        @Override
        protected SSLSessionContext engineGetClientSessionContext() {
            return delegate.getClientSessionContext();
        }

        @Override
        protected SSLSessionContext engineGetServerSessionContext() {
            return delegate.getServerSessionContext();
        }

        @Override
        protected void engineInit(KeyManager[] km, TrustManager[] tm, SecureRandom random) {
            // Delegation handled in constructor
        }
    }

    private static class FipsValidatingSSLSocketFactory extends SSLSocketFactory {
        private final SSLSocketFactory delegate;

        public FipsValidatingSSLSocketFactory(SSLSocketFactory delegate) {
            this.delegate = delegate;
        }

        @Override
        public Socket createSocket() throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket());
        }

        @Override
        public Socket createSocket(String host, int port) throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket(host, port));
        }

        @Override
        public Socket createSocket(String host, int port, InetAddress localHost, int localPort)
                throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket(host, port, localHost, localPort));
        }

        @Override
        public Socket createSocket(InetAddress host, int port) throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket(host, port));
        }

        @Override
        public Socket createSocket(InetAddress address, int port, InetAddress localAddress, int localPort)
                throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket(address, port, localAddress, localPort));
        }

        @Override
        public Socket createSocket(Socket s, String host, int port, boolean autoClose) throws IOException {
            return wrapSocket((SSLSocket) delegate.createSocket(s, host, port, autoClose));
        }

        @Override
        public String[] getDefaultCipherSuites() {
            return getFipsCompliantCiphers(delegate.getDefaultCipherSuites());
        }

        @Override
        public String[] getSupportedCipherSuites() {
            return getFipsCompliantCiphers(delegate.getSupportedCipherSuites());
        }

        private SSLSocket wrapSocket(SSLSocket socket) {
            String[] fipsProtocols = new String[] { "TLSv1.3", "TLSv1.2" };
            socket.setEnabledProtocols(
                    Arrays.stream(fipsProtocols)
                          .filter(p -> Arrays.asList(socket.getSupportedProtocols()).contains(p))
                          .toArray(String[]::new)
            );

            System.out.println("=== CIPHER DEBUG ===");
            System.out.println("Supported: " + Arrays.toString(socket.getSupportedCipherSuites()));
            
            String[] fipsCiphers = getFipsCompliantCiphers(socket.getSupportedCipherSuites());
            System.out.println("FIPS ciphers: " + Arrays.toString(fipsCiphers));
            
            if (fipsCiphers.length > 0) {
                socket.setEnabledCipherSuites(fipsCiphers);
                System.out.println("Enabled: " + Arrays.toString(socket.getEnabledCipherSuites()));
            } else {
                System.out.println("There are no FIPS-approved cipher suites!");
            }

            socket.addHandshakeCompletedListener(event -> {
                try {
                    validateFipsConnection(event.getSession());
                } catch (SQLException e) {
                    System.err.println("FIPS validation failed: " + e.getMessage());
                    try {
                        socket.close();
                    } catch (IOException ex) {
                        // ignore
                    }
                    throw new RuntimeException("FIPS validation failed", e);
                }
            });

            return socket;
        }

        private String[] getFipsCompliantCiphers(String[] supportedCiphers) {
            return java.util.Arrays.stream(supportedCiphers)
                    .filter(FIPS_APPROVED_CIPHERS::contains)
                    .toArray(String[]::new);
        }
    }

    public static class StrictFipsTrustManager implements X509TrustManager {
        private final X509TrustManager defaultTrustManager;

        public StrictFipsTrustManager() throws Exception {
            TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            tmf.init((KeyStore) null);

            for (TrustManager tm : tmf.getTrustManagers()) {
                if (tm instanceof X509TrustManager) {
                    this.defaultTrustManager = (X509TrustManager) tm;
                    return;
                }
            }
            throw new Exception("X509TrustManager not found");
        }

        @Override
        public void checkClientTrusted(X509Certificate[] chain, String authType) throws java.security.cert.CertificateException {
            validateFipsCertificate(authType);
            defaultTrustManager.checkClientTrusted(chain, authType);
        }

        @Override
        public void checkServerTrusted(X509Certificate[] chain, String authType) throws java.security.cert.CertificateException {
            validateFipsCertificate(authType);
            defaultTrustManager.checkServerTrusted(chain, authType);
        }

        @Override
        public X509Certificate[] getAcceptedIssuers() {
            return defaultTrustManager.getAcceptedIssuers();
        }

        private void validateFipsCertificate(String authType) throws java.security.cert.CertificateException {
            var types = List.of("RSA", "ECDSA",
                    // UNKNOWN is possible for TLSv1.3
                    "UNKNOWN");
            if (!types.contains(authType)) {
                throw new java.security.cert.CertificateException(
                        "FIPS VALIDATION FAILED: Invalid authentication algorithm: " + authType);
            }
        }
    }
}
