/**
 * 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.saml;

import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.List;

import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignContext;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.ExcC14NParameterSpec;

import org.apache.wss4j.common.SignatureActionToken;
import org.apache.wss4j.common.WSEncryptionPart;
import org.apache.wss4j.common.crypto.Crypto;
import org.apache.wss4j.common.crypto.CryptoType;
import org.apache.wss4j.common.ext.WSSecurityException;
import org.apache.wss4j.common.saml.OpenSAMLUtil;
import org.apache.wss4j.common.saml.SAMLKeyInfo;
import org.apache.wss4j.common.saml.SAMLUtil;
import org.apache.wss4j.common.saml.SamlAssertionWrapper;
import org.apache.wss4j.common.token.BinarySecurity;
import org.apache.wss4j.common.token.DOMX509Data;
import org.apache.wss4j.common.token.DOMX509IssuerSerial;
import org.apache.wss4j.common.token.Reference;
import org.apache.wss4j.common.token.SecurityTokenReference;
import org.apache.wss4j.common.token.X509Security;
import org.apache.wss4j.common.util.KeyUtils;
import org.apache.wss4j.dom.WSConstants;
import org.apache.wss4j.dom.WSDocInfo;
import org.apache.wss4j.dom.handler.RequestData;
import org.apache.wss4j.dom.message.WSSecHeader;
import org.apache.wss4j.dom.message.WSSecSignature;
import org.apache.wss4j.dom.transform.STRTransform;
import org.apache.wss4j.dom.util.WSSecurityUtil;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

public class WSSecSignatureSAML extends WSSecSignature {

    private static final org.slf4j.Logger LOG =
        org.slf4j.LoggerFactory.getLogger(WSSecSignatureSAML.class);
    private boolean senderVouches;
    private SecurityTokenReference secRefSaml;
    private String secRefID;
    private Element samlToken;
    private Crypto userCrypto;
    private Crypto issuerCrypto;
    private String issuerKeyName;
    private String issuerKeyPW;
    private boolean useDirectReferenceToAssertion;

    /**
     * Constructor.
     */
    public WSSecSignatureSAML(WSSecHeader securityHeader) {
        super(securityHeader);
    }

    public WSSecSignatureSAML(Document doc) {
        super(doc);
    }

    /**
     * Builds a signed soap envelope with SAML token.
     *
     * The method first gets an appropriate security header. According to the
     * defined parameters for certificate handling the signature elements are
     * constructed and inserted into the <code>wsse:Signature</code>
     *
     * @param uCrypto
     *            The user's Crypto instance
     * @param samlAssertion
     *            the complete SAML assertion
     * @param iCrypto
     *            An instance of the Crypto API to handle keystore SAML token
     *            issuer and to generate certificates
     * @param iKeyName
     *            Private key to use in case of "sender-Vouches"
     * @param iKeyPW
     *            Password for issuer private key
     * @return A signed SOAP envelope as <code>Document</code>
     * @throws WSSecurityException
     */
    public Document build(
        Crypto uCrypto, SamlAssertionWrapper samlAssertion,
        Crypto iCrypto, String iKeyName, String iKeyPW
    ) throws WSSecurityException {

        prepare(uCrypto, samlAssertion, iCrypto, iKeyName, iKeyPW);

        if (getParts().isEmpty()) {
            getParts().add(WSSecurityUtil.getDefaultEncryptionPart(getDocument()));
        } else {
            for (WSEncryptionPart part : getParts()) {
                if ("STRTransform".equals(part.getName()) && part.getId() == null) {
                    part.setId(strUri);
                }
            }
        }

        //
        // Add the STRTransform for the SecurityTokenReference to the SAML assertion
        // if it exists
        //
        if (secRefID != null) {
            String soapNamespace =
                WSSecurityUtil.getSOAPNamespace(getDocument().getDocumentElement());
            WSEncryptionPart encP =
                new WSEncryptionPart("STRTransform", soapNamespace, "Content");
            encP.setId(secRefID);
            getParts().add(encP);
        }

        List<javax.xml.crypto.dsig.Reference> referenceList = addReferencesToSign(getParts());

        prependSAMLElementsToHeader();

        if (senderVouches) {
            computeSignature(referenceList, secRefSaml.getElement());
        } else {
            computeSignature(referenceList, samlToken);
        }

        //
        // if we have a BST prepend it in front of the Signature according to
        // strict layout rules.
        //
        if (bstToken != null) {
            prependBSTElementToHeader();
        }

        return getDocument();
    }

    /**
     * Initialize a WSSec SAML Signature.
     *
     * The method sets up and initializes a WSSec SAML Signature structure after
     * the relevant information was set. After setup of the references to
     * elements to sign may be added. After all references are added they can be
     * signed.
     *
     * This method does not add the Signature element to the security header.
     * See <code>prependSignatureElementToHeader()</code> method.
     *
     * @param uCrypto
     *            The user's Crypto instance
     * @param samlAssertion
     *            the complete SAML assertion
     * @param iCrypto
     *            An instance of the Crypto API to handle keystore SAML token
     *            issuer and to generate certificates
     * @param iKeyName
     *            Private key to use in case of "sender-Vouches"
     * @param iKeyPW
     *            Password for issuer private key
     * @throws WSSecurityException
     */
    public void prepare(
        Crypto uCrypto, SamlAssertionWrapper samlAssertion, Crypto iCrypto,
        String iKeyName, String iKeyPW
    ) throws WSSecurityException {

        LOG.debug("Beginning ST signing...");

        userCrypto = uCrypto;
        issuerCrypto = iCrypto;
        issuerKeyName = iKeyName;
        issuerKeyPW = iKeyPW;

        samlToken = samlAssertion.toDOM(getDocument());

        //
        // Get some information about the SAML token content. This controls how
        // to deal with the whole stuff. First get the Authentication statement
        // (includes Subject), then get the _first_ confirmation method only
        // thats if "senderVouches" is true.
        //
        String confirmMethod = null;
        List<String> methods = samlAssertion.getConfirmationMethods();
        if (methods != null && !methods.isEmpty()) {
            confirmMethod = methods.get(0);
        }
        if (OpenSAMLUtil.isMethodSenderVouches(confirmMethod)) {
            senderVouches = true;
        }
        //
        // Gather some info about the document to process and store it for
        // retrieval
        //
        if (super.getWsDocInfo() == null) {
            WSDocInfo wsDocInfo = new WSDocInfo(getDocument());
            super.setWsDocInfo(wsDocInfo);
        }


        X509Certificate[] certs = null;
        PublicKey publicKey = null;

        if (senderVouches) {
            CryptoType cryptoType = new CryptoType(CryptoType.TYPE.ALIAS);
            cryptoType.setAlias(issuerKeyName);
            certs = issuerCrypto.getX509Certificates(cryptoType);
            getWsDocInfo().setCrypto(issuerCrypto);
        } else {
            //
            // in case of key holder: - get the user's certificate that _must_ be
            // included in the SAML token. To ensure the cert integrity the SAML
            // token must be signed (by the issuer).
            //
            if (userCrypto == null || !samlAssertion.isSigned()) {
                throw new WSSecurityException(
                    WSSecurityException.ErrorCode.FAILURE,
                    "invalidSAMLsecurity",
                    new Object[] {"for SAML Signature (Key Holder)"});
            }
            if (secretKey == null) {
                RequestData data = new RequestData();
                data.setWsDocInfo(getWsDocInfo());
                SignatureActionToken actionToken = new SignatureActionToken();
                data.setSignatureToken(actionToken);
                actionToken.setCrypto(userCrypto);
                SAMLKeyInfo samlKeyInfo =
                    SAMLUtil.getCredentialFromSubject(
                            samlAssertion, new WSSSAMLKeyInfoProcessor(data),
                            userCrypto, data.getCallbackHandler()
                    );
                if (samlKeyInfo != null) {
                    publicKey = samlKeyInfo.getPublicKey();
                    certs = samlKeyInfo.getCerts();
                    getWsDocInfo().setCrypto(userCrypto);
                }
            }
        }
        if ((certs == null || certs.length == 0 || certs[0] == null)
            && publicKey == null && secretKey == null) {
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE,
                "noCertsFound",
                new Object[] {"SAML signature"});
        }

        if (getSignatureAlgorithm() == null) {
            PublicKey key = null;
            if (certs != null && certs[0] != null) {
                key = certs[0].getPublicKey();
            } else if (publicKey != null) {
                key = publicKey;
            } else {
                throw new WSSecurityException(
                    WSSecurityException.ErrorCode.FAILURE, "unknownSignatureAlgorithm"
                );
            }

            String pubKeyAlgo = key.getAlgorithm();
            LOG.debug("automatic sig algo detection: {}", pubKeyAlgo);
            if (pubKeyAlgo.equalsIgnoreCase("DSA")) {
                setSignatureAlgorithm(WSConstants.DSA);
            } else if (pubKeyAlgo.equalsIgnoreCase("RSA")) {
                setSignatureAlgorithm(WSConstants.RSA);
            } else {
                throw new WSSecurityException(
                    WSSecurityException.ErrorCode.FAILURE,
                    "unknownSignatureAlgorithm",
                    new Object[] {pubKeyAlgo});
            }
        }
        sig = null;

        try {
            C14NMethodParameterSpec c14nSpec = null;
            if (isAddInclusivePrefixes() && getSigCanonicalization().equals(WSConstants.C14N_EXCL_OMIT_COMMENTS)) {
                Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
                List<String> prefixes = getInclusivePrefixes(securityHeaderElement, false);
                c14nSpec = new ExcC14NParameterSpec(prefixes);
            }

           c14nMethod =
               signatureFactory.newCanonicalizationMethod(getSigCanonicalization(), c14nSpec);
        } catch (Exception ex) {
            LOG.error("", ex);
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILED_SIGNATURE, ex, "noXMLSig"
            );
        }

        keyInfoUri = getIdAllocator().createSecureId("KeyId-", keyInfo);
        SecurityTokenReference secRef = new SecurityTokenReference(getDocument());
        strUri = getIdAllocator().createSecureId("STRId-", secRef);
        secRef.setID(strUri);
        setSecurityTokenReference(secRef);

        if (certs != null && certs.length != 0) {
            certUri = getIdAllocator().createSecureId("CertId-", certs[0]);
        }

        //
        // If the sender vouches, then we must sign the SAML token _and_ at
        // least one part of the message (usually the SOAP body). To do so we
        // need to - put in a reference to the SAML token. Thus we create a STR
        // and insert it into the wsse:Security header - set a reference of the
        // created STR to the signature and use STR Transform during the
        // signature
        //
        try {
            if (senderVouches) {
                secRefSaml = new SecurityTokenReference(getDocument());
                secRefID = getIdAllocator().createSecureId("STRSAMLId-", secRefSaml);
                secRefSaml.setID(secRefID);

                if (useDirectReferenceToAssertion) {
                    Reference ref = new Reference(getDocument());
                    ref.setURI("#" + samlAssertion.getId());
                    if (samlAssertion.getSaml1() != null) {
                        ref.setValueType(WSConstants.WSS_SAML_KI_VALUE_TYPE);
                        secRefSaml.addTokenType(WSConstants.WSS_SAML_TOKEN_TYPE);
                    } else if (samlAssertion.getSaml2() != null) {
                        secRefSaml.addTokenType(WSConstants.WSS_SAML2_TOKEN_TYPE);
                    }
                    secRefSaml.setReference(ref);
                } else {
                    Element keyId = getDocument().createElementNS(WSConstants.WSSE_NS, "wsse:KeyIdentifier");
                    String valueType = null;
                    if (samlAssertion.getSaml1() != null) {
                        valueType = WSConstants.WSS_SAML_KI_VALUE_TYPE;
                        secRefSaml.addTokenType(WSConstants.WSS_SAML_TOKEN_TYPE);
                    } else if (samlAssertion.getSaml2() != null) {
                        valueType = WSConstants.WSS_SAML2_KI_VALUE_TYPE;
                        secRefSaml.addTokenType(WSConstants.WSS_SAML2_TOKEN_TYPE);
                    }
                    keyId.setAttributeNS(
                        null, "ValueType", valueType
                    );
                    keyId.appendChild(getDocument().createTextNode(samlAssertion.getId()));
                    Element elem = secRefSaml.getElement();
                    elem.appendChild(keyId);
                }
                getWsDocInfo().addTokenElement(secRefSaml.getElement(), false);
            }
        } catch (Exception ex) {
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILED_SIGNATURE, ex, "noXMLSig"
            );
        }

        X509Certificate cert = certs != null ? certs[0] : null;
        configureKeyInfo(secRef, cert, iCrypto != null ? iCrypto : uCrypto, samlAssertion);

        getWsDocInfo().addTokenElement(samlToken, false);
    }

    private void configureKeyInfo(
        SecurityTokenReference secRef, X509Certificate cert,
        Crypto crypto, SamlAssertionWrapper samlAssertion
    ) throws WSSecurityException {
        if (getCustomKeyInfoElement() == null) {
            if (senderVouches) {
                switch (keyIdentifierType) {
                case WSConstants.BST_DIRECT_REFERENCE:
                    Reference ref = new Reference(getDocument());
                    ref.setURI("#" + certUri);
                    BinarySecurity binarySecurity = new X509Security(getDocument());
                    ((X509Security) binarySecurity).setX509Certificate(cert);
                    binarySecurity.setID(certUri);
                    bstToken = binarySecurity.getElement();
                    getWsDocInfo().addTokenElement(bstToken, false);
                    ref.setValueType(binarySecurity.getValueType());
                    secRef.setReference(ref);
                    break;

                case WSConstants.X509_KEY_IDENTIFIER :
                    secRef.setKeyIdentifier(cert);
                    break;

                case WSConstants.SKI_KEY_IDENTIFIER:
                    secRef.setKeyIdentifierSKI(cert, crypto);
                    break;

                case WSConstants.THUMBPRINT_IDENTIFIER:
                    secRef.setKeyIdentifierThumb(cert);
                    break;

                    case WSConstants.ISSUER_SERIAL:
                        addIssuerSerial(cert, secRef, false);
                        break;
                    case WSConstants.ISSUER_SERIAL_QUOTE_FORMAT:
                        addIssuerSerial(cert, secRef, true);
                        break;

                default:
                    throw new WSSecurityException(
                        WSSecurityException.ErrorCode.FAILURE, "unsupportedKeyId"
                    );
                }
            } else if (useDirectReferenceToAssertion) {
                Reference ref = new Reference(getDocument());
                ref.setURI("#" + samlAssertion.getId());
                if (samlAssertion.getSaml1() != null) {
                    ref.setValueType(WSConstants.WSS_SAML_KI_VALUE_TYPE);
                    secRef.addTokenType(WSConstants.WSS_SAML_TOKEN_TYPE);
                } else if (samlAssertion.getSaml2() != null) {
                    secRef.addTokenType(WSConstants.WSS_SAML2_TOKEN_TYPE);
                }
                secRef.setReference(ref);
            } else {
                Element keyId = getDocument().createElementNS(WSConstants.WSSE_NS, "wsse:KeyIdentifier");
                String valueType = null;
                if (samlAssertion.getSaml1() != null) {
                    valueType = WSConstants.WSS_SAML_KI_VALUE_TYPE;
                    secRef.addTokenType(WSConstants.WSS_SAML_TOKEN_TYPE);
                } else if (samlAssertion.getSaml2() != null) {
                    valueType = WSConstants.WSS_SAML2_KI_VALUE_TYPE;
                    secRef.addTokenType(WSConstants.WSS_SAML2_TOKEN_TYPE);
                }
                keyId.setAttributeNS(
                    null, "ValueType", valueType
                );
                keyId.appendChild(getDocument().createTextNode(samlAssertion.getId()));
                Element elem = secRef.getElement();
                elem.appendChild(keyId);
            }
        }

        marshalKeyInfo(getWsDocInfo());
    }

    private void addIssuerSerial(X509Certificate cert, SecurityTokenReference secRef, boolean isQuoteDelimited) {
        final String issuer = cert.getIssuerDN().getName();
        final java.math.BigInteger serialNumber = cert.getSerialNumber();
        final DOMX509IssuerSerial domIssuerSerial =
                new DOMX509IssuerSerial(getDocument(), issuer, serialNumber, isQuoteDelimited);
        final DOMX509Data domX509Data = new DOMX509Data(getDocument(), domIssuerSerial);
        secRef.setUnknownElement(domX509Data.getElement());
    }

    /**
     * Prepend the SAML elements to the elements already in the Security header.
     *
     * The method can be called any time after <code>prepare()</code>. This
     * allows to insert the SAML elements at any position in the Security
     * header.
     *
     * This methods first prepends the SAML security reference if mode is
     * <code>senderVouches</code>, then the SAML token itself,
     */
    public void prependSAMLElementsToHeader() {
        Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
        if (senderVouches) {
            WSSecurityUtil.prependChildElement(securityHeaderElement, secRefSaml.getElement());
        }

        WSSecurityUtil.prependChildElement(securityHeaderElement, samlToken);
    }


    /**
     * 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,
        Element siblingElement
    ) throws WSSecurityException {
        try {
            java.security.Key key;
            if (senderVouches) {
                key = issuerCrypto.getPrivateKey(issuerKeyName, issuerKeyPW);
            } else if (secretKey != null) {
                key = KeyUtils.prepareSecretKey(getSignatureAlgorithm(), secretKey);
            } else {
                key = userCrypto.getPrivateKey(user, password);
            }
            SignatureMethod signatureMethod =
                signatureFactory.newSignatureMethod(getSignatureAlgorithm(), null);
            SignedInfo signedInfo =
                signatureFactory.newSignedInfo(c14nMethod, signatureMethod, referenceList);

            sig = signatureFactory.newXMLSignature(
                    signedInfo,
                    keyInfo,
                    null,
                    getIdAllocator().createId("SIG-", null),
                    null);

            Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
            //
            // Prepend the signature element to the security header (after the assertion)
            //
            XMLSignContext signContext = null;
            if (siblingElement != null && siblingElement.getNextSibling() != null) {
                signContext =
                    new DOMSignContext(key, securityHeaderElement, siblingElement.getNextSibling());
            } else {
                signContext = new DOMSignContext(key, securityHeaderElement);
            }
            if (getSignatureProvider() != null) {
                signContext.setProperty("org.jcp.xml.dsig.internal.dom.SignatureProvider", getSignatureProvider());
            }

            signContext.putNamespacePrefix(WSConstants.SIG_NS, WSConstants.SIG_PREFIX);
            if (WSConstants.C14N_EXCL_OMIT_COMMENTS.equals(getSigCanonicalization())) {
                signContext.putNamespacePrefix(
                    WSConstants.C14N_EXCL_OMIT_COMMENTS,
                    WSConstants.C14N_EXCL_OMIT_COMMENTS_PREFIX
                );
            }
            signContext.setProperty(STRTransform.TRANSFORM_WS_DOC_INFO, getWsDocInfo());
            getWsDocInfo().setCallbackLookup(callbackLookup);

            // Add the elements to sign to the Signature Context
            getWsDocInfo().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
            );
        }
    }

    /**
     * Return whether a Direct Reference is to be used to reference the assertion. The
     * default is false.
     * @return whether a Direct Reference is to be used to reference the assertion
     */
    public boolean isUseDirectReferenceToAssertion() {
        return useDirectReferenceToAssertion;
    }

    /**
     * Set whether a Direct Reference is to be used to reference the assertion. The
     * default is false.
     * @param useDirectReferenceToAssertion whether a Direct Reference is to be used
     *        to reference the assertion
     */
    public void setUseDirectReferenceToAssertion(boolean useDirectReferenceToAssertion) {
        this.useDirectReferenceToAssertion = useDirectReferenceToAssertion;
    }

}