package WSS4J;

import com.ibm.xml.sax.StandardErrorHandler;
import java.io.ByteArrayInputStream;
import java.security.Key;
import java.security.KeyStore;
import java.security.cert.X509Certificate;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import org.w3c.dom.Node;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import com.ibm.xml.dsig.KeyInfo;
import com.ibm.xml.dsig.SignatureContext;
import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import org.w3c.dom.NamedNodeMap;

import com.ibm.xml.dsig.KeyInfo;
import com.ibm.xml.dsig.SignatureContext;
import com.ibm.xml.dsig.util.DOMParserNS;
import org.apache.xerces.parsers.DOMParser;
import org.w3c.dom.NodeList;

import org.xml.sax.InputSource;

public class SignatureToken implements Token
{
    private String signatureTemplateFile = "SignatureTemplate.xml"; 
    private static WSSMessage wssMessage = null;
    private static String tokenWSUId = "";
    private static KeyStore store = null;
    private static String userName = "";
    private static String keyName = "";

    protected static Key key = null;
    protected static X509Certificate certificate = null;
    protected static Element templateRootEl = null;

    public SignatureToken(
                     String keyStoreFileName,
                     String keyStorePassword,
                     String keyName,
                     String tokenWSUId,
                     WSSMessage wssMessage
                     ){

        this.keyName = keyName;
        this.tokenWSUId = tokenWSUId;
        this.wssMessage = wssMessage;

        try{
            store = KeyStore.getInstance("jceks");

            File f = new File ( keyStoreFileName );

            InputStream is = null;
            if (f.exists()){
                is = new FileInputStream ( f );
            }

            char[] storePass;
            storePass = new String( keyStorePassword ).toCharArray();
            store.load(is, storePass);

            certificate = (X509Certificate)store.getCertificate(keyName);

            loadTemplate();
        }catch( Exception e ) {
            System.out.println( 
              "Exception in SignatureToken Constructor1......" );
            e.printStackTrace();
        }

    }


    public SignatureToken( 
                 String tokenWSUId,
                 String userName,
                 WSSMessage wssMessage 
                 ){
         this.userName = userName;
         this.wssMessage = wssMessage;
         this.tokenWSUId = tokenWSUId;

         loadTemplate();

    }

    public String getWSUId()
    {
        return tokenWSUId;
    }//end of getWSUId method

    public String getType()
    {
        return null;
    }//end of getType method

    public WSSMessage getParentWSSMessage()
    {
        return wssMessage;
    }//end of getParentWSSMessage method

    public String getXMLString()
    {
        return null;
    }//end of getXMLString method

    public org.w3c.dom.Element getDOMObject()
    {
        return null;
    }//end of getDOMObject method

    public void setSecret(byte[] secretBytes)
    {
        try{
            String secret = new String( secretBytes );

            key = store.getKey(keyName, secret.toCharArray()); 
            if (key == null) {
                throw new RuntimeException( "Could not get a key: " + keyName );
            } 

        }catch( Exception e ) {
            System.out.println( 
                 "Exception in setSecret method...." );
            e.printStackTrace();
        }
    }//end of setSecret method

