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

import java.io.IOException;
import java.net.MalformedURLException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;

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

import org.apache.commons.lang3.StringUtils;
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;

import at.gv.egiz.eaaf.core.api.idp.IConfiguration;
import at.gv.egiz.eaaf.core.exceptions.EaafConfigurationException;
import at.gv.egiz.eaaf.core.impl.utils.FileUtils;
import at.gv.egiz.eaaf.core.impl.utils.KeyStoreUtils;
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;

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

  @Autowired(required = true)
  IConfiguration authConfig;
  private Key signPrivKey = null;
  private X509Certificate[] signCertChain = null;

  private Key encPrivKey = null;
  private X509Certificate[] encCertChain = null;

  private List<X509Certificate> trustedCerts = new ArrayList<>();

  private static JsonMapper mapper = new JsonMapper();

  @PostConstruct
  protected void initalize() {
    log.info("Initialize SL2.0 authentication security constrains ... ");
    try {
      if (getKeyStoreFilePath() != null) {
        final KeyStore keyStore = KeyStoreUtils.loadKeyStore(getKeyStoreFilePath(), getKeyStorePassword());

        // load signing key
        signPrivKey = keyStore.getKey(getSigningKeyAlias(), getSigningKeyPassword().toCharArray());
        final Certificate[] certChainSigning = keyStore.getCertificateChain(getSigningKeyAlias());
        signCertChain = new X509Certificate[certChainSigning.length];
        for (int i = 0; i < certChainSigning.length; i++) {
          if (certChainSigning[i] instanceof X509Certificate) {
            signCertChain[i] = (X509Certificate) certChainSigning[i];
          } else {
            log.warn("NO X509 certificate for signing: " + certChainSigning[i].getType());
          }

        }

        // load encryption key
        try {
          encPrivKey = keyStore.getKey(getEncryptionKeyAlias(), getEncryptionKeyPassword().toCharArray());
          if (encPrivKey != null) {
            final Certificate[] certChainEncryption = keyStore.getCertificateChain(getEncryptionKeyAlias());
            encCertChain = new X509Certificate[certChainEncryption.length];
            for (int i = 0; i < certChainEncryption.length; i++) {
              if (certChainEncryption[i] instanceof X509Certificate) {
                encCertChain[i] = (X509Certificate) certChainEncryption[i];
              } else {
                log.warn("NO X509 certificate for encryption: " + certChainEncryption[i].getType());
              }
            }
          } else {
            log.info("No encryption key for SL2.0 found. End-to-End encryption is not used.");
          }

        } catch (final Exception e) {
          log.warn("No encryption key for SL2.0 found. End-to-End encryption is not used. Reason: " + e.getMessage(),
              e);

        }

        // load trusted certificates
        trustedCerts = readCertsFromKeyStore(keyStore);

        // some short validation
        if (signPrivKey == null || !(signPrivKey instanceof PrivateKey)) {
          log.info("Can NOT open privateKey for SL2.0 signing. KeyStore=" + getKeyStoreFilePath());
          throw new SL20Exception("sl20.03", new Object[] { "Can NOT open private key for signing" });

        }

        if (signCertChain == null || signCertChain.length == 0) {
          log.info("NO certificate for SL2.0 signing. KeyStore=" + getKeyStoreFilePath());
          throw new SL20Exception("sl20.03", new Object[] { "NO certificate for SL2.0 signing" });

        }

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

      } else {
        log.info("NO SL2.0 authentication security configuration. Initialization was skipped");
      }

    } catch (final RuntimeException e) {
      throw e;

    } catch (final Exception e) {
      log.error("SL2.0 security constrains initialization FAILED.", e);

    }

  }

  @Override
  public String createSignature(final String payLoad) 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);
      jws.setKey(signPrivKey);

      // TODO:
      jws.setCertificateChainHeaderValue(signCertChain);
      jws.setX509CertSha256ThumbprintHeaderValue(signCertChain[0]);

      return jws.getCompactSerialization();

    } catch (final JoseException 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 = 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, trustedCerts, 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 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);

      // 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(encCertChain[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(encCertChain[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(encPrivKey);

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

    } catch (final JoseException 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() {
    // TODO: maybe update after SL2.0 update on encryption certificate parts
    if (encCertChain != null && encCertChain.length > 0) {
      return encCertChain[0];
    } else {
      return null;
    }
  }

  private String getKeyStoreFilePath() throws EaafConfigurationException, MalformedURLException {
    return FileUtils.makeAbsoluteUrl(authConfig.getBasicConfiguration(Constants.CONFIG_PROP_SECURITY_KEYSTORE_PATH),
        authConfig.getConfigurationRootDirectory());
  }

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

    return value;

  }

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

    return value;
  }

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

    return value;
  }

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

    return value;
  }

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

    return value;
  }

  @Nonnull
  private List<X509Certificate> readCertsFromKeyStore(@Nonnull final KeyStore keyStore) throws KeyStoreException {
    final List<X509Certificate> result = new ArrayList<>();

    final Enumeration<String> aliases = keyStore.aliases();
    while (aliases.hasMoreElements()) {
      final String el = aliases.nextElement();
      log.trace("Process TrustStoreEntry: " + el);
      if (keyStore.isCertificateEntry(el)) {
        final Certificate cert = keyStore.getCertificate(el);
        if (cert != null && cert instanceof X509Certificate) {
          result.add((X509Certificate) cert);
        } else {
          log.info("Can not process entry: {}. Reason: {}", el, cert != null ? cert.getType() : "cert is null");
        }

      }
    }

    return Collections.unmodifiableList(result);
  }

}