/*
 *  Copyright 2001-2005 Internet2
 * 
 * Licensed 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.
 */

/* SAMLSignedObject.cpp - base class for all signed SAML constructs

   Scott Cantor
   8/12/02

   $History:$
*/

#include "internal.h"

#include <xsec/enc/XSECCryptoProvider.hpp>
#include <xsec/enc/XSECKeyInfoResolverDefault.hpp>
#include <xsec/enc/XSECCryptoException.hpp>
#include <xsec/framework/XSECException.hpp>
#include <xsec/dsig/DSIGKeyInfo.hpp>
#include <xsec/dsig/DSIGKeyInfoX509.hpp>
#include <xsec/dsig/DSIGTransformC14n.hpp>
#include <xsec/dsig/DSIGReference.hpp>
#include <xsec/dsig/DSIGTransformList.hpp>

using namespace saml;
using namespace std;

struct DefaultKeyInfoResolver : public KeyInfoResolver
{
    DefaultKeyInfoResolver() {
        m_resolver=new XSECKeyInfoResolverDefault();
    }

    XSECCryptoX509* resolveCert(DSIGKeyInfoList* klist) {
    	for (unsigned int i=0; klist && i<klist->getSize(); i++) {
            DSIGKeyInfo* ki=klist->item(i);
            if (ki && ki->getKeyInfoType()==DSIGKeyInfo::KEYINFO_X509) {
                DSIGKeyInfoX509* kix509=static_cast<DSIGKeyInfoX509*>(ki);
                if (kix509->getCertificateListSize())
                    return kix509->getCertificateCryptoItem(0);
            }
        }
        return NULL;
    }
};

KeyInfoResolver* KeyInfoResolver::getInstance(const DOMElement* e)
{
    return new DefaultKeyInfoResolver();
}

KeyInfoResolver* KeyInfoResolver::getInstance(const char* type, const DOMElement* e)
{
    IPlugIn* p=SAMLConfig::getConfig().getPlugMgr().newPlugin(type,e);
    KeyInfoResolver* rc=dynamic_cast<KeyInfoResolver*>(p);
    if (rc)
        return rc;
    delete p;
    throw UnsupportedExtensionException("factory returned plugin type other than KeyInfoResolver");
}

SAMLSignedObject::SAMLSignedObject(istream& in) : SAMLObject(in), m_id(NULL), m_signature(NULL), m_sigElement(NULL) {}

SAMLSignedObject::SAMLSignedObject(istream& in, int minor)
    : SAMLObject(in,minor), m_id(NULL), m_signature(NULL), m_sigElement(NULL) {}

SAMLSignedObject::~SAMLSignedObject()
{
    if (m_bOwnStrings)
        XMLString::release(&m_id);
    if (m_signature)
        dynamic_cast<SAMLInternalConfig&>(SAMLConfig::getConfig()).m_xsec->releaseSignature(m_signature);
}

void SAMLSignedObject::ownStrings()
{
    if (!m_bOwnStrings) {
        m_id=XML::assign(m_id);
        m_bOwnStrings=true;
    }
}

void SAMLSignedObject::setDirty()
{
    SAMLObject::setDirty();
    unsign();
}

bool SAMLSignedObject::isSigned() const
{
    return (m_signature) ? true : false;
}

void SAMLSignedObject::unsign()
{
    if (isSigned())
    {
        m_sigElement->getParentNode()->removeChild(m_sigElement);
        m_sigElement->release();
        m_sigElement=NULL;
        dynamic_cast<SAMLInternalConfig&>(SAMLConfig::getConfig()).m_xsec->releaseSignature(m_signature);
        m_signature=NULL;
    }
}

void SAMLSignedObject::setId(const XMLCh* id)
{
    if (XML::isEmpty(id))
        throw SAMLException("id cannot be null or empty");
        
    if (m_bOwnStrings)
        XMLString::release(&m_id);
    else
        m_id=NULL;
    m_id=XML::assign(id);
    setDirty();
}