    public String sign(
            String wsuElementID,
            String digestAlgo,
            String signatureAlgo,
            String canonicalizationAlgo
             ){

          Element templateElement = null;
          try{

              /*****Step1*****/
              Document inputDoc = wssMessage.getWSSDocument();
              templateElement =  
                                 (Element) inputDoc.importNode(
                                 templateRootEl, 
                                 true
                                 );

              /*****Step2*****/
              Element reference = (Element) templateElement.
                                   getElementsByTagName("Reference").item(0);
              reference.setAttribute( "URI", "#" + wsuElementID );

              /*****Step3*****/
              Element digestMethod = 
                               (Element)templateElement.
                            getElementsByTagName("DigestMethod").item(0);
              digestMethod.setAttribute ( "Algorithm", digestAlgo );

              Element signatureMethod = 
                              (Element)templateElement.
                           getElementsByTagName("SignatureMethod").item(0);
              signatureMethod.setAttribute ( "Algorithm", signatureAlgo );

              Element canonicalizationMethod = 
                           (Element)templateElement.
                        getElementsByTagName("CanonicalizationMethod").item(0);
              canonicalizationMethod.setAttribute 
                                      ( "Algorithm", canonicalizationAlgo );

              /*****Step4*****/
              KeyInfo keyInfo = new KeyInfo();
              keyInfo.insertTo(templateElement);  

              /*****Step5*****/
              Element keyInfoEl = (Element)templateElement.
                                     getElementsByTagName("KeyInfo").item(0);

              Element securityTokenRef = templateElement.
                                    getOwnerDocument().createElementNS(
                               "http://schemas.xmlsoap.org/ws/2002/secext",
                               "SecurityTokenReference"
                               );

              securityTokenRef.setAttribute(
                               "xmlns",
                               "http://schemas.xmlsoap.org/ws/2002/secext"
                               );

              keyInfoEl.appendChild ( securityTokenRef );

              /*****Step6*****/     
              Element referenceEl = templateElement.
                              getOwnerDocument().createElementNS(
                           "http://schemas.xmlsoap.org/ws/2002/secext",
                           "Reference"
                           );

              referenceEl.setAttribute(
                           "URI",
                           "#" + tokenWSUId
                           );

              securityTokenRef.appendChild ( referenceEl );

              /*****Step7*****/
              SignatureContext sigContext = new SignatureContext();

              /*****Step8*****/
              sigContext.setIDResolver(new WSUIdResolver());

              /*****Step9*****/
              sigContext.sign(templateElement, key);

          }catch( Exception e ) {
             System.out.println( 
                 "Exception in SignatureToken's sign......" );
             e.printStackTrace();
          } // end of catch

        return getXMLString ( templateElement );   
    }// end of sign method

    public String signWithXPath(
            String XPathExpression,
            String digestAlgo,
            String signatureAlgo,
            String canonicalizationAlgo
             ){

         Element templateElement = null;
         try{

             /*****Step1*****/
             Document inputDoc = wssMessage.getWSSDocument();
             templateElement =  
                                  (Element) inputDoc.importNode(
                                   templateRootEl, 
                                   true
                                   );

             /*****Step2*****/
             String xpathTransform = 
                    "\n<Transforms>" +
                    "\n    <Transform " + 
                    "\nAlgorithm=\"http://www.w3.org/TR/1999/REC-xpath-19991116\">" +
                    "\n        <XPath xmlns=\"http://www.w3.org/2000/09/xmldsig#\">" +
                    "\n            " + XPathExpression +
                    "\n        </XPath>" + 
                    "\n    </Transform>" +
                    "\n<Transforms>\n";

             Document XPathTransformDoc = loadDocument(xpathTransform);
             Element referenceElementFirstChild = null;
             Element referenceElement = (Element) templateElement.
                                     getElementsByTagName ( "Reference" ).
                                     item(0);

             NodeList referenceElementChilds = referenceElement.
                                                          getChildNodes();
             for ( int i = 0; i < referenceElementChilds.getLength(); i++ ){
                 if ( referenceElementChilds.item(i).getNodeType() == 
                                                       Node.ELEMENT_NODE ){
                     referenceElementFirstChild = 
                                     (Element) referenceElementChilds.item(i);
                     break;  
                 }// end of if
             }// end of for
             Element XPathTransformElement = (Element) inputDoc.importNode(
                                 XPathTransformDoc.getDocumentElement(), 
                                 true
                                 );

             referenceElement.insertBefore(
                                   XPathTransformElement,
                                   referenceElementFirstChild
                                   );

             /*****Step3*****/
             Element digestMethod = 
                          (Element)templateElement.
                       getElementsByTagName("DigestMethod").item(0);
             digestMethod.setAttribute ( "Algorithm", digestAlgo );

             Element signatureMethod = 
                             (Element)templateElement.
                          getElementsByTagName("SignatureMethod").item(0);
             signatureMethod.setAttribute ( "Algorithm", signatureAlgo );

             Element canonicalizationMethod = 
                          (Element)templateElement.
                       getElementsByTagName("CanonicalizationMethod").item(0);
             canonicalizationMethod.setAttribute 
                                       ( "Algorithm", canonicalizationAlgo );

             /*****Step4*****/
             KeyInfo keyInfo = new KeyInfo();
             keyInfo.insertTo(templateElement);

             /*****Step5*****/
             Element keyInfoEl = (Element)templateElement.
                                      getElementsByTagName("KeyInfo").item(0);

             Element securityTokenRef = templateElement.
                                     getOwnerDocument().createElementNS(
                                "http://schemas.xmlsoap.org/ws/2002/secext",
                                "SecurityTokenReference"
                                );

             securityTokenRef.setAttribute(
                                "xmlns",
                                "http://schemas.xmlsoap.org/ws/2002/secext"
                                );

             keyInfoEl.appendChild ( securityTokenRef );

             /*****Step6*****/     
             Element reference = templateElement.
                             getOwnerDocument().createElementNS(
                          "http://schemas.xmlsoap.org/ws/2002/secext",
                          "Reference"
                          );

             reference.setAttribute(
                          "URI",
                          "#" + tokenWSUId
                          );

             securityTokenRef.appendChild ( reference );

             /*****Step7*****/
             SignatureContext sigContext = new SignatureContext();

             /*****Step8*****/

             /*****Step9*****/
             sigContext.sign(templateElement, key);

         }catch( Exception e ) {
            System.out.println( 
                "Exception in SignatureToken's sign......" );
            e.printStackTrace();
         } // end of catch

        return getXMLString ( templateElement );
    }// end of sign with xpath method

