package com.airlock.iam.oauth2.application.configuration.signature;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.misc.time.RealTimeProvider;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.common.domain.model.jwt.JwtClaimsSet;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.service.jwt.signature.JwtSignerException;
import com.airlock.iam.core.domain.model.http.HttpRequestHeader;
import com.airlock.iam.core.domain.model.http.HttpRequestParameter;
import com.airlock.iam.oauth2.api.domain.model.clientauthentication.ClientAuthenticationMethod;
import com.airlock.iam.oauth2.domain.model.client.OAuth2ClientSideCredentials;
import com.airlock.iam.oauth2.domain.model.client.communication.OAuth2TokenRequestClientSecretProcessor;
import com.airlock.iam.oauth2.domain.model.clientauthentication.method.KnownClientAuthenticationMethod;
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.property.BooleanProperty;
import com.airlock.iam.plugin.api.annotation.property.IntegerProperty;
import com.airlock.iam.plugin.api.annotation.property.StringProperty;
import com.airlock.iam.plugin.api.annotation.property.TestletProperty;
import com.airlock.iam.security.util.crypto.nonce.NonceGenerator;
import com.airlock.iam.security.util.crypto.nonce.SecureRandomNonceGenerator;
import com.google.common.annotations.VisibleForTesting;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTClaimsSet.Builder;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.File;
import java.net.URI;
import java.time.Instant;
import java.util.Date;
import java.util.List;
import java.util.Optional;
import org.apache.commons.codec.binary.Base64;

@Plugin(
   description = "Allows clients that have previously registered a public key on the authorization server to sign a JWT using the corresponding private key in order to authenticate with the authorization server.<p/>The JWT is formatted according to the <code>private_key_jwt</code> client authentication as specified by the <a target=\"_blank\" href=\"https://openid.net/specs/openid-connect-core-1_0.html#ClientAuthentication\">OpenID Connect Core Specification</a>.",
   displayName = "OIDC Private Key JWT Client Authentication"
)
@PluginLicenseRestriction(moduleAffiliation = "OAuthClient")
public final class OpenIDConnectPrivateKeyJwtClientAuthenticationConfig extends PrivateKeyJwtSignatureConfig implements OAuth2TokenRequestClientSecretProcessor {
   private static final String CLIENT_ASSERTION_KEY = "client_assertion";
   private static final String CLIENT_ASSERTION_TYPE_KEY = "client_assertion_type";
   private static final String CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
   @StringProperty(
      description = "Custom value used as audience claim (<code>aud</code>) of the JWT. According to the <a target=\"_blank\" href=\"https://openid.net/specs/openid-connect-core-1_0.html#ClientAuthentication\">OpenID Connect Core Specification</a>, the audience SHOULD be the URL of the authorization server's token endpoint. If not configured, the authorization server's token endpoint URL is used as \"aud\" claim.",
      exampleValues = {
            "https://airlock.com/auth/rest/oauth2/authorization-servers/as/token",
            "https://accounts.google.com/o/oauth2/token",
            "https://login.live.com/oauth20_token.srf"
      },
      optional = true
   )
   private String customAudience;
   @IntegerProperty(
      description = "The duration (in seconds) for which the generated JWT will be accepted by the authorization server. Will be used for calculating the <code>exp</code> of the JWT.",
      displayName = "JWT Validity [s]",
      optional = true,
      minValue = 1,
      defaultValue = 60,
      exampleValues = {15, 60, 600}
   )
   private int validityDuration;
   @IntegerProperty(
      description = "This claim identifies the time before which the JWT must not be accepted for processing. To determine the <code>nbf</code> claim value in the JWT, the number of seconds configured in this property are subtracted from the JWT issue time. The motivation to set a time in the past is to avoid clock synchronization problems with the authorization server.",
      optional = true,
      defaultValue = 5,
      minValue = 0,
      exampleValues = {5, 60, 3600}
   )
   private int validNotBeforeSkew;
   @BooleanProperty(
      description = "<p>If enabled the KID of the public key used to sign a JWT is added to the JWT header. Consumers of the JWT can\nuse the KID to identify the public key that is verifying the signature.</p>\n<p>The KID is ignored by some servers but is usually necessary when the client has multiple keys (or wants to\nrotate them) or when the authorization server needs to look up the corresponding public key on the client's\nJWKS endpoint.</p>\n<p>This property's testlet displays the KID that is included in the JWT header if the property is enabled.</p>\n",
      displayName = "Include KID",
      optional = true,
      defaultValue = "true"
   )
   @TestletProperty(buttonToolTip = "Displays the KID of the public key used to sign JWTs with this signer.", testletClass = OAuth2KidTestlet.class)
   private boolean includeKid;
   public static final String INCLUDE_KID_DISPLAY_NAME = "Include KID";
   public static final String INCLUDE_KID_FIELD = "includeKid";
   private NonceGenerator nonceGenerator = new SecureRandomNonceGenerator();
   private TimeProvider timeProvider = new RealTimeProvider();

   public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig(
      PrivateKeyJwtSignatureConfig.Algorithm algorithm,
      File keystoreFile,
      String keystorePassword,
      String signingKeyAlias,
      String signingKeyPassword,
      String customAudience,
      int validityDuration,
      int validNotBeforeSkew,
      boolean includeKid
   ) {
      super(algorithm, keystoreFile, keystorePassword, signingKeyAlias, signingKeyPassword);
      this.customAudience = customAudience;
      this.validityDuration = validityDuration;
      this.validNotBeforeSkew = validNotBeforeSkew;
      this.includeKid = includeKid;
   }

