nomyo4J/src/test/java/ai/nomyo/SecureCompletionClientTest.java

package ai.nomyo;

import ai.nomyo.errors.SecurityError;
import ai.nomyo.util.Pass2Key;
import org.junit.jupiter.api.*;
import org.junit.jupiter.api.io.TempDir;

import java.io.File;
import java.math.BigInteger;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.*;
import java.security.spec.RSAKeyGenParameterSpec;
import java.security.spec.RSAPrivateCrtKeySpec;

import static org.junit.jupiter.api.Assertions.*;

@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
class SecureCompletionClientTest {

    private static final String TEST_PASSWORD = "test-password-123";
    private static final String PLAINTEXT_PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC7o4qne60TBb\n-----END PRIVATE KEY-----";

    private SecureCompletionClient client;

    @BeforeEach
    void setUp() {
        client = new SecureCompletionClient();
    }

    @AfterEach
    void tearDown() {
        client = null;
    }

    // ── Key Generation Tests ──────────────────────────────────────────

    @Test
    @Order(1)
    @DisplayName("generateKeys should create 4096-bit RSA key pair")
    void generateKeys_shouldCreateValidKeyPair() {
        client.generateKeys(false);

        PrivateKey privateKey = client.getPrivateKey();
        String publicPemKey = client.getPublicPemKey();

        assertNotNull(privateKey, "Private key should not be null");
        assertNotNull(publicPemKey, "Public PEM key should not be null");
        assertEquals("RSA", privateKey.getAlgorithm(), "Key algorithm should be RSA");
        assertTrue(publicPemKey.contains("BEGIN PUBLIC KEY"), "Public key should be valid PEM");
    }

    @Test
    @Order(2)
    @DisplayName("generateKeys should produce keys with correct bit size")
    void generateKeys_shouldProduceCorrectKeySize() throws Exception {
        client.generateKeys(false);

        PrivateKey privateKey = client.getPrivateKey();
        java.security.KeyFactory kf = java.security.KeyFactory.getInstance("RSA");
        java.security.spec.RSAPrivateCrtKeySpec spec = kf.getKeySpec(privateKey, java.security.spec.RSAPrivateCrtKeySpec.class);
        int keySizeBits = spec.getModulus().bitLength();

        assertEquals(Constants.RSA_KEY_SIZE, keySizeBits,
                "RSA key should be " + Constants.RSA_KEY_SIZE + " bits");
    }

    @Test
    @Order(3)
    @DisplayName("generateKeys should create unique keys on each call")
    void generateKeys_shouldProduceUniqueKeys() {
        client.generateKeys(false);
        PrivateKey firstKey = client.getPrivateKey();

        SecureCompletionClient client2 = new SecureCompletionClient();
        client2.generateKeys(false);
        PrivateKey secondKey = client2.getPrivateKey();

        assertNotEquals(java.util.Arrays.hashCode(firstKey.getEncoded()), java.util.Arrays.hashCode(secondKey.getEncoded()),
                "Different keys should have different encoded content");
    }

    // ── Key Generation with File Save Tests ───────────────────────────

    @Test
    @Order(4)
    @DisplayName("generateKeys with saveToFile=true should create key files")
    void generateKeys_withSaveToFile_shouldCreateKeyFiles(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();

        client.generateKeys(true, keyDir.getAbsolutePath(), null);

        File privateKeyFile = new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
        File publicKeyFile = new File(keyDir, Constants.DEFAULT_PUBLIC_KEY_FILE);

        assertTrue(privateKeyFile.exists(), "Private key file should be created");
        assertTrue(publicKeyFile.exists(), "Public key file should be created");
        assertTrue(privateKeyFile.length() > 0, "Private key file should not be empty");
        assertTrue(publicKeyFile.length() > 0, "Public key file should not be empty");
    }

    @Test
    @Order(5)
    @DisplayName("generateKeys with password should encrypt private key file")
    void generateKeys_withPassword_shouldEncryptPrivateKey(@TempDir Path tempDir) throws Exception {
        File keyDir = tempDir.toFile();

        client.generateKeys(true, keyDir.getAbsolutePath(), TEST_PASSWORD);

        File privateKeyFile = new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
        String content = Files.readString(privateKeyFile.toPath());

        assertFalse(content.contains("BEGIN PRIVATE KEY"),
                "Encrypted private key should not contain PEM header");
        assertNotEquals(PLAINTEXT_PRIVATE_KEY, content,
                "Encrypted key should differ from plaintext");
    }