void SAMLSignedObject::sign(
    XSECCryptoKey* k,
    const Iterator<XSECCryptoX509*>& certs,
    const char* sigAlg,
    const char* digestAlg
    )
{
    if (isSigned())
        throw InvalidCryptoException("can't sign object a second time");

    // Map algorithms to constants.
    SAMLInternalConfig& conf=dynamic_cast<SAMLInternalConfig&>(SAMLConfig::getConfig());
    map<string,signatureMethod>::const_iterator i=conf.m_sigAlgFromURI.find(sigAlg);
    if (i==conf.m_sigAlgFromURI.end())
        throw InvalidCryptoException("unknown/unsupported signature algorithm: $1",params(1,sigAlg));
    signatureMethod sm=i->second;
    map<string,hashMethod>::const_iterator j=conf.m_digestAlgFromURI.find(digestAlg);
    if (j==conf.m_digestAlgFromURI.end())
        throw InvalidCryptoException("unknown/unsupported digest algorithm: $1",params(1,digestAlg));
    hashMethod hm=j->second;

    // Generate the DOM if not already built, and anchor the DOM in the document.
    toDOM();
    plantRoot();

    // Build the empty signature.
    static const XMLCh ds[]={chLatin_d, chLatin_s, chNull};
    m_signature=conf.m_xsec->newSignature();
    m_signature->setDSIGNSPrefix(ds);
    m_sigElement=m_signature->createBlankSignature(m_root->getOwnerDocument(),CANON_C14NE_NOC,sm,hm);
    
    // Have the object place it in the proper place.
    insertSignature();    
    
    try {
        DSIGReference* ref=NULL;
        if (SAMLConfig::getConfig().compatibility_mode) {
            static const XMLCh empty[]={chNull};
            ref=m_signature->createReference(empty);    
        }
        else {
            XMLCh* buf=new XMLCh[XMLString::stringLen(getId()) + 2];
            buf[0]=chPound;
            buf[1]=chNull;
            XMLString::catString(buf,getId());
            try {
                ref=m_signature->createReference(buf);
                delete[] buf;
            }
            catch(...) {
                delete[] buf;
                throw;
            }
        }

        ref->appendEnvelopedSignatureTransform();
        DSIGTransformC14n* c14n=ref->appendCanonicalizationTransform(CANON_C14NE_NOC);
        c14n->setInclusiveNamespaces(
            conf.wide_inclusive_namespace_prefixes
            );
        
        // Append any certs.
        if (certs.size()) {
            DSIGKeyInfoX509* x509Data=m_signature->appendX509Data();
            while (certs.hasNext()) {
                safeBuffer& buf=certs.next()->getDEREncodingSB();
                x509Data->appendX509Certificate(buf.sbStrToXMLCh());
            }
        }
        
        // Finally, sign the thing.
        m_signature->setSigningKey(k);
        m_signature->sign();
    }
    catch(XSECException& e) {
        if (m_signature) {
            conf.m_xsec->releaseSignature(m_signature);
            m_signature=NULL;
        }
        if (m_sigElement) {
            m_sigElement->getParentNode()->removeChild(m_sigElement);
            m_sigElement->release();
            m_sigElement=NULL;
        }
        auto_ptr_char temp(e.getMsg());
        SAML_log.error("caught an XMLSec exception: %s",temp.get());
        throw InvalidCryptoException("caught an XMLSec exception while signing: $1",params(1,temp.get()));
    }
    catch(XSECCryptoException& e) {
        if (m_signature) {
            conf.m_xsec->releaseSignature(m_signature);
            m_signature=NULL;
        }
        if (m_sigElement) {
            m_sigElement->getParentNode()->removeChild(m_sigElement);
            m_sigElement->release();
            m_sigElement=NULL;
        }
        SAML_log.error("caught an XMLSec crypto exception: %s",e.getMsg());
        throw InvalidCryptoException("caught an XMLSec crypto exception while signing: $1",params(1,e.getMsg()));
    }
}

