//from https://docs.aws.amazon.com/cloudhsm/latest/userguide/java-sample-ec-key.html

package com.amazonaws.cloudhsm.examples.key.asymmetric;

import java.security.InvalidAlgorithmParameterException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;

import com.amazonaws.cloudhsm.examples.operations.LoginLogoutExample;
import com.cavium.cfm2.CFM2Exception;
import com.cavium.cfm2.Util;
import com.cavium.key.CaviumECPrivateKey;
import com.cavium.key.CaviumECPublicKey;
import com.cavium.key.CaviumKey;
import com.cavium.key.parameter.CaviumECGenParameterSpec;

public class ECAsymmetricKeyGeneration {

  String[] supportedSpec = {"prime256v1", "secp256r1", "secp384r1"}; 

  public static void main(String[] z) {
    LoginLogoutExample.loginWithEnvVars();
    ECAsymmetricKeyGeneration obj = new ECAsymmetricKeyGeneration();
    obj.generateECKeyPair("secp256r1", true);
    obj.generateECKeyPair("secp256r1", "MyECPublicKey" , "MyECPrivateKey" , false ,true);
    LoginLogoutExample.logout();
  }

  public KeyPair generateECKeyPair(String spec, boolean isPersistent) {
    KeyPairGenerator keyPairGen;
    try {
    
      // Create an instance of the provider.
      keyPairGen = KeyPairGenerator.getInstance("EC", "Cavium");
      
      // Generate the key pair.       
      CaviumECGenParameterSpec paramSpec = new CaviumECGenParameterSpec(spec);
      keyPairGen.initialize(paramSpec);
      KeyPair kp = keyPairGen.generateKeyPair();
      if (kp == null) {
        System.out.println("Failed to generate keypair");
      }

      // Get the key pair.
      PrivateKey privKey = kp.getPrivate();
      PublicKey pubKey = kp.getPublic();
      System.out.println("Generated EC Key Pair!");
      
      if (privKey instanceof CaviumECPrivateKey) {
        CaviumECPrivateKey cavEcPrivateKey = (CaviumECPrivateKey) privKey;
        
        // Save the private key handle. You'll need this to perform future encryption and decryption operations.
        System.out.println("Private Key Handle = " + cavEcPrivateKey.getHandle());
        
        // Get the private key label generated by the SDK.
        System.out.println("Private Key Label = " + cavEcPrivateKey.getLabel());
        
        // Get the Extractable property of the private key.
        System.out.println("Is Private Key Extractalbe = " + cavEcPrivateKey.isExtractable());
        
        // Get the Persistent property of the private key.
        System.out.println("Is Private Key Persistent = " + cavEcPrivateKey.isPersistent());
        
        // By default, keys are not persistent. Make them Persistent here.
        if(isPersistent) {
          System.out.println("Setting Private Key as Persistent:");
          makeKeyPersistent(cavEcPrivateKey);
          System.out.println("Added EC Private Key to HSM");
        }
        System.out.println("Is Private Key Persistent = " + cavEcPrivateKey.isPersistent());
      }

      if(pubKey instanceof CaviumECPublicKey) {
        CaviumECPublicKey cavEcPublicKey = (CaviumECPublicKey) pubKey;
        
        // Save the public key handle. You'll need this to perform future encryption and decryption operations.
        System.out.println("Public Key Handle = " + cavEcPublicKey.getHandle());
        
        // Get the public key label generated by the SDK.
        System.out.println("Public Key Label = " + cavEcPublicKey.getLabel());
        
        // Get the Extractable property of the public key.
        System.out.println("Is Public Key Extractalbe = " +cavEcPublicKey.isExtractable());
        
        // Get the Persistent property of the public key.
        System.out.println("Is Public Key Persistent = " + cavEcPublicKey.isPersistent());
        
        // By default, keys are not persistent. Make them Persistent here.
        if(isPersistent) {
          System.out.println("Setting Public Key as Persistent:");
          makeKeyPersistent(cavEcPublicKey);
          System.out.println("Added EC Public Key to HSM");
        }
        System.out.println("Is Private Key Persistent = " + cavEcPublicKey.isPersistent());
      }
      
      return kp;
      
    } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
      e.printStackTrace();
    } catch (InvalidAlgorithmParameterException e) {
      e.printStackTrace();
    }
    return null;
  }
  
  //If invoking this method, you can specify public key label, private key label, if private key can be extracted and if key pair is Persistent
  // This method allows you to specify the following parameters:
  //   - The EC parameter spec.
  //   - The public and private key labels.
  //   - A Boolean value that specifies whether the keys can be extracted.
  //   - A Boolean value that specifies whether the keys should be saved to an HSM.
  //  
  public KeyPair generateECKeyPair(String spec, String publicKeyLabel, String privateKeyLabel, boolean isExtractable, boolean isPersistent) {
    KeyPairGenerator keyPairGen;
    try {
    
      // Create an instance of the provider.
      keyPairGen = KeyPairGenerator.getInstance("EC", "Cavium");
      
      // Generate the key pair.     
      CaviumECGenParameterSpec paramSpec = new CaviumECGenParameterSpec(spec, publicKeyLabel, privateKeyLabel, isExtractable, isPersistent);
      keyPairGen.initialize(paramSpec);
      KeyPair kp = keyPairGen.generateKeyPair();
      if (kp == null) {
        System.out.println("Failed to generate keypair");
      }

      // Get the key pair.
      PrivateKey privKey = kp.getPrivate();
      PublicKey pubKey = kp.getPublic();
      System.out.println("Generated EC Key Pair!");
      
      if (privKey instanceof CaviumECPrivateKey) {
        CaviumECPrivateKey cavEcPrivateKey = (CaviumECPrivateKey) privKey;
        
        // Save the private key handle. You'll need this to perform future encryption and decryption operations.
        System.out.println("Private Key Handle = " + cavEcPrivateKey.getHandle());
        
        // Get the private key label generated by the SDK.
        System.out.println("Private Key Label = " + cavEcPrivateKey.getLabel());
        
        // Get the Extractable property of the private key.
        System.out.println("Is Private Key Extractalbe = " + cavEcPrivateKey.isExtractable());
        
        // Get the Persistent property of the private key.
        System.out.println("Is Private Key Persistent = " + cavEcPrivateKey.isPersistent());
      }

      if(pubKey instanceof CaviumECPublicKey) {
        CaviumECPublicKey cavEcPublicKey = (CaviumECPublicKey) pubKey;
        
        // Save the public key handle. You'll need this to perform future encryption and decryption operations.
        System.out.println("Public Key Handle = " + cavEcPublicKey.getHandle());
        
        // Get the public key label generated by the SDK.
        System.out.println("Public Key Label = " + cavEcPublicKey.getLabel());

        // Get the Extractable property of the public key.
        System.out.println("Is Public Key Extractalbe = " +cavEcPublicKey.isExtractable());
        
        // Get the Persistent property of the public key.
        System.out.println("Is Public Key Persistent = " + cavEcPublicKey.isPersistent());
      }
      
      return kp;
      
    } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
      e.printStackTrace();
    } catch (InvalidAlgorithmParameterException e) {
      e.printStackTrace();
    }
    return null;
  }

  // Save the key to the HSM.
  protected void makeKeyPersistent(CaviumKey key) {
    CaviumKey rsaKey = (CaviumKey) key;
    try {
      Util.persistKey(rsaKey);
    } 
    catch (CFM2Exception e) {
      e.printStackTrace();
    }
  }
}
