/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package de.christofreichardt.jaxws.securitypriceservice;

import de.christofreichardt.diagnosis.AbstractTracer;
import de.christofreichardt.diagnosis.LogLevel;
import de.christofreichardt.diagnosis.Traceable;
import de.christofreichardt.diagnosis.TracerFactory;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.namespace.QName;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPHeader;
import javax.xml.soap.SOAPMessage;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;
import org.w3c.dom.Element;

/**
 *
 * @author student
 */
public class SignatureHandler implements SOAPHandler<SOAPMessageContext>, Traceable {
  static public final String STORE_PASS = "changeit";
  static public final String KEY_PASS = "changeit";
  static public final String KEYSTORE_PATH = ".." + File.separator + "mykeystore.jks";
  
  final private KeyStore keyStore;

  public SignatureHandler() throws GeneralSecurityException, IOException {
    this.keyStore = KeyStore.getInstance("JKS");
    try (FileInputStream inputStream = new FileInputStream(KEYSTORE_PATH)) {
      this.keyStore.load(inputStream, STORE_PASS.toCharArray());
    }
  }

  public SignatureHandler(File keystoreFile) throws GeneralSecurityException, IOException {
    this.keyStore = KeyStore.getInstance("JKS");
    try (FileInputStream inputStream = new FileInputStream(keystoreFile)) {
      this.keyStore.load(inputStream, STORE_PASS.toCharArray());
    }
  }
  
  private KeyStore.PrivateKeyEntry retrieveKey(String alias) throws GeneralSecurityException {
    AbstractTracer tracer = getCurrentTracer();
    tracer.entry("KeyStore.PrivateKeyEntry", this, "retrieveKey(File keyStoreFile)");
    
    try {
      KeyStore.PasswordProtection passwordProtection = new KeyStore.PasswordProtection(KEY_PASS.toCharArray());
      KeyStore.Entry keyEntry = this.keyStore.getEntry(alias, passwordProtection);
      
      return (KeyStore.PrivateKeyEntry) keyEntry;
    }
    finally {
      tracer.wayout();
    }
  }
  
  @Override
  public boolean handleMessage(SOAPMessageContext soapMessageContext) {
    AbstractTracer tracer = getCurrentTracer();
    tracer.entry("boolean", this, "handleMessage(SOAPMessageContext messageContext)");
    
    try {
      tracer.out().printfIndentln("%s = %b", SOAPMessageContext.MESSAGE_OUTBOUND_PROPERTY, soapMessageContext.get(SOAPMessageContext.MESSAGE_OUTBOUND_PROPERTY));

      boolean outbound = (boolean) soapMessageContext.get(SOAPMessageContext.MESSAGE_OUTBOUND_PROPERTY);
      if (outbound) {
        SOAPMessage soapMessage = soapMessageContext.getMessage();
        try {
          SOAPHeader soapHeader = soapMessage.getSOAPHeader();
          SOAPBody soapBody = soapMessage.getSOAPBody();
          String userId = soapBody.getElementsByTagNameNS("http://de.christofreichardt/SecurityPriceService", "UserId").item(0).getTextContent();
          addSignature(soapHeader, userId);
        }
        catch (SOAPException | GeneralSecurityException | MarshalException | XMLSignatureException ex) {
          tracer.logException(LogLevel.ERROR, ex, getClass(), "handleMessage(SOAPMessageContext soapMessageContext)");
        }
      }
      
      return true;
    }
    finally {
      tracer.wayout();
    }
  }
  
  private void traceReferences(List<?> references) {
    AbstractTracer tracer = getCurrentTracer();
    tracer.entry("void", this, "traceReferences(List<?> references)");
    
    try {
      try {
        Iterator iter = references.iterator();
        while (iter.hasNext()) {
          Reference reference = (Reference) iter.next();
          try (InputStream inputStream = reference.getDigestInputStream();
              InputStreamReader inputStreamReader = new InputStreamReader(inputStream);
              BufferedReader bufferedReader = new BufferedReader(inputStreamReader)) {
            do {
              String line = bufferedReader.readLine();
              if (line == null) {
                break;
              }
              tracer.out().printfIndentln("%s", line);
            } while (true);
          }
        }
      }
      catch (IOException ex) {
        tracer.logException(LogLevel.ERROR, ex, getClass(), "traceReference(List<?> references)");
      }
    }
    finally {
      tracer.wayout();
    }
  }
  
  private void addSignature(Element soapHeader, String userId) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, MarshalException, XMLSignatureException, GeneralSecurityException {
    AbstractTracer tracer = getCurrentTracer();
    tracer.entry("void", this, "addSignature(Element soapHeader, String userId)");

    try {
      XMLSignatureFactory xmlSignatureFactory = XMLSignatureFactory.getInstance();
      DigestMethod digestMethod = xmlSignatureFactory.newDigestMethod(DigestMethod.SHA256, null);
      Reference reference = xmlSignatureFactory.newReference("#xpointer(/S:Envelope/S:Body)", digestMethod);
      CanonicalizationMethod canonicalizationMethod = xmlSignatureFactory.newCanonicalizationMethod(CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS, (C14NMethodParameterSpec) null);
      SignatureMethod signatureMethod = xmlSignatureFactory.newSignatureMethod("http://www.w3.org/2001/04/xmldsig-more#rsa-sha256", null);
      SignedInfo signedInfo = xmlSignatureFactory.newSignedInfo(canonicalizationMethod, signatureMethod, Collections.singletonList(reference));
      XMLSignature xmlSignature = xmlSignatureFactory.newXMLSignature(signedInfo, null);
      DOMSignContext domSignContext = new DOMSignContext(retrieveKey(userId).getPrivateKey(), soapHeader);
      domSignContext.setURIDereferencer(new SoapBodyDereferencer());
      domSignContext.setProperty("javax.xml.crypto.dsig.cacheReference", Boolean.TRUE);
      xmlSignature.sign(domSignContext);
      List<?> references = xmlSignature.getSignedInfo().getReferences();
      traceReferences(references);
    }
    finally {
      tracer.wayout();
    }
  }
  
  @Override
  public Set<QName> getHeaders() {
    return Collections.emptySet();
  }
  
  @Override
  public boolean handleFault(SOAPMessageContext messageContext) {
    return true;
  }
  
  @Override
  public void close(MessageContext context) {
  }

  @Override
  public AbstractTracer getCurrentTracer() {
    return TracerFactory.getInstance().getCurrentPoolTracer();
  }
  
}
