package at.gv.egiz.eaaf.modules.auth.sl20.utils;

import java.io.IOException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.Provider;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.Collections;
import java.util.List;

import javax.annotation.Nonnull;
import javax.annotation.PostConstruct;

import at.gv.egiz.eaaf.core.api.idp.IConfiguration;
import at.gv.egiz.eaaf.core.exception.EaafKeyAccessException;
import at.gv.egiz.eaaf.core.exceptions.EaafConfigurationException;
import at.gv.egiz.eaaf.core.impl.credential.EaafKeyStoreFactory;
import at.gv.egiz.eaaf.core.impl.credential.EaafKeyStoreUtils;
import at.gv.egiz.eaaf.core.impl.credential.KeyStoreConfiguration;
import at.gv.egiz.eaaf.core.impl.credential.KeyStoreConfiguration.KeyStoreType;
import at.gv.egiz.eaaf.core.impl.data.Pair;
import at.gv.egiz.eaaf.core.impl.utils.X509Utils;
import at.gv.egiz.eaaf.modules.auth.sl20.Constants;
import at.gv.egiz.eaaf.modules.auth.sl20.data.VerificationResult;
import at.gv.egiz.eaaf.modules.auth.sl20.exceptions.SL20Exception;
import at.gv.egiz.eaaf.modules.auth.sl20.exceptions.SL20SecurityException;
import at.gv.egiz.eaaf.modules.auth.sl20.exceptions.SlCommandoBuildException;
import at.gv.egiz.eaaf.modules.auth.sl20.exceptions.SlCommandoParserException;

import org.apache.commons.lang3.StringUtils;
import org.jose4j.jca.ProviderContext;
import org.jose4j.jwa.AlgorithmConstraints;
import org.jose4j.jwa.AlgorithmConstraints.ConstraintType;
import org.jose4j.jwe.JsonWebEncryption;
import org.jose4j.jws.AlgorithmIdentifiers;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwx.JsonWebStructure;
import org.jose4j.keys.X509Util;
import org.jose4j.keys.resolvers.X509VerificationKeyResolver;
import org.jose4j.lang.JoseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Service;
import org.springframework.util.Base64Utils;

import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.databind.JsonNode;

@Service
public class JsonSecurityUtils implements IJoseTools {
  private static final Logger log = LoggerFactory.getLogger(JsonSecurityUtils.class);

  private static final String FRIENDLYNAME_KEYSTORE = "SL2.0 KeyStore";
  private static final String FRIENDLYNAME_TRUSTSTORE = "SL2.0 TrustStore";

  @Autowired(required = true)
  IConfiguration authConfig;
  @Autowired(required = true)
  EaafKeyStoreFactory keystoreFactory;

  private Pair<KeyStore, Provider> keyStore;
  private Pair<KeyStore, Provider> trustStore;

  private static JsonMapper mapper = new JsonMapper();

  @PostConstruct
  protected void initalize() throws SL20Exception {
    log.info("Initialize SL2.0 authentication security constrains ... ");
    try {
      // load KeyStore
      final KeyStoreConfiguration keyStoreConfig = buildKeyStoreConfiguration();
      keyStore = keystoreFactory.buildNewKeyStore(keyStoreConfig);

      // load TrustStore
      final KeyStoreConfiguration trustStoreConfig = buildTrustStoreConfiguration();
      trustStore = keystoreFactory.buildNewKeyStore(trustStoreConfig);

      // validate KeyStore entries
      EaafKeyStoreUtils.getPrivateKeyAndCertificates(keyStore.getFirst(), getSigningKeyAlias(),
          getSigningKeyPassword(), true, FRIENDLYNAME_KEYSTORE);
      final Pair<Key, X509Certificate[]> encCredentials =
          EaafKeyStoreUtils.getPrivateKeyAndCertificates(keyStore.getFirst(), getEncryptionKeyAlias(),
              getEncryptionKeyPassword(), false, FRIENDLYNAME_TRUSTSTORE);
      if (encCredentials == null) {
        log.info("No encryption key for SL2.0 found. End-to-End encryption is not used.");

      }

      // validate TrustStore
      final List<X509Certificate> trustedCerts = EaafKeyStoreUtils.readCertsFromKeyStore(trustStore
          .getFirst());
      if (trustedCerts.isEmpty()) {
        log.info("No certificates in TrustStore: {}. Signature validation will FAIL!",
            FRIENDLYNAME_TRUSTSTORE);

      } else {
        log.info("Find #{} certificates in TrustStore: {}",
            trustedCerts.size(), FRIENDLYNAME_TRUSTSTORE);

      }

      log.info("SL2.0 authentication security constrains initialized.");

    } catch (final RuntimeException e) {
      throw e;

    } catch (final Exception e) {
      log.error("SL2.0 security constrains initialization FAILED.");
      throw new SL20Exception("sl20.11", new Object[] { e.getMessage() }, e);

    }

  }

