nomyo4J/src/main/java/ai/nomyo/SecureCompletionClient.java

package ai.nomyo;

import ai.nomyo.errors.*;
import ai.nomyo.util.PEMConverter;
import ai.nomyo.util.Pass2Key;
import lombok.Getter;

import javax.crypto.*;
import java.io.FileWriter;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.nio.file.attribute.PosixFilePermission;
import java.nio.file.attribute.PosixFilePermissions;
import java.security.*;
import java.security.spec.RSAKeyGenParameterSpec;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Low-level client: key management, hybrid encryption, HTTP with retry, response decryption. Used by {@link SecureChatCompletion}.
 */
public class SecureCompletionClient {

    /**
     * NOMYO router base URL (trailing slash stripped).
     */
    @Getter
    private final String routerUrl;

    /**
     * Permit HTTP (non-HTTPS) URLs.
     */
    @Getter
    private final boolean allowHttp;

    /**
     * RSA key size in bits ({@link Constants#RSA_KEY_SIZE}).
     */
    @Getter
    private final int keySize;

    /**
     * Max retries on retryable errors.
     */
    @Getter
    private final int maxRetries;

    /**
     * Secure memory operations active.
     */
    @Getter
    private final boolean useSecureMemory;

    /**
     * Lock for double-checked key initialization.
     */
    private final ReentrantLock keyInitLock = new ReentrantLock();

    /**
     * RSA private key ({@code null} until loaded/generated).
     */
    @Getter
    private PrivateKey privateKey;

    /**
     * PEM-encoded public key ({@code null} until loaded/generated).
     */
    @Getter
    private String publicPemKey;

    /**
     * Keys initialized.
     */
    private volatile boolean keysInitialized = false;

    /**
     * Default settings: {@code https://api.nomyo.ai}, HTTPS-only, secure memory, 2 retries.
     */
    public SecureCompletionClient() {
        this(Constants.DEFAULT_BASE_URL, false, true, Constants.DEFAULT_MAX_RETRIES);
    }

    /**
     * @param routerUrl    NOMYO router base URL
     * @param allowHttp    permit HTTP URLs
     * @param secureMemory enable memory locking/zeroing
     * @param maxRetries   retries on retryable errors
     */
    public SecureCompletionClient(String routerUrl, boolean allowHttp, boolean secureMemory, int maxRetries) {
        this.routerUrl = routerUrl != null ? routerUrl.replaceAll("/+$", "") : Constants.DEFAULT_BASE_URL;
        this.allowHttp = allowHttp;
        this.useSecureMemory = secureMemory;
        this.keySize = Constants.RSA_KEY_SIZE;
        this.maxRetries = maxRetries;
    }

    private static String readFileContent(String filePath) throws IOException {
        return Files.readString(Path.of(filePath));
    }

    /**
     * Generates a 4096-bit RSA key pair (exponent 65537). Saves to disk if {@code saveToFile}.
     */
    public void generateKeys(boolean saveToFile, String keyDir, String password) {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
            generator.initialize(new RSAKeyGenParameterSpec(Constants.RSA_KEY_SIZE, BigInteger.valueOf(Constants.RSA_PUBLIC_EXPONENT)));

            KeyPair pair = generator.generateKeyPair();

            String privatePem = PEMConverter.toPEM(pair.getPrivate().getEncoded(), true);
            String publicPem = PEMConverter.toPEM(pair.getPublic().getEncoded(), false);

              if (saveToFile) {
                Path keyFolder = Path.of(keyDir);
                if (!Files.exists(keyFolder)) {
                    try {
                        Files.createDirectories(keyFolder);
                    } catch (IOException e) {
                        throw new IOException("Failed to create key directory: " + keyDir, e);
                    }
                }

                Path privateKeyPath = Path.of(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
                if (!Files.exists(privateKeyPath)) {
                    Set<PosixFilePermission> filePermissions = PosixFilePermissions.fromString(Constants.PRIVATE_KEY_FILE_MODE);
                    Files.createFile(privateKeyPath, PosixFilePermissions.asFileAttribute(filePermissions));

                    try (var writer = Files.newBufferedWriter(privateKeyPath)) {
                        if (password == null || password.isEmpty()) {
                            System.out.println("WARNING: Saving keys in plaintext!");
                        } else {
                            try {
                                privatePem = Pass2Key.encrypt("AES/GCM/NoPadding", privatePem, password);
                            } catch (NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException |
                                     InvalidKeyException e) {
                                throw new RuntimeException(e);
                            }
                        }
                        writer.write(privatePem);
                    }
                }

                Path publicKeyPath = Path.of(keyDir, Constants.DEFAULT_PUBLIC_KEY_FILE);
                if (!Files.exists(publicKeyPath)) {
                    Set<PosixFilePermission> publicPermissions = PosixFilePermissions.fromString(Constants.PUBLIC_KEY_FILE_MODE);
                    Files.createFile(publicKeyPath, PosixFilePermissions.asFileAttribute(publicPermissions));

                    try (var writer = Files.newBufferedWriter(publicKeyPath)) {
                        writer.write(publicPem);
                    }
                }
            }

            this.privateKey = pair.getPrivate();
            this.publicPemKey = publicPem;

        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("RSA not available: " + e.getMessage(), e);
        } catch (InvalidAlgorithmParameterException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException("Failed to save keys: " + e.getMessage(), e);
        }
    }

