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

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import javax.crypto.KeyGenerator;
import javax.crypto.spec.SecretKeySpec;
import java.io.Serializable;
import java.security.Key;
import java.security.Provider;
import java.security.Security;

import static org.assertj.core.api.Assertions.*;

/**
 * @author Roman Pedchenko
 * @date 2025.08.07
 */
public class TestBasicCryptorFips {
    private static final String TEST_ALGORITHM = "TESTAES";
    private static final String TEST_PROVIDER = "MockFipsProvider";
    private Key testKey;

    @BeforeEach
    public void setUp() throws Exception {
        testKey = new SecretKeySpec(new byte[16], TEST_ALGORITHM);

        Provider mockProvider = new MockFipsProvider();
        Security.addProvider(mockProvider);

        FipsUtil.setFipsMode(false, null);
    }

    @AfterEach
    public void tearDown() {
        FipsUtil.setFipsMode(false, null);
        Security.removeProvider(TEST_PROVIDER);
    }

    @Test
    public void testInitializationWithValidProvider() throws Exception {
        BasicCryptor cryptor = new BasicCryptor(testKey, testKey, TEST_ALGORITHM);
        assertThat(cryptor).isNotNull();
    }

    @Test
    public void testInitializationWithInvalidProvider() {
        try {
            FipsUtil.validateProvider("NonExistentProvider");
            fail("Expected CryptoException was not thrown");
        } catch (CryptoException e) {
            assertThat(e.getMessage()).contains("Security provider 'NonExistentProvider' not found");
        }
    }

    @Test
    public void testEncryptDecrypt() throws Exception {
        BasicCryptor cryptor = new BasicCryptor(testKey, testKey, TEST_ALGORITHM);

        String original = "Test data for encryption";
        byte[] encrypted = cryptor.encrypt(original.getBytes());
        byte[] decrypted = cryptor.decrypt(encrypted);

        assertThat(new String(decrypted)).isEqualTo(original);
    }

    @Test
    public void testDecryptFailsInFipsMode() throws Exception {
        MockFailingCryptor failingCryptor = new MockFailingCryptor(testKey, testKey, TEST_ALGORITHM);

        assertThatThrownBy(() ->
                failingCryptor.decrypt("test data".getBytes())
        ).isInstanceOf(CryptoException.class);
    }

    @Test
    public void testEncryptFailsInFipsMode() throws Exception {
        MockFailingCryptor failingCryptor = new MockFailingCryptor(testKey, testKey, TEST_ALGORITHM);

        assertThatThrownBy(() ->
                failingCryptor.encrypt("test data".getBytes())
        ).isInstanceOf(CryptoException.class);
    }

    @Test
    public void testSealUnsealObject() throws Exception {
        BasicCryptor cryptor = new BasicCryptor(testKey, testKey, TEST_ALGORITHM);

        TestSerializable original = new TestSerializable("Test object");
        Object sealed = cryptor.sealObject(original);
        Object unsealed = cryptor.unsealObject(sealed);

        assertThat(unsealed).isInstanceOf(TestSerializable.class);
        assertThat(((TestSerializable) unsealed).getValue()).isEqualTo(original.getValue());
    }

    @Test
    public void testUnsealObjectFailsInFipsMode() throws Exception {
        BasicCryptor cryptor = new BasicCryptor(testKey, testKey, TEST_ALGORITHM);

        TestSerializable original = new TestSerializable("Test object");
        Object sealed = cryptor.sealObject(original);

        // Create a mock cryptor that will throw an exception during unsealing
        MockFailingCryptor failingCryptor = new MockFailingCryptor(testKey, testKey, TEST_ALGORITHM);

        // This will throw a CryptoException regardless of FIPS mode because our mock is designed to fail
        assertThatThrownBy(() ->
                failingCryptor.unsealObject(sealed)
        ).isInstanceOf(CryptoException.class);
    }

    private static class MockFipsProvider extends Provider {
        private static final long serialVersionUID = 0x1E510000L;