   @VisibleForTesting
   OpenIDConnectPrivateKeyJwtClientAuthenticationConfig setTimeProvider(TimeProvider timeProvider) {
      this.timeProvider = timeProvider;
      return this;
   }

   @VisibleForTesting
   OpenIDConnectPrivateKeyJwtClientAuthenticationConfig setNonceGenerator(NonceGenerator nonceGenerator) {
      this.nonceGenerator = nonceGenerator;
      return this;
   }

   @Override
   protected boolean isIncludeKid() {
      return this.includeKid;
   }

   @Override
   public void putClientSecretInRequestData(
      URI baseUri, List<HttpRequestParameter> params, List<HttpRequestHeader> headers, OAuth2ClientSideCredentials clientCredentials
   ) {
      String audience = Optional.ofNullable(this.customAudience).orElse(baseUri.toASCIIString());
      JsonWebSignatureToken signedToken = this.generateJws(clientCredentials.clientId().id(), audience);
      params.add(HttpRequestParameter.httpRequestParameter("client_assertion_type", "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"));
      params.add(HttpRequestParameter.httpRequestParameter("client_assertion", signedToken.tokenString()));
   }

   @Override
   public ClientAuthenticationMethod clientAuthenticationMethod() {
      return KnownClientAuthenticationMethod.PRIVATE_KEY_JWT;
   }

   private JsonWebSignatureToken generateJws(String clientId, String audience) {
      Instant now = this.timeProvider.getCurrentInstant();
      Date createdAt = Date.from(now);
      Date expiresAt = Date.from(now.plusSeconds(this.validityDuration));
      Date validNotBefore = Date.from(now.minusSeconds(this.validNotBeforeSkew));
      Builder claimSetBuilder = new Builder();
      JWTClaimsSet claimsSet = claimSetBuilder.issuer(clientId)
         .subject(clientId)
         .audience(audience)
         .jwtID(Base64.encodeBase64String(this.nonceGenerator.generateNonce()))
         .expirationTime(expiresAt)
         .issueTime(createdAt)
         .notBeforeTime(validNotBefore)
         .build();

      try {
         return this.getSigner().sign(JwtClaimsSet.of(claimsSet.toString()));
      } catch (JwtSignerException e) {
         throw new FailureException("Failed to sign the jwt used for authenticating the OAuth 2 Client.", e);
      }
   }

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

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

   @SuppressFBWarnings(justification = "generated code")
   public static class OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private PrivateKeyJwtSignatureConfig.Algorithm algorithm;
      @SuppressFBWarnings(justification = "generated code")
      private File keystoreFile;
      @SuppressFBWarnings(justification = "generated code")
      private String keystorePassword;
      @SuppressFBWarnings(justification = "generated code")
      private String signingKeyAlias;
      @SuppressFBWarnings(justification = "generated code")
      private String signingKeyPassword;
      @SuppressFBWarnings(justification = "generated code")
      private String customAudience;
      @SuppressFBWarnings(justification = "generated code")
      private int validityDuration;
      @SuppressFBWarnings(justification = "generated code")
      private int validNotBeforeSkew;
      @SuppressFBWarnings(justification = "generated code")
      private boolean includeKid;

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

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder algorithm(
         PrivateKeyJwtSignatureConfig.Algorithm algorithm
      ) {
         this.algorithm = algorithm;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder keystoreFile(File keystoreFile) {
         this.keystoreFile = keystoreFile;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder keystorePassword(
         String keystorePassword
      ) {
         this.keystorePassword = keystorePassword;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder signingKeyAlias(
         String signingKeyAlias
      ) {
         this.signingKeyAlias = signingKeyAlias;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder signingKeyPassword(
         String signingKeyPassword
      ) {
         this.signingKeyPassword = signingKeyPassword;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder customAudience(
         String customAudience
      ) {
         this.customAudience = customAudience;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder validityDuration(
         int validityDuration
      ) {
         this.validityDuration = validityDuration;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder validNotBeforeSkew(
         int validNotBeforeSkew
      ) {
         this.validNotBeforeSkew = validNotBeforeSkew;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder includeKid(boolean includeKid) {
         this.includeKid = includeKid;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectPrivateKeyJwtClientAuthenticationConfig build() {
         return new OpenIDConnectPrivateKeyJwtClientAuthenticationConfig(
            this.algorithm,
            this.keystoreFile,
            this.keystorePassword,
            this.signingKeyAlias,
            this.signingKeyPassword,
            this.customAudience,
            this.validityDuration,
            this.validNotBeforeSkew,
            this.includeKid
         );
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "OpenIDConnectPrivateKeyJwtClientAuthenticationConfig.OpenIDConnectPrivateKeyJwtClientAuthenticationConfigBuilder(algorithm="
            + this.algorithm
            + ", keystoreFile="
            + this.keystoreFile
            + ", keystorePassword="
            + this.keystorePassword
            + ", signingKeyAlias="
            + this.signingKeyAlias
            + ", signingKeyPassword="
            + this.signingKeyPassword
            + ", customAudience="
            + this.customAudience
            + ", validityDuration="
            + this.validityDuration
            + ", validNotBeforeSkew="
            + this.validNotBeforeSkew
            + ", includeKid="
            + this.includeKid
            + ")";
      }
   }
}