package org.owasp.saml;

import org.w3c.dom.Document;

import javax.xml.XMLConstants;
import javax.xml.namespace.NamespaceContext;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.logging.Logger;

public class SamlNamespaceResolver implements NamespaceContext {

    private final static Logger LOG = Logger.getLogger(SamlNamespaceResolver.class.getName());

    private Document sourceDocument;

    /**
     * This constructor stores the source document to search the namespaces in
     * it.
     *
     * @param document
     *            source document
     */
    public SamlNamespaceResolver(Document document) {
        sourceDocument = document;
    }

    /**
     * The lookup for the namespace uris is delegated to the stored document.
     *
     * @param prefix
     *            to search for
     * @return uri
     */
    public String getNamespaceURI(String prefix) {

        String uri = null;
        final String SAML10 = "urn:oasis:names:tc:SAML:1.0:assertion";
        final String SAML20 = "urn:oasis:names:tc:SAML:2.0:assertion";
        final String WSSE = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd";
        final String DS = "http://www.w3.org/2000/09/xmldsig#";
        final String WSU = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";

        uri =  sourceDocument.lookupNamespaceURI(prefix);

        if (uri == null) {

            final Map<String, String> hm = new HashMap<String, String>();
            hm.put("saml2", "urn:oasis:names:tc:SAML:2.0:assertion");
            hm.put("saml", "urn:oasis:names:tc:SAML:1.0:assertion");
            hm.put("wsse", "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd");
            hm.put("ds", "http://www.w3.org/2000/09/xmldsig#");
            hm.put(XMLConstants.DEFAULT_NS_PREFIX, "http://www.w3.org/2000/09/xmldsig#");
            hm.put("wsu", "urn:oasis:names:tc:SAML:2.0:assertion");
            hm.put("saml2", "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd");
            hm.put("xs", "http://www.w3.org/2001/XMLSchema");
            hm.put("wsp", "http://schemas.xmlsoap.org/ws/2004/09/policy");
            hm.put("xsa", "http://schemas.xmlsoap.org/ws/2004/08/addressing");
            hm.put("wst", "http://schemas.xmlsoap.org/ws/2005/02/trust");

            uri = hm.get(prefix);
        }

        if (uri == null ) {
            LOG.warning("getNamespaceURI prefix=" + prefix + " returns=" + uri);
        } else {
            LOG.info("getNamespaceURI prefix=" + prefix + " returns=" + uri);
        }

        return uri;
    }

    public String getPrefix(String namespaceURI) {
        return sourceDocument.lookupPrefix(namespaceURI);
    }

    public Iterator getPrefixes(String namespaceURI) {
        // not implemented yet
        return null;
    }

}

@@ -864,32 +864,46 @@
 urn uri%0A     */%0A
+    @Override%0A
     public Strin
@@ -950,281 +950,90 @@
     
-String uri = null;%0A        final String SAML10 = %22urn:oasis:names:tc:SAML:1.0:assertion%22;%0A        final String SAML20 = %22urn:oasis:names:tc:SAML:2.0:assertion%22;%0A        final String WSSE = %22http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd%22;
+// First try to automatically resolve the namespace URI based on document contents
 %0A   
@@ -1041,236 +1041,128 @@
     
-final 
 String 
-DS = %22http://www.w3.org/2000/09/xmldsig#%22;%0A        final String WSU = %22http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd%22;%0A%0A        uri =  sourceDocument.lookupNamespaceURI(prefix);%0A
+uri = this.sourceDocument.lookupNamespaceURI(prefix);%0A%0A        // If this doesn't work resort to hardcoded values
 %0A   
@@ -1580,98 +1580,8 @@
 %22);%0A
-            hm.put(XMLConstants.DEFAULT_NS_PREFIX, %22http://www.w3.org/2000/09/xmldsig#%22);%0A
     
@@ -2046,24 +2046,271 @@
 2/trust%22);%0A%0A
+            // This worked for me, most likely won't work for you%0A            // default prefix is just empty prefix, so %3CSignature%3E and not %3Cds:Signature%3E%0A            hm.put(XMLConstants.DEFAULT_NS_PREFIX, %22http://www.w3.org/2000/09/xmldsig#%22);%0A%0A
             
@@ -2576,32 +2576,46 @@
 urn uri;%0A    %7D%0A%0A
+    @Override%0A
     public Strin
@@ -2710,24 +2710,38 @@
 RI);%0A    %7D%0A%0A
+    @Override%0A
     public I

add @Override, cleanup