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.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.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.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<X509Certificate>();
	
	private static JsonMapper mapper = new JsonMapper();
	
	@PostConstruct
	protected void initalize() {
		log.info("Initialize SL2.0 authentication security constrains ... ");
		try {
			if (getKeyStoreFilePath() != null) { 			
				KeyStore keyStore = KeyStoreUtils.loadKeyStore(getKeyStoreFilePath(), 
						getKeyStorePassword());
				
				//load signing key
				signPrivKey = keyStore.getKey(getSigningKeyAlias(), getSigningKeyPassword().toCharArray());
				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) {
						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 (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
				Enumeration<String> aliases = keyStore.aliases();
				while(aliases.hasMoreElements()) {
					String el = aliases.nextElement();
					log.trace("Process TrustStoreEntry: " + el);
					if (keyStore.isCertificateEntry(el)) {
						Certificate cert = keyStore.getCertificate(el); 
						if (cert != null && cert instanceof X509Certificate)
							trustedCerts.add((X509Certificate) cert);
						else
							log.info("Can not process entry: " + el + ". Reason: " + cert.toString());
						
					}
				}
	
				//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 ( Exception e) {
			log.error("SL2.0 security constrains initialization FAILED.", e);
			
		}
		
	}
	
	
	@Override
	public String createSignature(String payLoad) throws SLCommandoBuildException {
		try {
			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 (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(String serializedContent) throws SL20Exception {
		try {
			JsonWebSignature jws = new JsonWebSignature();
			//set payload
			jws.setCompactSerialization(serializedContent);
				
			//set security constrains
			jws.setAlgorithmConstraints(new AlgorithmConstraints(ConstraintType.WHITELIST,   
					SL20Constants.SL20_ALGORITHM_WHITELIST_SIGNING.toArray(new String[SL20Constants.SL20_ALGORITHM_WHITELIST_SIGNING.size()])));
			
			//load signinc certs
			Key selectedKey = null;
			List<X509Certificate> x5cCerts = jws.getCertificateChainHeaderValue();
			String x5t256 = jws.getX509CertSha256ThumbprintHeaderValue();
			if (x5cCerts != null) {
				log.debug("Found x509 certificate in JOSE header ... ");
			log.trace("Sorting received X509 certificates ... ");
				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.");
					log.debug("JOSE certificate: " + sortedX5cCerts.get(0).toString());
					try {
						log.debug("Cert: " + Base64Utils.encodeToString(sortedX5cCerts.get(0).getEncoded()));
					} catch (CertificateEncodingException e) {
						e.printStackTrace();
					}
					
				}
				
			} else if (StringUtils.isNotEmpty(x5t256)) {
				log.debug("Found x5t256 fingerprint in JOSE header .... ");
				X509VerificationKeyResolver x509VerificationKeyResolver = new X509VerificationKeyResolver(trustedCerts);
				selectedKey = x509VerificationKeyResolver.resolveKey(jws, Collections.<JsonWebStructure>emptyList());
				
			} 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");
				
			}
					
			if (selectedKey == null) {
				log.info("Can NOT select verification key for JWS. Signature verification FAILED.");
				throw new SLCommandoParserException("Can NOT select verification key for JWS. Signature verification FAILED");
				
			}
			
			//set verification key
			jws.setKey(selectedKey);
			
			//validate signature
			boolean valid = jws.verifySignature();
			if (!valid) {
				log.info("JWS signature invalide. Stopping authentication process ...");
				log.debug("Received JWS msg: " + serializedContent);
				throw new SL20SecurityException("JWS signature invalide.");
				
			}
			
			
			//load payLoad
			log.debug("SL2.0 commando signature validation sucessfull");
			JsonNode sl20Req = mapper.getMapper().readTree(jws.getPayload());
			
			return new VerificationResult(sl20Req, null, valid) ;
			
		} catch (JoseException | JsonParseException e) {
			log.warn("SL2.0 commando signature validation FAILED", e);
			throw new SL20SecurityException(new Object[]{e.getMessage()}, e);
			
		} catch (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(String compactSerialization) throws SL20Exception {
		try {			
			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
			List<X509Certificate> x5cCerts = receiverJwe.getCertificateChainHeaderValue();
			String x5t256 = receiverJwe.getX509CertSha256ThumbprintHeaderValue();
			if (x5cCerts != null) {
				log.debug("Found x509 certificate in JOSE header ... ");
				log.trace("Sorting received X509 certificates ... ");
				List<X509Certificate> sortedX5cCerts  = X509Utils.sortCertificates(x5cCerts);
		
				if (!sortedX5cCerts.get(0).equals(encCertChain[0])) {
					log.info("Certificate from JOSE header does NOT match encryption certificate");
					log.debug("JOSE certificate: " + sortedX5cCerts.get(0).toString());
					
					try {
						log.debug("Cert: " + Base64Utils.encode(sortedX5cCerts.get(0).getEncoded()));
					} catch (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 .... ");
				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 (JoseException e) {
			log.warn("SL2.0 result decryption FAILED", e);
			throw new SL20SecurityException(new Object[]{e.getMessage()}, e);
			
		} catch ( 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 (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;
	}
	
}