  @Override
  public String createSignature(final String payLoad) throws SlCommandoBuildException {
    return createSignature(payLoad, true);

  }

  @Override
  public String createSignature(final String payLoad, boolean addFullCertChain) throws SlCommandoBuildException {
    try {
      final JsonWebSignature jws = new JsonWebSignature();

      // set payload
      jws.setPayload(payLoad);

      // set basic header
      jws.setContentTypeHeaderValue(SL20Constants.SL20_CONTENTTYPE_SIGNED_COMMAND);

      // set signing information
      jws.setAlgorithmHeaderValue(AlgorithmIdentifiers.RSA_USING_SHA256);
      final Pair<Key, X509Certificate[]> signingCred = EaafKeyStoreUtils.getPrivateKeyAndCertificates(
          keyStore.getFirst(), getSigningKeyAlias(), getSigningKeyPassword(), true, FRIENDLYNAME_KEYSTORE);
      jws.setKey(signingCred.getFirst());

      // set special provider if required
      if (keyStore.getSecond() != null) {
        log.trace("Injecting special Java Security Provider: {}", keyStore.getSecond().getName());
        final ProviderContext providerCtx = new ProviderContext();
        providerCtx.getSuppliedKeyProviderContext().setSignatureProvider(
            keyStore.getSecond().getName());
        jws.setProviderContext(providerCtx);

      }

      if (addFullCertChain) {
        jws.setCertificateChainHeaderValue(signingCred.getSecond());

      }

      jws.setX509CertSha256ThumbprintHeaderValue(signingCred.getSecond()[0]);

      return jws.getCompactSerialization();

    } catch (final JoseException | EaafKeyAccessException e) {
      log.warn("Can NOT sign SL2.0 command.", e);
      throw new SlCommandoBuildException("Can NOT sign SL2.0 command.", e);

    }

  }

  @Override
  public VerificationResult validateSignature(final String serializedContent, final KeyStore trustStore,
      final AlgorithmConstraints algconstraints) throws JoseException, IOException, KeyStoreException {
    final List<X509Certificate> trustedCertificates = EaafKeyStoreUtils.readCertsFromKeyStore(trustStore);
    return validateSignature(serializedContent, trustedCertificates, algconstraints);

  }

  @Override
  @NonNull
  public VerificationResult validateSignature(@Nonnull final String serializedContent,
      @Nonnull final List<X509Certificate> trustedCerts, @Nonnull final AlgorithmConstraints constraints)
      throws JoseException, IOException {
    final JsonWebSignature jws = new JsonWebSignature();
    // set payload
    jws.setCompactSerialization(serializedContent);

    // set security constrains
    jws.setAlgorithmConstraints(constraints);

    // load signinc certs
    Key selectedKey = null;
    final List<X509Certificate> x5cCerts = jws.getCertificateChainHeaderValue();
    final String x5t256 = jws.getX509CertSha256ThumbprintHeaderValue();
    if (x5cCerts != null) {
      log.debug("Found x509 certificate in JOSE header ... ");
      log.trace("Sorting received X509 certificates ... ");
      final List<X509Certificate> sortedX5cCerts = X509Utils.sortCertificates(x5cCerts);

      if (trustedCerts.contains(sortedX5cCerts.get(0))) {
        selectedKey = sortedX5cCerts.get(0).getPublicKey();

      } else {
        log.info("Can NOT find JOSE certificate in truststore.");
        try {
          log.debug("Cert: " + Base64Utils.encodeToString(sortedX5cCerts.get(0).getEncoded()));

        } catch (final CertificateEncodingException e) {
          e.printStackTrace();

        }

      }

    } else if (StringUtils.isNotEmpty(x5t256)) {
      log.debug("Found x5t256 fingerprint in JOSE header .... ");
      final X509VerificationKeyResolver x509VerificationKeyResolver = new X509VerificationKeyResolver(
          trustedCerts);
      selectedKey = x509VerificationKeyResolver.resolveKey(jws, Collections.<JsonWebStructure>emptyList());

    } else {
      throw new JoseException("JWS contains NO signature certificate or NO certificate fingerprint");

    }

    if (selectedKey == null) {
      throw new JoseException("Can NOT select verification key for JWS. Signature verification FAILED");

    }

    // set verification key
    jws.setKey(selectedKey);

    // load payLoad
    return new VerificationResult(mapper.getMapper().readTree(jws.getPayload()), null, jws.verifySignature());

  }

