package com.example.crypto.algorithms;

//import org.bouncycastle.jce.provider.BouncyCastleProvider;
import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.util.Base64;
import java.util.Arrays;
import javax.crypto.Mac;

/**
 * This class demonstrates encryption schemes using elliptic-curve
 * Diffie-Hellman (ECDH)
 * and hybrid encryption methods, including a post-quantum hybrid scheme.
 *
 * SAST/CBOM Classification:
 * 
 * 1. EC Key Generation & ECDH Key Agreement:
 * - Parent Classification: Asymmetric Key Generation / Key Agreement.
 * - SAST: Secure when using established curves (secp256r1) and reputable
 * providers (BouncyCastle).
 *
 * 2. ECDH Hybrid Encryption:
 * - Parent Classification: Hybrid Cryptosystem (ECDH + AEAD).
 * - SAST: Uses ECDH for key agreement and AES/GCM for encryption. However, the
 * derivation of an AES key
 * by applying a single SHA-256 hash to the shared secret may be flagged as a
 * weak key derivation method.
 * A dedicated KDF (e.g., HKDF) is recommended.
 *
 * 3. Post-Quantum Hybrid Encryption:
 * - Parent Classification: Hybrid Cryptosystem (Classical ECDH + Post-Quantum
 * Secret + KDF + AEAD).
 * - SAST: Combining classical and post-quantum components is advanced and
 * secure if implemented properly.
 * The custom HKDF expand function provided here is simplistic and may be
 * flagged in a CBOM analysis;
 * a standard HKDF library should be used in production.
 */
public class Encryption2 {

    // static {
    //     Security.addProvider(new BouncyCastleProvider());
    // }

    /**
     * Generates an Elliptic Curve (EC) key pair using the secp256r1 curve.
     *
     * SAST/CBOM Notes:
     * - Algorithm: EC key pair generation.
     * - Parent Classification: Asymmetric Key Generation.
     * - SAST: Considered secure when using strong randomness and a reputable
     * provider.
     *
     * @return an EC KeyPair.
     */
    public KeyPair generateECKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC", "BC");
        keyPairGenerator.initialize(new ECGenParameterSpec("secp256r1"), new SecureRandom());
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * Derives a shared secret using Elliptic Curve Diffie-Hellman (ECDH).
     *
     * SAST/CBOM Notes:
     * - Algorithm: ECDH key agreement.
     * - Parent Classification: Asymmetric Key Agreement.
     * - SAST: Secure when both parties use strong EC keys and proper randomness.
     *
     * @param privateKey the private key of one party.
     * @param publicKey  the public key of the other party.
     * @return the derived shared secret as a byte array.
     */
    public byte[] deriveSharedSecret(PrivateKey privateKey, PublicKey publicKey) throws Exception {
        KeyAgreement keyAgreement = KeyAgreement.getInstance("ECDH", "BC");
        keyAgreement.init(privateKey);
        keyAgreement.doPhase(publicKey, true);
        return keyAgreement.generateSecret();
    }