    /**
     * Generates a 4096-bit RSA key pair and saves to the default directory.
     */
    public void generateKeys(boolean saveToFile) {
        generateKeys(saveToFile, Constants.DEFAULT_KEY_DIR, null);
    }

    /**
     * Loads RSA private key from disk. If {@code publicPemKeyPath} is {@code null}, derives public key.
     * Validates key size >= {@link Constants#MIN_RSA_KEY_SIZE} bits.
     *
     * @param privateKeyPath   private key PEM path
     * @param publicPemKeyPath optional public key PEM path
     * @param password         optional password for encrypted private key
     */
    public void loadKeys(String privateKeyPath, String publicPemKeyPath, String password) {
        Path keyPath = Path.of(privateKeyPath);
        if (!Files.exists(keyPath)) {
            throw new RuntimeException("Private key file not found: " + privateKeyPath);
        }

        String keyContent;
        if (password != null && !password.isEmpty()) {
            try {
                keyContent = readFileContent(privateKeyPath);
            } catch (IOException e) {
                throw new RuntimeException("Failed to read private key file: " + e.getMessage(), e);
            }

            try {
                keyContent = Pass2Key.decrypt("AES/GCM/NoPadding", keyContent, password);
            } catch (NoSuchPaddingException | NoSuchAlgorithmException | BadPaddingException |
                     IllegalBlockSizeException | InvalidAlgorithmParameterException | InvalidKeyException e) {
                System.out.println("Wrong password!");
                return;
            }
        } else {
            try {
                keyContent = readFileContent(privateKeyPath);
            } catch (IOException e) {
                throw new RuntimeException("Failed to read private key file: " + e.getMessage(), e);
            }
        }

        try {
            this.privateKey = Pass2Key.convertStringToPrivateKey(keyContent);
        } catch (Exception e) {
            throw new RuntimeException("Failed to load private key: " + e.getMessage(), e);
        }
    }

    /**
     * Loads RSA private key from disk, deriving public key.
     */
    public void loadKeys(String privateKeyPath, String password) {
        loadKeys(privateKeyPath, null, password);
    }