  @Override
  @Nonnull
  public VerificationResult validateSignature(@Nonnull final String serializedContent) throws SL20Exception {
    try {
      final AlgorithmConstraints algConstraints = new AlgorithmConstraints(ConstraintType.WHITELIST,
          SL20Constants.SL20_ALGORITHM_WHITELIST_SIGNING
              .toArray(new String[SL20Constants.SL20_ALGORITHM_WHITELIST_SIGNING.size()]));

      final VerificationResult result =
          validateSignature(serializedContent, EaafKeyStoreUtils.readCertsFromKeyStore(trustStore.getFirst()),
              algConstraints);

      if (!result.isValidSigned()) {
        log.info("JWS signature invalide. Stopping authentication process ...");
        log.debug("Received JWS msg: " + serializedContent);
        throw new SL20SecurityException("JWS signature invalide.");

      }

      log.debug("SL2.0 commando signature validation sucessfull");
      return result;

    } catch (JoseException | JsonParseException | KeyStoreException e) {
      log.warn("SL2.0 commando signature validation FAILED", e);
      throw new SL20SecurityException(new Object[] { e.getMessage() }, e);

    } catch (final IOException e) {
      log.warn("Decrypted SL2.0 result can not be parsed.", e);
      throw new SlCommandoParserException("Decrypted SL2.0 result can not be parsed", e);

    }

  }

  @Override
  public JsonNode decryptPayload(final String compactSerialization) throws SL20Exception {
    try {
      final JsonWebEncryption receiverJwe = new JsonWebEncryption();

      // set security constrains
      receiverJwe.setAlgorithmConstraints(
          new AlgorithmConstraints(ConstraintType.WHITELIST,
              SL20Constants.SL20_ALGORITHM_WHITELIST_KEYENCRYPTION
                  .toArray(new String[SL20Constants.SL20_ALGORITHM_WHITELIST_KEYENCRYPTION.size()])));
      receiverJwe.setContentEncryptionAlgorithmConstraints(
          new AlgorithmConstraints(ConstraintType.WHITELIST, SL20Constants.SL20_ALGORITHM_WHITELIST_ENCRYPTION
              .toArray(new String[SL20Constants.SL20_ALGORITHM_WHITELIST_ENCRYPTION.size()])));

      // set payload
      receiverJwe.setCompactSerialization(compactSerialization);

      final Pair<Key, X509Certificate[]> encryptionCred = EaafKeyStoreUtils.getPrivateKeyAndCertificates(
          keyStore.getFirst(), getEncryptionKeyAlias(), getEncryptionKeyPassword(), true,
          FRIENDLYNAME_KEYSTORE);

      // validate key from header against key from config
      final List<X509Certificate> x5cCerts = receiverJwe.getCertificateChainHeaderValue();
      final String x5t256 = receiverJwe.getX509CertSha256ThumbprintHeaderValue();
      if (x5cCerts != null) {
        log.debug("Found x509 certificate in JOSE header ... ");
        log.trace("Sorting received X509 certificates ... ");
        final List<X509Certificate> sortedX5cCerts = X509Utils.sortCertificates(x5cCerts);

        if (!sortedX5cCerts.get(0).equals(encryptionCred.getSecond()[0])) {
          log.info("Certificate from JOSE header does NOT match encryption certificate");

          try {

            log.debug("JOSE certificate: {}", Base64Utils.encode(sortedX5cCerts.get(0).getEncoded()));
          } catch (final CertificateEncodingException e) {
            e.printStackTrace();
          }
          throw new SL20Exception("sl20.05",
              new Object[] { "Certificate from JOSE header does NOT match encryption certificate" });
        }

      } else if (StringUtils.isNotEmpty(x5t256)) {
        log.debug("Found x5t256 fingerprint in JOSE header .... ");
        final String certFingerPrint = X509Util.x5tS256(encryptionCred.getSecond()[0]);
        if (!certFingerPrint.equals(x5t256)) {
          log.info("X5t256 from JOSE header does NOT match encryption certificate");
          log.debug("X5t256 from JOSE header: " + x5t256 + " Encrytption cert: " + certFingerPrint);
          throw new SL20Exception("sl20.05",
              new Object[] { "X5t256 from JOSE header does NOT match encryption certificate" });

        }

      } else {
        log.info("Signed SL2.0 response contains NO signature certificate or NO certificate fingerprint");
        throw new SlCommandoParserException(
            "Signed SL2.0 response contains NO signature certificate or NO certificate fingerprint");

      }

      // set key
      receiverJwe.setKey(encryptionCred.getFirst());

      // decrypt payload
      return mapper.getMapper().readTree(receiverJwe.getPlaintextString());

    } catch (final JoseException | EaafKeyAccessException e) {
      log.warn("SL2.0 result decryption FAILED", e);
      throw new SL20SecurityException(new Object[] { e.getMessage() }, e);

    } catch (final JsonParseException e) {
      log.warn("Decrypted SL2.0 result is NOT a valid JSON.", e);
      throw new SlCommandoParserException("Decrypted SL2.0 result is NOT a valid JSON.", e);

    } catch (final IOException e) {
      log.warn("Decrypted SL2.0 result can not be parsed.", e);
      throw new SlCommandoParserException("Decrypted SL2.0 result can not be parsed", e);

    }
  }