    public Document loadDocument ( String XMLString ) {
       DOMParser dp = new DOMParser();
       InputStream byteData = null;
       byteData = new ByteArrayInputStream (XMLString.getBytes());
         try
         {
            dp.parse ( new InputSource ( byteData ) );
         }catch ( Exception e )
         {
            System.out.println ( "The XMLString is not loaded." );
            e.printStackTrace();
         }

         return dp.getDocument();
    }//loadDocument

    public void decrypt(String encryptedData)
    {
    }//end of decrypt method

    public String getXMLString ( Element elementToText ) {
       String elementString = new String();

       if ( elementToText != null ){
           elementString += "<" + elementToText.getTagName();
           NamedNodeMap elementAttributes = elementToText.getAttributes();
           int totalAttributes = elementAttributes.getLength();

           for ( int j = 0; j < totalAttributes; j++ )
           {
               elementString += " " + elementAttributes.item(j).getNodeName() + "=\"";
               elementString += elementAttributes.item(j).getNodeValue() + "\"";
          }//for

           elementString += ">";

           for ( int i = 0; i < elementToText.getChildNodes().getLength(); i++ )
           {
             Node child = elementToText.getChildNodes().item(i);
             if ( child.getNodeType() == Node.ELEMENT_NODE )
                 elementString += getXMLString( (Element)child );
             if ( child.getNodeType() == Node.TEXT_NODE )
                 elementString += child.getNodeValue();
           }//for

           elementString += "";

       }//if
       return elementString;
    }// end of getXMLString private method

    private void loadTemplate()
    {
        try{ 
            DocumentBuilder builder = DOMParserNS.createBuilder();
            builder.setErrorHandler(new StandardErrorHandler());
   
            templateRootEl = builder.parse(signatureTemplateFile).
                                                           getDocumentElement();
        }catch( Exception e )
        {
            System.out.println ( "Exception in loadTemplate method......" );
            e.printStackTrace();
        }
    }//end of loadTemplate 
}//end of SignatureToken class