void SAMLSignedObject::verify(XSECCryptoKey* k) const
{
    if (!isSigned())
        throw InvalidCryptoException("can't verify unsigned object");
    
    bool valid=false;

    DSIGReferenceList* refs=m_signature->getReferenceList();
    if (refs && refs->getSize()==1) {
        DSIGReference* ref=refs->item(0);
        if (ref) {
            const XMLCh* URI=ref->getURI();
            if (URI==NULL || *URI==0 || (*URI==chPound && !XMLString::compareString(URI+1,getId()))) {
                DSIGTransformList* tlist=ref->getTransforms();
                for (unsigned int i=0; tlist && i<tlist->getSize(); i++) {
                    if (tlist->item(i)->getTransformType()==TRANSFORM_ENVELOPED_SIGNATURE)
                        valid=true;
                    else if (tlist->item(i)->getTransformType()!=TRANSFORM_EXC_C14N &&
                             tlist->item(i)->getTransformType()!=TRANSFORM_C14N) {
                        valid=false;
                        break;
                    }
                }
            }
        }
    }
    
    if (!valid)
        throw InvalidCryptoException("detected an invalid signature profile while verifying signature");

    try {
        if (k) {
            m_signature->setSigningKey(k);
        }
        else {
            // Clear any existing key and install a default resolver.
            m_signature->setSigningKey(NULL);
            XSECKeyInfoResolverDefault resolver;
            m_signature->setKeyInfoResolver(&resolver); // It will clone the resolver for us.
        }
            
        if (!m_signature->verify()) {
            auto_ptr_char temp(m_signature->getErrMsgs());
            SAML_log.error("signature failed to verify, error messages follow:\n%s",temp.get());
            throw InvalidCryptoException("failed to verify signature value: $1",params(1,temp.get()));
        }
    }
    catch(XSECException& e) {
        auto_ptr_char temp(e.getMsg());
        SAML_log.error("caught an XMLSec exception: %s",temp.get());
        throw InvalidCryptoException("caught an XMLSec exception while verifying signature: $1",params(1,temp.get()));
    }
    catch(XSECCryptoException& e) {
        SAML_log.error("caught an XMLSec crypto exception: %s",e.getMsg());
        throw InvalidCryptoException("caught an XMLSec crypto exception while verifying signature: $1",params(1,e.getMsg()));
    }
}

void SAMLSignedObject::verify(XSECCryptoX509& cert) const
{
    verify(cert.clonePublicKey());
}

void SAMLSignedObject::verify(XSECKeyInfoResolver* r) const
{
    if (m_signature)
        m_signature->setKeyInfoResolver(r);
    verify();
}

const char* SAMLSignedObject::getSignatureAlgorithm() const
{
    if (isSigned()) {
        if (signatureHashMethod2URI(m_safebuf,m_signature->getSignatureMethod(),m_signature->getHashMethod()))
            return m_safebuf.rawCharBuffer();
    }
    return NULL;
}

const char* SAMLSignedObject::getDigestAlgorithm() const
{
    if (isSigned()) {
        if (hashMethod2URI(m_safebuf,m_signature->getHashMethod()))
            return m_safebuf.rawCharBuffer();
    }
    return NULL;
}

unsigned int SAMLSignedObject::getX509CertificateCount() const
{
    if (isSigned()) {
        DSIGKeyInfoList* klist=m_signature->getKeyInfoList();
        for (unsigned int i=0; klist && i<klist->getSize(); i++) {
            if (klist->item(i)->getKeyInfoType()==DSIGKeyInfo::KEYINFO_X509) {
                if (static_cast<DSIGKeyInfoX509*>(klist->item(i))->getCertificateListSize())
                    return static_cast<DSIGKeyInfoX509*>(klist->item(i))->getCertificateListSize();
            }
        }
        return 0;
    }
    throw InvalidCryptoException("can't examine certificates in unsigned object");
}

const XMLCh* SAMLSignedObject::getX509Certificate(unsigned int index) const
{
    if (isSigned()) {
        DSIGKeyInfoList* klist=m_signature->getKeyInfoList();
        for (unsigned int i=0; klist && i<klist->getSize(); i++) {
            if (klist->item(i)->getKeyInfoType()==DSIGKeyInfo::KEYINFO_X509) {
                if (index < static_cast<DSIGKeyInfoX509*>(klist->item(i))->getCertificateListSize())
                    return static_cast<DSIGKeyInfoX509*>(klist->item(i))->getCertificateItem(index);
            }
        }
        throw InvalidCryptoException("can't locate suitable ds:X509Data");
    }
    throw InvalidCryptoException("can't examine certificates in unsigned object");
}