        public MockFipsProvider() {
            super(TEST_PROVIDER, "1.0", "Mock FIPS Provider for Testing");
            put("Cipher.TESTAES", MockCipher.class.getName());
            put("KeyGenerator.TESTAES", MockKeyGenerator.class.getName());
            put("SecureRandom.SHA1PRNG", MockSecureRandom.class.getName());
        }
    }

    public static class MockCipher extends javax.crypto.CipherSpi {
        @Override
        protected byte[] engineDoFinal(byte[] input, int inputOffset, int inputLen) {
            byte[] result = new byte[inputLen];
            System.arraycopy(input, inputOffset, result, 0, inputLen);
            return result;
        }

        @Override
        protected int engineDoFinal(byte[] input, int inputOffset, int inputLen, byte[] output, int outputOffset) {
            System.arraycopy(input, inputOffset, output, outputOffset, inputLen);
            return inputLen;
        }

        @Override
        protected int engineGetBlockSize() {
            return 16;
        }

        @Override
        protected byte[] engineGetIV() {
            return new byte[16];
        }

        @Override
        protected int engineGetOutputSize(int inputLen) {
            return inputLen;
        }

        @Override
        protected java.security.AlgorithmParameters engineGetParameters() {
            return null;
        }

        @Override
        protected void engineInit(int opmode, Key key, java.security.spec.AlgorithmParameterSpec params, java.security.SecureRandom random) {
        }

        @Override
        protected void engineInit(int opmode, Key key, java.security.AlgorithmParameters params, java.security.SecureRandom random) {
        }

        @Override
        protected void engineInit(int opmode, Key key, java.security.SecureRandom random) {
        }

        @Override
        protected void engineSetMode(String mode) {
        }

        @Override
        protected void engineSetPadding(String padding) {
        }

        @Override
        protected byte[] engineUpdate(byte[] input, int inputOffset, int inputLen) {
            byte[] result = new byte[inputLen];
            System.arraycopy(input, inputOffset, result, 0, inputLen);
            return result;
        }

        @Override
        protected int engineUpdate(byte[] input, int inputOffset, int inputLen, byte[] output, int outputOffset) {
            System.arraycopy(input, inputOffset, output, outputOffset, inputLen);
            return inputLen;
        }
    }

    public static class MockKeyGenerator extends KeyGenerator {
        public MockKeyGenerator() {
            super(new MockKeyGeneratorSpi(), new MockFipsProvider(), TEST_ALGORITHM);
        }
    }

    public static class MockKeyGeneratorSpi extends javax.crypto.KeyGeneratorSpi {
        @Override
        protected SecretKeySpec engineGenerateKey() {
            return new SecretKeySpec(new byte[16], TEST_ALGORITHM);
        }

        @Override
        protected void engineInit(int keysize, java.security.SecureRandom random) {
        }

        @Override
        protected void engineInit(java.security.spec.AlgorithmParameterSpec params, java.security.SecureRandom random) {
        }

        @Override
        protected void engineInit(java.security.SecureRandom random) {
        }
    }

    public static class MockSecureRandom extends java.security.SecureRandom {
        private static final long serialVersionUID = 0x1E510001L;

        @Override
        public void nextBytes(byte[] bytes) {
            for (int i = 0; i < bytes.length; i++) {
                bytes[i] = (byte) i;
            }
        }
    }

    private static class MockFailingCryptor extends BasicCryptor {
        public MockFailingCryptor(Key encryptKey, Key decryptKey, String algorithm) throws CryptoException {
            super(encryptKey, decryptKey, algorithm);
        }

        @Override
        public synchronized byte[] decrypt(byte[] ciphertext) throws CryptoException {
            throw new CryptoException("Simulated decryption failure in FIPS mode");
        }

        @Override
        public synchronized byte[] encrypt(byte[] buffer, int offset, int length) throws CryptoException {
            throw new CryptoException("Simulated encryption failure in FIPS mode");
        }

        @Override
        public synchronized Object unsealObject(Object object) throws CryptoException {
            throw new CryptoException("Simulated unsealing failure in FIPS mode");
        }
    }

    private static class TestSerializable implements Serializable {
        private static final long serialVersionUID = 0x1E510002L;
        private final String value;

        public TestSerializable(String value) {
            this.value = value;
        }

        public String getValue() {
            return value;
        }
    }
}
