import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.DocumentBuilder;
import org.apache.xerces.parsers.DOMParser;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.w3c.dom.NamedNodeMap;
import org.xml.sax.InputSource;
import org.apache.xml.serialize.OutputFormat;
import org.apache.xml.serialize.XMLSerializer;

import com.ibm.xml.enc.AlgorithmFactoryExtn;
import com.ibm.xml.enc.EncryptionContext;
import com.ibm.xml.enc.KeyInfoResolver;
import com.ibm.xml.enc.util.KeyStoreKeyInfoResolver;

import java.io.InputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.ByteArrayInputStream;
import java.util.Vector;
import java.security.KeyStore;
import org.apache.xpath.XPathAPI;

public class EncryptedKeyToken implements Token
{
   private Element encryptionTemplate = null;
   private String encryptionTemplateFile = "EncryptionTemplate.xml";
   private KeyStoreKeyInfoResolver kiRes = null;
   private AlgorithmFactoryExtn af = null;
   private WSSMessage wssMessage = null;
   private String keyName = "";
   private String wsuKeyID = "";
   private Element tokenElement = null;


   public EncryptedKeyToken (
                  String keyStoreFileName,
                  String keyStorePassword,
                  String keyName,
                  String wsuKeyID,
                  WSSMessage wssMessage
                  ){

      this.wssMessage = wssMessage;
      this.wsuKeyID = wsuKeyID;
      this.keyName = keyName;

      DOMParser dp = new DOMParser();

      try {
         dp.parse ( encryptionTemplateFile );
         encryptionTemplate = dp.getDocument().getDocumentElement();

         KeyStore store = null;
         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);

         // Create a new KeyStoreKeyInfoResolver object.
         kiRes = new KeyStoreKeyInfoResolver( store );

         // Instantiate an algo factory.
         af = new AlgorithmFactoryExtn();

         // Set the algo factory for use by the key resolver.
         kiRes.setAlgorithmFactory(af);

      }//try
      catch ( Exception e ) {
         System.out.println ( "Exception in EncryptedKeyTokon's Constructor......." );
         e.printStackTrace();
      }//catch

   }//Constructor


   // Return the value of the wsu:Id attribute of the token.
   public String getWSUId()
   {
      return wsuKeyID;
   }//getWSUId


   // Return a string identifier for the type of token.
   public String getType()
   {
      return "EncryptedKeyToken";
   }//getType


   // Return the parent WSSMessage object that wraps this token.
   public WSSMessage getParentWSSMessage()
   {
      return wssMessage;
   }//getParentWSSMessage


   // Return the WSS-compliant XML structure of this token.
   public String getXMLString()
   {
      return getXMLString( tokenElement );
   }//getXMLString


   // Return the W3C DOM object representation of the token object.
   public org.w3c.dom.Element getDOMObject()
   {
      return null;
   }//getDOMObject


   // The Token instance will receive an array of bytes (secretBytes) that
   // it will treat as the secret to be used to
   // produce signatures and decrypt encrypted data.
   public void setSecret(byte[] secretBytes)
   {
      String secret = new String ( secretBytes );

      // The key resolver needs to know the key pass.
      char keyPass[];
      keyPass = new String ( secret ).toCharArray();
      kiRes.putAliasAndPassword( keyName, keyPass );

   }//setSecret


   public Element encrypt(
                  String wsuEncryptedElementID,
                  String encryptionAlgo,
                  Element elementWantToEncrypt
                          )
   {

      /**** Step 1 *********/
      // Set the mode of operation of the key resolver to "encryption".
      kiRes.setOperationMode(KeyInfoResolver.ENCRYPT_MODE);


      /**** Step 2 *********/
      Element dataReference = (Element)encryptionTemplate.
                 getElementsByTagName("DataReference").item(0);
      dataReference.setAttribute(
         "URI",
         "#"+wsuEncryptedElementID
      );


      /**** Step 3 *********/
      Element encryptedData = 
         (Element)encryptionTemplate.getOwnerDocument().
            getElementsByTagName("EncryptedData").item(0);
      encryptedData.setAttribute ( "wsu:Id", wsuEncryptedElementID );


      /**** Step 4 *********/
      Element encAlgo = (Element)encryptionTemplate.getElementsByTagName("EncryptionMethod").item(0);
      encAlgo.setAttribute ( "Algorithm", encryptionAlgo );


      /**** Step 5 *********/
      //EncryptedKey element in the encryption template is the actual token.
      //We need to set its ID.
      tokenElement = (Element)encryptionTemplate.getElementsByTagName("EncryptedKey").item(0);
      tokenElement.setAttribute ( "wsu:Id", getWSUId() );


      /**** Step 6 *********/
      EncryptionContext ec = new EncryptionContext();
      ec.setAlgorithmFactory (af);
      ec.setKeyInfoResolver (kiRes);
      ec.setData (elementWantToEncrypt);
      ec.setEncryptedType ( (Element)encryptionTemplate.cloneNode(true),
                            null, null, null);

      try {
         ec.encrypt();
         ec.replace();
      }
      catch(Exception e) {
           System.out.println ( "Error in encryptionContext" );
           e.printStackTrace();
      }


      /**** Step 7 *********/
      Document wssMessage = elementWantToEncrypt.getOwnerDocument();
      Element encryptedDataEl = getElementByWSUId ( wssMessage, wsuEncryptedElementID );
      tokenElement = (Element)encryptedDataEl.getElementsByTagName("EncryptedKey").item(0).cloneNode(true);
      Element encryptedKey = (Element)encryptedDataEl.getElementsByTagName("EncryptedKey").item(0);
      // We have stored the EncryptedKey element as tokenElement object.
      // But we don't need the EncryptedKey to reside inside the EncryptedData.
      // Therefore, remove the EncryptedKey node and normalize the DOM tree.
      encryptedKey.getParentNode().removeChild(encryptedKey);
      wssMessage.normalize();

      /**** Step 8 *********/
      Element keyInfo = (Element)encryptedDataEl.getElementsByTagName("KeyInfo").item(0);
      // The KeyInfo element in the EncryptedData is unnecessary,
      // so remove it and normalize the DOM tree.
      keyInfo.getParentNode().removeChild(keyInfo);
      wssMessage.normalize();

      return encryptedDataEl;

   }// End of encrypt Element


   private String getXMLString ( Element elementToText ) {

      String s = new String();

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

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

         s += ">";

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

         s += "";

      }//if
      return s;
   }//getXMLString


   private Element getElementByWSUId (
                               Document document,
                               String wsuEncryptedElementID
                                      ) {

      Element root = document.getDocumentElement();
      String xpathString = "//*[@Id=\"" + wsuEncryptedElementID + "\"]";
      Element plainTextElement = null;
      try {
         plainTextElement = (Element)XPathAPI.selectSingleNode(root, xpathString);
      }
      catch ( Exception e ) {
         System.out.println ( "Exception in getting the plainTextElement." );
         e.printStackTrace();
      }

      return plainTextElement;

   }//getElementByWSUId 


   // The sign, signWithXPath, and decrypt methods are included here
   // because the Token knows the secret and the secret is needed to
   // produce signatures and for decryption.
   // Therefore, the WSSMessage, Signature, and EncryptedData classes will
   // use these methods internally.
   // The Token should never tell the secret to any other class.

   public String sign (
                         String wsuElementID,
                         String digestAlgo,
                         String signatureAlgo,
                         String canonicalizationAlgo
                        )
   {
      return null;
   }


   public String signWithXPath (
                            String XPathExpression,
                            String digestAlgo,
                            String signatureAlgo,
                            String canonicalizationAlgo
                         )
   {
      return null;
   }


   public void decrypt(String encryptedData)
   {
   }


}//EncryptedKeyToken