  @Override
  public X509Certificate getEncryptionCertificate() {
    Pair<Key, X509Certificate[]> encryptionCred;
    try {
      encryptionCred = EaafKeyStoreUtils.getPrivateKeyAndCertificates(keyStore.getFirst(),
          getEncryptionKeyAlias(), getEncryptionKeyPassword(), false, FRIENDLYNAME_KEYSTORE);
      if (encryptionCred != null && encryptionCred.getSecond().length > 0) {
        return encryptionCred.getSecond()[0];

      }

    } catch (final EaafKeyAccessException e) {
      log.trace("Exception is skipped because Encryption is not mandatory on this level", e);

    }

    return null;

  }

  private KeyStoreConfiguration buildKeyStoreConfiguration() throws EaafConfigurationException {
    final KeyStoreConfiguration config = new KeyStoreConfiguration();
    config.setFriendlyName(FRIENDLYNAME_KEYSTORE);

    config.setKeyStoreType(authConfig.getBasicConfiguration(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_TYPE),
        KeyStoreType.JKS.getKeyStoreType()));
    config.setKeyStoreName(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_NAME));
    config.setSoftKeyStoreFilePath(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_PATH));
    config.setSoftKeyStorePassword(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_PASSWORD));

    // validate configuration state
    config.validate();

    return config;

  }

  private KeyStoreConfiguration buildTrustStoreConfiguration() throws EaafConfigurationException {
    final KeyStoreConfiguration config = new KeyStoreConfiguration();
    config.setFriendlyName(FRIENDLYNAME_TRUSTSTORE);

    config.setKeyStoreType(authConfig.getBasicConfiguration(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_TRUSTSTORE_TYPE),
        KeyStoreType.JKS.getKeyStoreType()));
    config.setKeyStoreName(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_TRUSTSTORE_NAME));
    config.setSoftKeyStoreFilePath(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_TRUSTSTORE_PATH));
    config.setSoftKeyStorePassword(
        authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_TRUSTSTORE_PASSWORD));

    // validate configuration state
    config.validate();

    return config;
  }

  private String getSigningKeyAlias() {
    String value = authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_KEY_SIGN_ALIAS);
    if (value != null) {
      value = value.trim();
    }

    return value;
  }

  private char[] getSigningKeyPassword() {
    final String value = authConfig.getBasicConfiguration(
        Constants.CONFIG_PROP_SECURITY_KEYSTORE_KEY_SIGN_PASSWORD);
    if (value != null) {
      return value.trim().toCharArray();
    }

    return null;
  }

  private String getEncryptionKeyAlias() {
    String value = authConfig.getBasicConfiguration(
        Constants.CONFIG_PROP_SECURITY_KEYSTORE_KEY_ENCRYPTION_ALIAS);
    if (value != null) {
      value = value.trim();
    }

    return value;
  }

  private char[] getEncryptionKeyPassword() {
    final String value = authConfig.getBasicConfiguration(
        Constants.CONFIG_PROP_SECURITY_KEYSTORE_KEY_ENCRYPTION_PASSWORD);
    if (value != null) {
      return value.trim().toCharArray();
    }

    return null;
  }

}