    @Test
    @Order(6)
    @DisplayName("generateKeys should not overwrite existing key files")
    void generateKeys_shouldNotOverwriteExistingKeys(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();

        client.generateKeys(true, keyDir.getAbsolutePath(), null);
        File privateKeyFile = new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
        long firstSize = privateKeyFile.length();

        client.generateKeys(true, keyDir.getAbsolutePath(), null);
        long secondSize = privateKeyFile.length();

        assertEquals(firstSize, secondSize,
                "Existing key files should not be overwritten");
    }

    // ── Key Loading Tests ─────────────────────────────────────────────

    @Test
    @Order(7)
    @DisplayName("loadKeys should load plaintext private key from file")
    void loadKeys_plaintext_shouldLoadPrivateKey(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();
        client.generateKeys(true, keyDir.getAbsolutePath(), null);

        PrivateKey originalKey = client.getPrivateKey();

        SecureCompletionClient loadClient = new SecureCompletionClient();
        loadClient.loadKeys(
                new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE).getAbsolutePath(),
                null
        );

        PrivateKey loadedKey = loadClient.getPrivateKey();
        assertNotNull(loadedKey, "Loaded private key should not be null");
        assertEquals(originalKey.getEncoded().length, loadedKey.getEncoded().length,
                "Loaded key should have same size as original");
    }

    @Test
    @Order(8)
    @DisplayName("loadKeys should load encrypted private key with correct password")
    void loadKeys_encrypted_correctPassword_shouldLoadPrivateKey(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();
        client.generateKeys(true, keyDir.getAbsolutePath(), TEST_PASSWORD);

        PrivateKey originalKey = client.getPrivateKey();

        SecureCompletionClient loadClient = new SecureCompletionClient();
        loadClient.loadKeys(
                new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE).getAbsolutePath(),
                null,
                TEST_PASSWORD
        );

        PrivateKey loadedKey = loadClient.getPrivateKey();
        assertNotNull(loadedKey, "Loaded private key should not be null");
        assertEquals(originalKey.getEncoded().length, loadedKey.getEncoded().length,
                "Loaded key should have same size as original");
    }

    @Test
    @Order(9)
    @DisplayName("loadKeys should handle wrong password gracefully")
    void loadKeys_encrypted_wrongPassword_shouldHandleGracefully(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();
        client.generateKeys(true, keyDir.getAbsolutePath(), TEST_PASSWORD);

        SecureCompletionClient loadClient = new SecureCompletionClient();

        assertDoesNotThrow(() ->
                loadClient.loadKeys(
                        new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE).getAbsolutePath(),
                        null,
                        "wrong-password"
                ),
                "Wrong password should not throw exception"
        );
    }

    @Test
    @Order(10)
    @DisplayName("loadKeys should throw exception for non-existent file")
    void loadKeys_nonExistentFile_shouldThrowException() {
        SecureCompletionClient loadClient = new SecureCompletionClient();

        RuntimeException error = assertThrows(RuntimeException.class, () ->
                loadClient.loadKeys("/non/existent/path/private_key.pem", null, null));

        assertTrue(error.getMessage().contains("not found"),
                "Error message should mention file not found");
    }

    // ── Key Validation Tests ──────────────────────────────────────────

    @Test
    @Order(11)
    @DisplayName("validateRsaKey should accept valid 4096-bit key")
    void validateRsaKey_validKey_shouldPass() {
        client.generateKeys(false);
        PrivateKey key = client.getPrivateKey();

        assertDoesNotThrow(() -> client.validateRsaKey(key),
                "Valid 4096-bit key should pass validation");
    }

    @Test
    @Order(12)
    @DisplayName("validateRsaKey should reject null key")
    void validateRsaKey_nullKey_shouldThrowSecurityError() {
        SecurityError error = assertThrows(SecurityError.class, () ->
                client.validateRsaKey(null));

        assertTrue(error.getMessage().contains("null"),
                "Error message should mention null key");
    }

    @Test
    @Order(13)
    @DisplayName("validateRsaKey should reject keys below minimum size")
    void validateRsaKey_tooSmallKey_shouldThrowSecurityError() throws Exception {
        KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
        generator.initialize(1024);
        KeyPair pair = generator.generateKeyPair();
        PrivateKey smallKey = pair.getPrivate();

        SecurityError error = assertThrows(SecurityError.class, () ->
                client.validateRsaKey(smallKey));

        assertTrue(error.getMessage().contains("below minimum"),
                "Error message should mention minimum size requirement");
    }

    @Test
    @Order(14)
    @DisplayName("validateRsaKey should accept minimum size key (2048 bits)")
    void validateRsaKey_minimumSizeKey_shouldPass() throws Exception {
        KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
        generator.initialize(new RSAKeyGenParameterSpec(2048, BigInteger.valueOf(65537)));
        KeyPair pair = generator.generateKeyPair();
        PrivateKey minKey = pair.getPrivate();

        assertDoesNotThrow(() -> client.validateRsaKey(minKey),
                "2048-bit key should pass validation (minimum)");
    }

    // ── Key Roundtrip Tests ───────────────────────────────────────────

    @Test
    @Order(15)
    @DisplayName("Full roundtrip: generate, save, load should produce same key")
    void roundtrip_generateSaveLoad_shouldProduceSameKey(@TempDir Path tempDir) {
        File keyDir = tempDir.toFile();

        // Generate and save
        client.generateKeys(true, keyDir.getAbsolutePath(), TEST_PASSWORD);
        PrivateKey originalKey = client.getPrivateKey();

        // Load into new client
        SecureCompletionClient loadClient = new SecureCompletionClient();
        loadClient.loadKeys(
                new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE).getAbsolutePath(),
                null,
                TEST_PASSWORD
        );
        PrivateKey loadedKey = loadClient.getPrivateKey();

        // Verify keys match
        assertNotNull(loadedKey, "Loaded key should not be null");
        assertArrayEquals(originalKey.getEncoded(), loadedKey.getEncoded(),
                "Loaded key should match original key bytes");
    }

    @Test
    @Order(16)
    @DisplayName("Multiple generate/load cycles should work correctly")
    void multipleCycles_shouldWorkCorrectly(@TempDir Path tempDir) throws Exception {
        File keyDir = tempDir.toFile();

        for (int i = 0; i < 3; i++) {
            SecureCompletionClient cycleClient = new SecureCompletionClient();
            cycleClient.generateKeys(true, keyDir.getAbsolutePath(), TEST_PASSWORD);
            PrivateKey generatedKey = cycleClient.getPrivateKey();

            SecureCompletionClient loadClient = new SecureCompletionClient();
            loadClient.loadKeys(
                    new File(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE).getAbsolutePath(),
                    null,
                    TEST_PASSWORD
            );
            PrivateKey loadedKey = loadClient.getPrivateKey();

            assertNotNull(loadedKey, "Cycle " + i + ": loaded key should not be null");
            KeyFactory kf = KeyFactory.getInstance("RSA");
            int generatedBits = kf.getKeySpec(generatedKey, RSAPrivateCrtKeySpec.class).getModulus().bitLength();
            int loadedBits = kf.getKeySpec(loadedKey, RSAPrivateCrtKeySpec.class).getModulus().bitLength();
            assertEquals(generatedBits, loadedBits,
                    "Cycle " + i + ": loaded key should have same modulus bit length as generated key");
        }
    }

    // ── Utility Method Tests ──────────────────────────────────────────

    @Test
    @Order(17)
    @DisplayName("urlEncodePublicKey should properly encode PEM keys")
    void urlEncodePublicKey_shouldEncodeCorrectly() {
        String pemKey = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A\n-----END PUBLIC KEY-----";

        String encoded = client.urlEncodePublicKey(pemKey);

        assertNotNull(encoded, "Encoded key should not be null");
        assertTrue(encoded.contains("%0A"), "Encoded key should contain URL-encoded newlines");
        assertTrue(encoded.contains("BEGIN+PUBLIC+KEY"), "Encoded key should contain encoded header");
    }

    @Test
    @Order(18)
    @DisplayName("urlEncodePublicKey should handle empty string")
    void urlEncodePublicKey_emptyString_shouldReturnEmpty() {
        String encoded = client.urlEncodePublicKey("");
        assertEquals("", encoded, "Empty string should encode to empty string");
    }

    // ── Pass2Key Encryption/Decryption Tests ──────────────────────────

    @Test
    @Order(19)
    @DisplayName("Pass2Key encrypt/decrypt should preserve plaintext")
    void pass2Key_encryptDecrypt_shouldPreservePlaintext() throws Exception {
        String plaintext = "Test plaintext content for encryption";
        String password = "my-secret-password";

        String encrypted = Pass2Key.encrypt("AES/GCM/NoPadding", plaintext, password);

        assertNotNull(encrypted, "Encrypted text should not be null");
        assertNotEquals(plaintext, encrypted, "Encrypted text should differ from plaintext");

        String decrypted = Pass2Key.decrypt("AES/GCM/NoPadding", encrypted, password);

        assertEquals(plaintext, decrypted, "Decrypted text should match original plaintext");
    }

    @Test
    @Order(20)
    @DisplayName("Pass2Key should produce different ciphertext for same plaintext")
    void pass2Key_shouldProduceDifferentCiphertext() throws Exception {
        String plaintext = "Same plaintext";
        String password = "same-password";

        String encrypted1 = Pass2Key.encrypt("AES/GCM/NoPadding", plaintext, password);
        String encrypted2 = Pass2Key.encrypt("AES/GCM/NoPadding", plaintext, password);

        assertNotEquals(encrypted1, encrypted2,
                "Same plaintext with same password should produce different ciphertext (due to random salt/IV)");
    }

    @Test
    @Order(21)
    @DisplayName("Pass2Key decrypt should fail with wrong password")
    void pass2Key_wrongPassword_shouldFail() throws Exception {
        String plaintext = "Secret content";
        String correctPassword = "correct-password";
        String wrongPassword = "wrong-password";

        String encrypted = Pass2Key.encrypt("AES/GCM/NoPadding", plaintext, correctPassword);

        assertThrows(Exception.class, () ->
                Pass2Key.decrypt("AES/GCM/NoPadding", encrypted, wrongPassword),
                "Decryption with wrong password should throw exception"
        );
    }

    @Test
    @Order(22)
    @DisplayName("Pass2Key convertStringToPrivateKey should parse PEM correctly")
    void convertStringToPrivateKey_shouldParsePEM() throws Exception {
        SecureCompletionClient tempClient = new SecureCompletionClient();
        tempClient.generateKeys(false);
        PrivateKey originalKey = tempClient.getPrivateKey();

        String pem = ai.nomyo.util.PEMConverter.toPEM(originalKey.getEncoded(), true);
        PrivateKey parsedKey = Pass2Key.convertStringToPrivateKey(pem);

        assertNotNull(parsedKey, "Parsed private key should not be null");
        assertEquals("RSA", parsedKey.getAlgorithm(), "Parsed key should be RSA");
        assertArrayEquals(originalKey.getEncoded(), parsedKey.getEncoded(),
                "Parsed key should match original key bytes");
    }

    @Test
    @Order(23)
    @DisplayName("Pass2Key convertStringToPrivateKey should handle PEM with whitespace")
    void convertStringToPrivateKey_shouldHandleWhitespace() throws Exception {
        SecureCompletionClient tempClient = new SecureCompletionClient();
        tempClient.generateKeys(false);
        PrivateKey originalKey = tempClient.getPrivateKey();

        String pem = "-----BEGIN PRIVATE KEY-----\n  " +
                originalKey.getEncoded().length + "lines\n" +
                "-----END PRIVATE KEY-----";

        String formattedPem = ai.nomyo.util.PEMConverter.toPEM(originalKey.getEncoded(), true);
        String pemWithWhitespace = formattedPem.replace("\n", "\n  ");

        PrivateKey parsedKey = Pass2Key.convertStringToPrivateKey(pemWithWhitespace);

        assertNotNull(parsedKey, "Parsed private key should not be null even with whitespace");
        assertArrayEquals(originalKey.getEncoded(), parsedKey.getEncoded(),
                "Parsed key with whitespace should match original key bytes");
    }
}