    /**
     * GET {@code {routerUrl}/pki/public_key}. Returns server PEM public key.
     */
    public CompletableFuture<String> fetchServerPublicKey() {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Hybrid encryption: AES-256-GCM for payload, RSA-OAEP-SHA256 for AES key wrapping.
     *
     * @param payload OpenAI-compatible chat parameters
     * @return encrypted bytes (JSON package)
     * @throws SecurityError if encryption fails or keys not loaded
     */
    public CompletableFuture<byte[]> encryptPayload(Map<String, Object> payload) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Core hybrid encryption: AES-256-GCM encrypts {@code payloadBytes} with {@code aesKey}.
     */
    public CompletableFuture<byte[]> doEncrypt(byte[] payloadBytes, byte[] aesKey) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Decrypts server response.
     */
    public CompletableFuture<Map<String, Object>> decryptResponse(byte[] encryptedResponse, String payloadId) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * encrypt → POST {routerUrl}/v1/chat/secure_completion → retry → decrypt → return.
     * <p>Headers: Content-Type=octet-stream, X-Payload-ID, X-Public-Key, Authorization (Bearer), X-Security-Tier.
     * Retryable: 429, 500, 502, 503, 504 + network errors. Backoff: 2^(attempt-1)s.
     * <p>Status mapping: 200→return, 400→InvalidRequestError, 401→AuthenticationError, 403→ForbiddenError,
     * 404→APIError, 429→RateLimitError, 500→ServerError, 503→ServiceUnavailableError,
     * 502/504→APIError(retryable), network→APIConnectionError.
     *
     * @param payload      OpenAI-compatible chat parameters
     * @param payloadId    unique payload identifier
     * @param apiKey       optional API key
     * @param securityTier optional: "standard", "high", "maximum"
     * @return decrypted response map
     * @throws SecurityError           encryption/decryption failure
     * @throws APIConnectionError      network error
     * @throws InvalidRequestError     HTTP 400
     * @throws AuthenticationError     HTTP 401
     * @throws ForbiddenError          HTTP 403
     * @throws RateLimitError          HTTP 429
     * @throws ServerError             HTTP 500
     * @throws ServiceUnavailableError HTTP 503
     * @throws APIError                other errors
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey, String securityTier) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Without security tier.
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey) {
        return sendSecureRequest(payload, payloadId, apiKey, null);
    }

    /**
     * No API key or security tier.
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId) {
        return sendSecureRequest(payload, payloadId, null, null);
    }

    /**
     * Thread-safe key init via double-checked locking. Loads from disk if {@code keyDir} set, else generates.
     *
     * @param keyDir key directory or {@code null} for ephemeral
     */
    public void ensureKeys(String keyDir) {
        if (keysInitialized) return;
        keyInitLock.lock();
        try {
            if (keysInitialized) return;
            if (keyDir == null || keyDir.isEmpty()) {
                generateKeys(false);
            } else {
                generateKeys(true);
            }
            keysInitialized = true;
        } finally {
            keyInitLock.unlock();
        }
    }

    /**
     * Validates RSA key size >= {@link Constants#MIN_RSA_KEY_SIZE} bits.
     */
    public void validateRsaKey(PrivateKey key) throws SecurityError {
        if (key == null) {
            throw new SecurityError("RSA key is null");
        }
        int keySize = extractKeySize(key);

        if (keySize < Constants.MIN_RSA_KEY_SIZE) {
            throw new SecurityError("RSA key size " + keySize + " bits is below minimum " + Constants.MIN_RSA_KEY_SIZE + " bits");
        }
    }

    private int extractKeySize(PrivateKey key) {
        try {
            var kf = KeyFactory.getInstance("RSA");
            try {
                var crtSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateCrtKeySpec.class);
                return crtSpec.getModulus().bitLength();
            } catch (Exception ignored) {
                var privSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateKeySpec.class);
                return privSpec.getModulus().bitLength();
            }
        } catch (Exception ignored) {
            if (key.getEncoded() != null) {
                return key.getEncoded().length * 8;
            }
        }
        return 0;
    }

    /**
     * Maps HTTP status code to exception (200→null).
     */
    public Exception mapHttpStatus(int statusCode, String responseBody) {
        return switch (statusCode) {
            case 200 -> null;
            case 400 ->
                    new InvalidRequestError("Invalid request: " + (responseBody != null ? responseBody : "no body"));
            case 401 ->
                    new AuthenticationError("Authentication failed: " + (responseBody != null ? responseBody : "no body"));
            case 403 -> new ForbiddenError("Access forbidden: " + (responseBody != null ? responseBody : "no body"));
            case 404 -> new APIError("Not found: " + (responseBody != null ? responseBody : "no body"));
            case 429 -> new RateLimitError("Rate limit exceeded: " + (responseBody != null ? responseBody : "no body"));
            case 500 -> new ServerError("Internal server error: " + (responseBody != null ? responseBody : "no body"));
            case 503 ->
                    new ServiceUnavailableError("Service unavailable: " + (responseBody != null ? responseBody : "no body"));
            case 502, 504 -> new APIError("Gateway error: " + (responseBody != null ? responseBody : "no body"));
            default ->
                    new APIError("Unexpected status " + statusCode + ": " + (responseBody != null ? responseBody : "no body"));
        };
    }

    /**
     * URL-encodes PEM key for {@code X-Public-Key} header.
     */
    public String urlEncodePublicKey(String pemKey) {
        return java.net.URLEncoder.encode(pemKey, StandardCharsets.UTF_8);
    }

    /**
     * Delegates to resource cleanup (stub).
     */
    public void close() {
        throw new UnsupportedOperationException("Not yet implemented");
    }
}