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

import com.datavirtuality.dv.jdbc.Driver;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.teiid.adminapi.Model;
import org.teiid.adminapi.impl.ModelMetaData;
import org.teiid.common.buffer.BufferManagerFactory;
import org.teiid.core.crypto.FipsUtil;
import org.teiid.jdbc.FakeServer;
import org.teiid.jdbc.FipsConnectionValidator;
import org.teiid.jdbc.FipsManager;
import org.teiid.net.TeiidURL;
import org.teiid.net.socket.SocketServerConnectionFactory;
import org.teiid.runtime.EmbeddedConfiguration;

import javax.net.ssl.SSLContext;
import java.net.InetSocketAddress;
import java.security.Security;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Properties;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

/**
 * Test class for JDBC connections with FIPS mode enabled.
 *
 * @author Roman Pedchenko
 * @date 2025.09.12
 */
public class TestJDBCFipsMode {
    public record Store(String name, String password, String type, String keyPassword) {}

    private static final Store BCFKS = new Store("src/test/resources/jdbc/server.bcfks", "Datavirtuality123!", "bcfks", null);
    private static final Store SHORT_KEY = new Store("src/test/resources/jdbc/server_short_key.bcfks", "Datavirtuality123!", "bcfks", null);
    private static final Store BCFKS_KP = new Store("src/test/resources/jdbc/server_kp.bcfks", "Datavirtuality123!", "bcfks", "keyPassword1!");
    private TestServerData serverData;

    @BeforeAll
    public static void oneTimeSetup() throws Exception {
        Security.setProperty("jdk.tls.disabledAlgorithms", "TLSv1, TLSv1.1, SSLv3, RC4, DES, MD5withRSA, 3DES_EDE_CBC, DESede, DH keySize < 1024, EC keySize < 224, NULL");

        // SSL debug logging
//        java.util.logging.ConsoleHandler handler = new java.util.logging.ConsoleHandler();
//        handler.setLevel(java.util.logging.Level.ALL);
//        java.util.logging.Logger.getLogger("").addHandler(handler);
//        System.setProperty("javax.net.debug", "ssl,handshake,record,plaintext,trustmanager,keymanager,session");
    }

    @AfterEach
    public void oneTimeTearDown() throws Exception {
        if (serverData != null) {
            if (serverData.jdbcTransport != null) {
                serverData.jdbcTransport.stop();
            }
            if (serverData.server != null) {
                serverData.server.stop();
            }
        }
        FipsManager.removeFipsProvider();
        SocketServerConnectionFactory.testClearInstance();
    }

    @Test
    public void testStrictFipsConnection() throws Exception {
        setTrustStore(BCFKS);
        serverData = setuoJdbcServer(BCFKS);
        testFipsClient(FipsConnectionValidator.createStrictFipsSSLContext());
    }

    @Test
    public void testShortKeyFipsConnection() throws Exception {
        setTrustStore(SHORT_KEY);
        serverData = setuoJdbcServer(SHORT_KEY);
        assertThatThrownBy(() -> testFipsClient(FipsConnectionValidator.createStrictFipsSSLContext()));
    }

    @Test
    public void testStrictKeyPasswordFipsConnection() throws Exception {
        setTrustStore(BCFKS_KP);
        serverData = setuoJdbcServer(BCFKS_KP);
        testFipsClient(FipsConnectionValidator.createStrictFipsSSLContext());
    }

    private void testFipsClient(SSLContext sslContext) throws Exception {
        SSLContext.setDefault(sslContext);

        Properties connProps = new Properties();
        connProps.setProperty("user", "testuser");
        connProps.setProperty("password", "testpassword");
        connProps.setProperty(TeiidURL.CONNECTION.FIPS, "enabled");
        connProps.setProperty(TeiidURL.CONNECTION.FIPS_PROVIDER, "BCFIPS");
        connProps.setProperty("ssl", "true");

        try (Connection fipsConn = Driver.getInstance().connect(
                "jdbc:datavirtuality:testvdb@mms://" + serverData.addr.getHostName() + ":" + serverData.jdbcTransport.getPort(),
                connProps)) {

            assertThat(FipsManager.isFipsModeEnabled()).isEqualTo(true);
            assertThat(FipsManager.getCurrentFipsProvider()).isEqualTo(FipsUtil.BCFIPS_PROVIDER);

            try (Statement stmt = fipsConn.createStatement()) {
                try (ResultSet rs = stmt.executeQuery("SELECT col1 FROM TestView")) {
                    assertThat(rs.next()).isTrue();
                    assertThat(rs.getInt(1)).isEqualTo(1);
                }
            }
        }
    }

    private static void setTrustStore(Store store) {
        System.setProperty("javax.net.ssl.trustStore", store.name());
        System.setProperty("javax.net.ssl.trustStorePassword", store.password());
        System.setProperty("javax.net.ssl.trustStoreType", store.type());
    }

    public static TestServerData setuoJdbcServer(Store store) throws Exception {
        TestServerData data = new TestServerData();
        final SSLConfiguration sslConfig = getServerSslConfiguration(store);

        SocketConfiguration config = new SocketConfiguration();
        config.setSSLConfiguration(sslConfig);

        data.addr = new InetSocketAddress(0);
        config.setBindAddress(data.addr.getHostName());
        config.setPortNumber(0);

        EmbeddedConfiguration dqpConfig = new EmbeddedConfiguration();
        dqpConfig.setMaxActivePlans(2);
        data.server = new FakeServer(false);
        data.server.start(dqpConfig, false);
        ModelMetaData testModel = new ModelMetaData();
        testModel.setName("test");
        testModel.setModelType(Model.Type.VIRTUAL);
        testModel.setSchemaSourceType("ddl");
        testModel.setSchemaText("CREATE VIEW TestView AS SELECT 1 AS col1");
        data.server.deployVDB("testvdb", testModel);

        data.jdbcTransport = new SocketListener(data.addr, config, data.server.getClientServiceRegistry(),
                BufferManagerFactory.getStandaloneBufferManager());
        data.jdbcTransport.setMaxMessageSize(100000);
        data.jdbcTransport.setMaxLobSize(10000);
        return data;
    }

    private static SSLConfiguration getServerSslConfiguration(Store store) throws SQLException {
        SSLConfiguration sslConfig = new SSLConfiguration();
        sslConfig.setMode(SSLConfiguration.ENABLED);

        sslConfig.setKeystoreFilename(store.name());
        sslConfig.setKeystorePassword(store.password());
        sslConfig.setKeystoreType(store.type());
        sslConfig.setKeystoreKeyPassword(store.keyPassword());

        Properties fipsProps = new Properties();
        fipsProps.setProperty(TeiidURL.CONNECTION.FIPS, "enabled");
        fipsProps.setProperty(TeiidURL.CONNECTION.FIPS_PROVIDER, "BCFIPS");
        FipsManager.initializeFipsMode(fipsProps);

        String fipsCiphers = "TLS_AES_256_GCM_SHA384";
        sslConfig.setEnabledCipherSuites(fipsCiphers);

        return sslConfig;
    }

    public static class TestServerData {
        InetSocketAddress addr;
        SocketListener jdbcTransport;
        FakeServer server;
    }
}