    /**
     * Performs hybrid encryption using ECDH to derive a shared secret, then derives
     * an AES key
     * by hashing the shared secret with SHA-256, and finally encrypts the data with
     * AES-GCM.
     *
     * SAST/CBOM Notes:
     * - Parent Classification: Hybrid Cryptosystem (ECDH + AES-GCM).
     * - SAST: While ECDH and AES-GCM are secure, the key derivation method here (a
     * single SHA-256 hash)
     * is not as robust as using a dedicated KDF. This approach may be flagged and
     * is recommended for
     * improvement.
     *
     * @param recipientPublicKey the recipient's public EC key.
     * @param data               the plaintext data to encrypt.
     */
    public void ecdhHybridEncryption(PublicKey recipientPublicKey, String data) throws Exception {
        // Generate an ephemeral EC key pair for the sender.
        KeyPair senderKeyPair = generateECKeyPair();
        // Derive the shared secret using ECDH.
        byte[] sharedSecret = deriveSharedSecret(senderKeyPair.getPrivate(), recipientPublicKey);

        // Derive an AES key by hashing the shared secret with SHA-256.
        // SAST Note: Using a direct hash for key derivation is simplistic and may be
        // flagged.
        MessageDigest sha256 = MessageDigest.getInstance("SHA-256");
        byte[] aesKeyBytes = sha256.digest(sharedSecret);
        // Use the first 16 bytes (128 bits) as the AES key.
        SecretKey aesKey = new SecretKeySpec(aesKeyBytes, 0, 16, "AES");

        // Encrypt the data using AES-GCM.
        Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
        byte[] iv = new byte[12]; // 12-byte IV recommended for GCM.
        new SecureRandom().nextBytes(iv);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv); // 128-bit authentication tag.
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey, gcmSpec);
        byte[] encryptedData = aesCipher.doFinal(data.getBytes());

        System.out.println(
                "ECDH Hybrid Encryption - Encrypted Data: " + Base64.getEncoder().encodeToString(encryptedData));
    }

    /**
     * Performs post-quantum hybrid encryption by combining a classical ECDH-derived
     * secret with a
     * post-quantum shared secret. The two secrets are combined using a custom HKDF
     * expansion, and the
     * derived key is used to encrypt data with AES-GCM.
     *
     * SAST/CBOM Notes:
     * - Parent Classification: Hybrid Cryptosystem (Classical ECDH + Post-Quantum
     * Secret + KDF + AES-GCM).
     * - SAST: The combination of classical and post-quantum secrets is a modern
     * approach. However, the
     * custom HKDF expand function is simplistic and may be flagged as insecure. Use
     * a standard HKDF
     * implementation in production.
     *
     * @param ecPublicKey    the recipient's EC public key.
     * @param pqSharedSecret the post-quantum shared secret from a separate
     *                       algorithm.
     */
    public void postQuantumHybridEncryption(PublicKey ecPublicKey, byte[] pqSharedSecret) throws Exception {
        // Step 1: Perform classical ECDH key agreement to derive a shared secret.
        byte[] ecdhSharedSecret = deriveSharedSecret(generateECKeyPair().getPrivate(), ecPublicKey);

        // Step 2: Combine the ECDH secret and the post-quantum secret using a
        // simplified HKDF expansion.
        // SAST Note: This custom HKDF implementation is minimal and does not follow the
        // full HKDF spec.
        byte[] combinedSecret = hkdfExpand(ecdhSharedSecret, pqSharedSecret, 32);
        // Use the first 16 bytes as the AES key (128-bit key).
        SecretKey aesKey = new SecretKeySpec(combinedSecret, 0, 16, "AES");

        // Step 3: Encrypt the data using AES-GCM.
        Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
        byte[] iv = new byte[12]; // 12-byte IV recommended for GCM.
        new SecureRandom().nextBytes(iv);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey, gcmSpec);
        byte[] encryptedData = aesCipher.doFinal("Post-Quantum Hybrid Encryption Data".getBytes());

        System.out.println("Post-Quantum Hybrid Encryption - Encrypted Data: "
                + Base64.getEncoder().encodeToString(encryptedData));
    }

    /**
     * A simplified HKDF expansion function that uses HMAC-SHA256 to derive a key of
     * a desired length.
     *
     * SAST/CBOM Notes:
     * - Parent Classification: Key Derivation Function (KDF).
     * - SAST: Custom KDF implementations are risky if not thoroughly vetted. This
     * simple HKDF expand
     * function lacks the full HKDF mechanism (e.g., multiple iterations, info, and
     * context parameters)
     * and may be flagged. It is recommended to use a standardized HKDF library.
     *
     * @param inputKey the input key material.
     * @param salt     a salt value (here, the post-quantum shared secret is used as
     *                 the salt).
     * @param length   the desired length of the derived key.
     * @return a derived key of the specified length.
     */
    private byte[] hkdfExpand(byte[] inputKey, byte[] salt, int length) throws Exception {
        Mac hmac = Mac.getInstance("HmacSHA256");
        SecretKey secretKey = new SecretKeySpec(salt, "HmacSHA256");
        hmac.init(secretKey);
        byte[] extractedKey = hmac.doFinal(inputKey);
        return Arrays.copyOf(extractedKey, length);
    }
}