/**
    * Licensed to the Apache Software Foundation (ASF) under one
    * or more contributor license agreements. See the NOTICE file
    * distributed with this work for additional information
    * regarding copyright ownership. The ASF licenses this file
    * to you under the Apache License, Version 2.0 (the
    * "License"); you may not use this file except in compliance
    * with the License. You may obtain a copy of the License at
    *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied. See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */
  
  package org.apache.wss4j.dom.message;
  
  import java.security.NoSuchProviderException;
  import java.security.Provider;
  import java.util.List;
  
  import javax.xml.crypto.XMLStructure;
  import javax.xml.crypto.dom.DOMStructure;
  import javax.xml.crypto.dsig.CanonicalizationMethod;
  import javax.xml.crypto.dsig.SignatureMethod;
  import javax.xml.crypto.dsig.SignedInfo;
  import javax.xml.crypto.dsig.XMLSignContext;
  import javax.xml.crypto.dsig.XMLSignature;
  import javax.xml.crypto.dsig.XMLSignatureFactory;
  import javax.xml.crypto.dsig.dom.DOMSignContext;
  import javax.xml.crypto.dsig.keyinfo.KeyInfo;
  import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
  import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
  import javax.xml.crypto.dsig.spec.ExcC14NParameterSpec;
  
  import org.w3c.dom.Document;
  import org.w3c.dom.Element;
  
  import org.apache.wss4j.common.WSEncryptionPart;
  import org.apache.wss4j.common.derivedKey.ConversationConstants;
  import org.apache.wss4j.common.ext.WSSecurityException;
  import org.apache.wss4j.common.token.Reference;
  import org.apache.wss4j.common.token.SecurityTokenReference;
  import org.apache.wss4j.common.util.KeyUtils;
  import org.apache.wss4j.common.util.XMLUtils;
  import org.apache.wss4j.dom.WSConstants;
  import org.apache.wss4j.dom.WSDocInfo;
  import org.apache.wss4j.dom.transform.STRTransform;
  import org.apache.wss4j.dom.util.WSSecurityUtil;
  
  /**
   * Builder to sign with derived keys
   */
  public class WSSecDKSign extends WSSecDerivedKeyBase {
  
      private static final org.slf4j.Logger LOG =
          org.slf4j.LoggerFactory.getLogger(WSSecDKSign.class);
  
      private String sigAlgo = WSConstants.HMAC_SHA1;
      private String digestAlgo = WSConstants.SHA1;
      private String canonAlgo = WSConstants.C14N_EXCL_OMIT_COMMENTS;
      private byte[] signatureValue;
  
      private String keyInfoUri;
      private SecurityTokenReference secRef;
      private String strUri;
      private WSDocInfo wsDocInfo;
  
      private XMLSignatureFactory signatureFactory;
      private XMLSignature sig;
      private KeyInfo keyInfo;
      private CanonicalizationMethod c14nMethod;
      private int derivedKeyLength = -1;
      private boolean addInclusivePrefixes = true;
  
      public WSSecDKSign(WSSecHeader securityHeader) {
          super(securityHeader);
          init(null);
      }
  
      public WSSecDKSign(Document doc) {
          this(doc, null);
      }
  
      public WSSecDKSign(Document doc, Provider provider) {
          super(doc);
          init(provider);
      }
  
      private void init(Provider provider) {
          if (provider == null) {
              // Try to install the Santuario Provider - fall back to the JDK provider if this does
              // not work
              try {
                  signatureFactory = XMLSignatureFactory.getInstance("DOM", "ApacheXMLDSig");
             } catch (NoSuchProviderException ex) {
                 signatureFactory = XMLSignatureFactory.getInstance("DOM");
             }
         } else {
             signatureFactory = XMLSignatureFactory.getInstance("DOM", provider);
         }
     }
 
     public Document build(byte[] ephemeralKey) throws WSSecurityException {
 
         prepare(ephemeralKey);
         if (getParts().isEmpty()) {
             getParts().add(WSSecurityUtil.getDefaultEncryptionPart(getDocument()));
         } else {
             for (WSEncryptionPart part : getParts()) {
                 if ("STRTransform".equals(part.getName()) && part.getId() == null) {
                     part.setId(strUri);
                 }
             }
         }
 
         List<javax.xml.crypto.dsig.Reference> referenceList = addReferencesToSign(getParts());
         computeSignature(referenceList);
 
         //
         // prepend elements in the right order to the security header
         //
         prependDKElementToHeader();
 
         return getDocument();
     }
 
     public void prepare(byte[] ephemeralKey) throws WSSecurityException {
         super.prepare(ephemeralKey);
         wsDocInfo = new WSDocInfo(getDocument());
         sig = null;
 
         try {
             C14NMethodParameterSpec c14nSpec = null;
             if (addInclusivePrefixes && canonAlgo.equals(WSConstants.C14N_EXCL_OMIT_COMMENTS)) {
                 Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
                 List<String> prefixes =
                     getInclusivePrefixes(securityHeaderElement, false);
                 c14nSpec = new ExcC14NParameterSpec(prefixes);
             }
 
            c14nMethod = signatureFactory.newCanonicalizationMethod(canonAlgo, c14nSpec);
         } catch (Exception ex) {
             LOG.error("", ex);
             throw new WSSecurityException(
                 WSSecurityException.ErrorCode.FAILED_SIGNATURE, ex, "noXMLSig"
             );
         }
 
         keyInfoUri = getIdAllocator().createSecureId("KI-", keyInfo);
 
         secRef = new SecurityTokenReference(getDocument());
         strUri = getIdAllocator().createSecureId("STR-", secRef);
         secRef.setID(strUri);
         if (addWSUNamespace) {
             secRef.addWSUNamespace();
         }
 
         Reference ref = new Reference(getDocument());
         ref.setURI("#" + getId());
         String ns =
             ConversationConstants.getWSCNs(getWscVersion())
             + ConversationConstants.TOKEN_TYPE_DERIVED_KEY_TOKEN;
         ref.setValueType(ns);
         secRef.setReference(ref);
 
         XMLStructure structure = new DOMStructure(secRef.getElement());
         wsDocInfo.addTokenElement(secRef.getElement(), false);
         KeyInfoFactory keyInfoFactory = signatureFactory.getKeyInfoFactory();
         keyInfo =
             keyInfoFactory.newKeyInfo(
                 java.util.Collections.singletonList(structure), keyInfoUri
             );
 
     }
 
     /**
      * Returns the SignatureElement.
      * The method can be called any time after <code>prepare()</code>.
      * @return The DOM Element of the signature.
      */
     public Element getSignatureElement() {
         Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
         return
             XMLUtils.getDirectChildElement(
                 securityHeaderElement, WSConstants.SIG_LN, WSConstants.SIG_NS
             );
     }
 
     /**
      * This method adds references to the Signature.
      *
      * @param references The list of references to sign
      * @throws WSSecurityException
      */
     public List<javax.xml.crypto.dsig.Reference> addReferencesToSign(
         List<WSEncryptionPart> references
     ) throws WSSecurityException {
         return
             addReferencesToSign(
                 getDocument(),
                 references,
                 wsDocInfo,
                 signatureFactory,
                 addInclusivePrefixes,
                 digestAlgo
             );
     }
 
     /**
      * Compute the Signature over the references.
      *
      * After references are set this method computes the Signature for them.
      * This method can be called any time after the references were set. See
      * <code>addReferencesToSign()</code>.
      *
      * @throws WSSecurityException
      */
     public void computeSignature(
         List<javax.xml.crypto.dsig.Reference> referenceList
     ) throws WSSecurityException {
         computeSignature(referenceList, true, null);
     }
 
     /**
      * Compute the Signature over the references.
      *
      * After references are set this method computes the Signature for them.
      * This method can be called any time after the references were set. See
      * <code>addReferencesToSign()</code>.
      *
      * @throws WSSecurityException
      */
     public void computeSignature(
         List<javax.xml.crypto.dsig.Reference> referenceList,
         boolean prepend,
         Element siblingElement
     ) throws WSSecurityException {
         try {
             java.security.Key key = getDerivedKey(sigAlgo);
             SignatureMethod signatureMethod =
                 signatureFactory.newSignatureMethod(sigAlgo, null);
             SignedInfo signedInfo =
                 signatureFactory.newSignedInfo(c14nMethod, signatureMethod, referenceList);
 
             sig = signatureFactory.newXMLSignature(
                     signedInfo,
                     keyInfo,
                     null,
                     getIdAllocator().createId("SIG-", null),
                     null);
 
             //
             // Figure out where to insert the signature element
             //
             XMLSignContext signContext = null;
             Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
             if (prepend) {
                 if (siblingElement == null) {
                     siblingElement = (Element)securityHeaderElement.getFirstChild();
                 }
                 if (siblingElement == null) {
                     signContext = new DOMSignContext(key, securityHeaderElement);
                 } else {
                     signContext = new DOMSignContext(key, securityHeaderElement, siblingElement);
                 }
             } else {
                 signContext = new DOMSignContext(key, securityHeaderElement);
             }
 
             signContext.putNamespacePrefix(WSConstants.SIG_NS, WSConstants.SIG_PREFIX);
             if (WSConstants.C14N_EXCL_OMIT_COMMENTS.equals(canonAlgo)) {
                 signContext.putNamespacePrefix(
                     WSConstants.C14N_EXCL_OMIT_COMMENTS,
                     WSConstants.C14N_EXCL_OMIT_COMMENTS_PREFIX
                 );
             }
             signContext.setProperty(STRTransform.TRANSFORM_WS_DOC_INFO, wsDocInfo);
             wsDocInfo.setCallbackLookup(callbackLookup);
 
             // Add the elements to sign to the Signature Context
             wsDocInfo.setTokensOnContext((DOMSignContext)signContext);
 
             sig.sign(signContext);
 
             signatureValue = sig.getSignatureValue().getValue();
         } catch (Exception ex) {
             LOG.error(ex.getMessage(), ex);
             throw new WSSecurityException(
                 WSSecurityException.ErrorCode.FAILED_SIGNATURE, ex
             );
         }
     }
 
     protected int getDerivedKeyLength() throws WSSecurityException {
         return derivedKeyLength > 0 ? derivedKeyLength : KeyUtils.getKeyLength(sigAlgo);
     }
 
     public void setDerivedKeyLength(int keyLength) {
         derivedKeyLength = keyLength;
     }
 
     /**
      * Set the signature algorithm to use. The default is WSConstants.SHA1.
      * @param algorithm the signature algorithm to use.
      */
     public void setSignatureAlgorithm(String algorithm) {
         sigAlgo = algorithm;
     }
 
     /**
      * @return the signature algorithm to use
      */
     public String getSignatureAlgorithm() {
         return sigAlgo;
     }
 
     /**
      * Returns the the value of wsu:Id attribute of the Signature element.
      *
      * @return Return the wsu:Id of this token or null if the signature has not been generated.
      */
     public String getSignatureId() {
         if (sig == null) {
             return null;
         }
         return sig.getId();
     }
 
     /**
      * Set the digest algorithm to use. The default is WSConstants.SHA1.
      * @param algorithm the digest algorithm to use.
      */
     public void setDigestAlgorithm(String algorithm) {
         digestAlgo = algorithm;
     }
 
     /**
      * @return the digest algorithm to use
      */
     public String getDigestAlgorithm() {
         return digestAlgo;
     }
 
     /**
      * @return Returns the signatureValue.
      */
     public byte[] getSignatureValue() {
         return signatureValue;
     }
 
     /**
      * Set the canonicalization method to use.
      *
      * If the canonicalization method is not set then the recommended Exclusive
      * XML Canonicalization is used by default Refer to WSConstants which
      * algorithms are supported.
      *
      * @param algo Is the name of the signature algorithm
      * @see WSConstants#C14N_OMIT_COMMENTS
      * @see WSConstants#C14N_WITH_COMMENTS
      * @see WSConstants#C14N_EXCL_OMIT_COMMENTS
      * @see WSConstants#C14N_EXCL_WITH_COMMENTS
      */
     public void setSigCanonicalization(String algo) {
         canonAlgo = algo;
     }
 
     /**
      * Get the canonicalization method.
      *
      * If the canonicalization method was not set then Exclusive XML
      * Canonicalization is used by default.
      *
      * @return The string describing the canonicalization algorithm.
      */
     public String getSigCanonicalization() {
         return canonAlgo;
     }
 
     public boolean isAddInclusivePrefixes() {
         return addInclusivePrefixes;
     }
 
     public void setAddInclusivePrefixes(boolean addInclusivePrefixes) {
         this.addInclusivePrefixes = addInclusivePrefixes;
     }
 }