package com.airlock.iam.common.application.configuration.e2ee;

import com.airlock.iam.common.api.application.configuration.jwt.JWKPairConfig;
import com.airlock.iam.plugin.api.annotation.config.ConfigurationValidator;
import com.airlock.iam.plugin.api.annotation.config.Plugin;
import com.airlock.iam.plugin.api.annotation.license.PluginLicenseRestriction;
import com.airlock.iam.plugin.api.annotation.lifecycle.AfterConfigActivation;
import com.airlock.iam.plugin.api.annotation.property.PluginProperty;
import com.airlock.iam.plugin.api.annotation.property.StringProperty;
import com.airlock.iam.plugin.api.config.ConfigurationContext;
import com.airlock.iam.plugin.api.config.ConfigurationEnvironment;
import com.airlock.iam.plugin.api.config.Environment;
import com.airlock.iam.plugin.api.config.PluginConfiguration;
import com.airlock.iam.plugin.api.exception.ConfigurationException;
import com.airlock.iam.plugin.api.exception.InitializeException;
import com.airlock.iam.plugin.api.exception.PropertyConfigurationException;
import com.airlock.iam.plugin.api.valuereader.StringPropertyValueReader;
import com.fasterxml.jackson.core.JsonPointer;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWEDecrypter;
import com.nimbusds.jose.crypto.ECDHDecrypter;
import com.nimbusds.jose.crypto.RSADecrypter;
import com.nimbusds.jose.crypto.X25519Decrypter;
import com.nimbusds.jose.jwk.ECKey;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.KeyType;
import com.nimbusds.jose.jwk.OctetKeyPair;
import com.nimbusds.jose.jwk.RSAKey;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;

@Plugin(
   description = "<p>Decryption service accepting encrypted passwords in JWE format (as specified in <a href=\"https://tools.ietf.org/html/rfc7516\" target=\"_blank\">https://tools.ietf.org/html/rfc7516</a>).</p><p>The encryption information contained in the step responses (additional data \"e2eEncryptionInformation\") consists of a type (always \"type\"=\"JWE\"), a nonce (element \"nonce\") and the public key of the configured key pair in JWK format (element \"publicKey\").</p><p>The password sent by the client for verification must be encrypted in JWE format and include the nonce in the header. The JWE payload must include the password as a JSON string. The location of the password in the payload JSON structure can be configured. By default the JWE payload is expected to be in the following format:</p><pre>{\n    \"header\": {\n        \"nonce\":\"si04fHDORRELOO0T4nJad8mz9DgPPE9GhArD2reQ2Dk=\",\n        \"alg\":\"RSA-OAEP-256\",\n        \"enc\":\"A128GCM\"\n     },\n    \"password\":\"userPasswordInPlaintext\"\n}</pre><p>For decryption, the private key of the configured \"Key Pair\" is used.</p>",
   displayName = "JWE Password Decryption"
)
@PluginLicenseRestriction(moduleAffiliation = "EndToEndPasswordEncryption")
public class JwePasswordDecryptionConfig implements PasswordDecryptionConfig {
   @StringProperty(
      description = "The path which points to the password contained in the payload of the JWE in JSON Pointer format as specified in <a href=\"https://tools.ietf.org/html/rfc6901\" target=\"_blank\">https://tools.ietf.org/html/rfc6901</a> The referenced plaintext password must be a JSON string.<p>If for example the payload of the JWE looks as follows:</p><pre>{\n  \"user\": {\n    \"id\": \"8331-1212-1233\",\n    \"password\": \"userPasswordInPlaintext\"\n    }\n}</pre>Then the following JSON Pointer should be configured in order to use \"userPasswordInPlaintext\" for the password check: /user/password ",
      exampleValues = {"/password", "/user/passwords/0"},
      defaultValue = "/password",
      optional = true
   )
   private String passwordJsonPath;
   private static final String PASSWORD_JSON_POINTER = "passwordJsonPath";
   @PluginProperty(
      description = "The configuration of the public/private keypair used for encrypting/decrypting the passwords. The public key information, needed for encrypting the JWE, will be returned by the flow step.",
      optional = false
   )
   private JWKPairConfig keyPair;
   private JsonPointer passwordJsonPointer;
   private JWEDecrypter jweDecrypter;

   public JwePasswordDecryptionConfig(String passwordJsonPath, JWKPairConfig keyPair) {
      this.passwordJsonPath = passwordJsonPath;
      this.keyPair = keyPair;
   }

   @AfterConfigActivation
   void compileJsonPath() throws InitializeException {
      this.passwordJsonPointer = JsonPointer.compile(this.passwordJsonPath);
      this.jweDecrypter = this.initializeJweDecrypter();
   }

   @ConfigurationValidator(localOnly = true)
   public static void validateConfig(PluginConfiguration config, ConfigurationEnvironment env, ConfigurationContext context, Environment validationEnvironment) throws ConfigurationException {
      String passwordPath = StringPropertyValueReader.effectiveStringConfigValue(config, "passwordJsonPath", env, context);

      try {
         JsonPointer.compile(passwordPath);
      } catch (IllegalArgumentException e) {
         throw new PropertyConfigurationException(
            config, env, context, "The configured JSON Path is invalid: " + e.getMessage(), e, new String[]{"passwordJsonPath"}
         );
      }
   }

   private JWEDecrypter initializeJweDecrypter() throws InitializeException {
      JWK jwkPair = this.keyPair.getKeyPair();
      KeyType keyType = jwkPair.getKeyType();

      try {
         if (KeyType.RSA.equals(keyType)) {
            return new RSADecrypter((RSAKey)jwkPair);
         }

         if (KeyType.EC.equals(keyType)) {
            return new ECDHDecrypter((ECKey)jwkPair);
         }

         if (KeyType.OKP.equals(keyType)) {
            return new X25519Decrypter((OctetKeyPair)jwkPair);
         }
      } catch (JOSEException e) {
         throw new InitializeException("Could not cast the provided jwk pair", e);
      }

      throw new InitializeException(
         String.format("The provided key type (%s) is not in the allowed list: [%s, %s, %s]", keyType, KeyType.RSA, KeyType.EC, KeyType.OKP)
      );
   }

   public JWKPairConfig getKeyPair() {
      return this.keyPair;
   }

   public JsonPointer getPasswordJsonPointer() {
      return this.passwordJsonPointer;
   }

   public JWEDecrypter getJweDecrypter() {
      return this.jweDecrypter;
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwePasswordDecryptionConfig.JwePasswordDecryptionConfigBuilder builder() {
      return new JwePasswordDecryptionConfig.JwePasswordDecryptionConfigBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public JwePasswordDecryptionConfig() {
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwePasswordDecryptionConfigBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private String passwordJsonPath;
      @SuppressFBWarnings(justification = "generated code")
      private JWKPairConfig keyPair;

      @SuppressFBWarnings(justification = "generated code")
      JwePasswordDecryptionConfigBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwePasswordDecryptionConfig.JwePasswordDecryptionConfigBuilder passwordJsonPath(String passwordJsonPath) {
         this.passwordJsonPath = passwordJsonPath;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwePasswordDecryptionConfig.JwePasswordDecryptionConfigBuilder keyPair(JWKPairConfig keyPair) {
         this.keyPair = keyPair;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwePasswordDecryptionConfig build() {
         return new JwePasswordDecryptionConfig(this.passwordJsonPath, this.keyPair);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwePasswordDecryptionConfig.JwePasswordDecryptionConfigBuilder(passwordJsonPath=" + this.passwordJsonPath + ", keyPair=" + this.keyPair + ")";
      }
   }
}
