package demo.pkcs.pkcs11.provider;

import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateKey;
import java.util.Collections;
import java.util.Enumeration;
import java.util.Properties;

import javax.servlet.ServletException;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;

import iaik.pkcs.pkcs11.provider.IAIKPkcs11;
import iaik.xml.crypto.XSecProvider;

/**
 * This servlet create an XML signature using XSECT and the IAIK PKCS#11 provider.
 * The servlet takes three init parameters:
 <ul>
 <li>MODULE<br/>The name of the PKCS#11 module, e.h. C:\HSM\pkcs11.dll</li>
 <li>SLOT<br/>The ID of the slot/token to use, e.g. 0.</li>
 <li>PIN<br/>The user PIN of the token, e.g. 0000.</li>
 </ul>
 * The request parameter <code>dataurl</code> gives the HTTP URL of the data to be
 * signed.
 <p>
 * This servlet was developed and tested with Apache Tomcat 5.5 and SUN JDK 1.5 under
 * Windows.
 */
 public class CreateSignatureServlet extends HttpServlet {
   
  /**
   * Serialization ID for compatibility.
   */
  private static final long serialVersionUID = 5707729184558821929L;

  /**
    * This is the instance of the PKCS#11 provider. It provides access to the
    * smart card.
    */
   private IAIKPkcs11 provider_;

   /**
    * This is the selected signature key. Usually, it does not contain the actual
    * key material. It is a mere proxy object.
    */
   private PrivateKey signatureKey_;

   /**
    * This is the signing certificate that is associated with the signing key.
    * The demo includes this certificate in the XML signature.
    */
   private X509Certificate signingCertificate_;

  /**
   * The default constructor. 
   */
  public CreateSignatureServlet() {
    super();
  
  
  /* (non-Javadoc)
   * @see javax.servlet.GenericServlet#init()
   */
  public void init() throws ServletException {
    super.init();
    
    // configure and install the PKCS#11 Provider
    Properties pkcs11config = new Properties();
    pkcs11config.put("PKCS11_NATIVE_MODULE", getInitParameter("MODULE"));
    pkcs11config.put("SLOT_ID", getInitParameter("SLOT"));
    pkcs11config.put("USER_PIN", getInitParameter("PIN"));
    IAIKPkcs11 provider = new IAIKPkcs11(pkcs11config);
    Security.addProvider(provider);
    provider_ = provider;
    
    // install IAIK XML Security Provider (XSECT)
    XSecProvider xsecProvider = new XSecProvider();
    XSecProvider.setDelegationProvider("Signature.SHA1withRSA", provider_.getName());
    Security.addProvider(xsecProvider);    
  }
  
  /* (non-Javadoc)
   * @see javax.servlet.Servlet#destroy()
   */
  public void destroy() {
    // remove and finalize the PKCS#11 provider
    Security.removeProvider(provider_.getName());
    try {
      provider_.getTokenManager().finalize();
    catch (Throwable ex) {
      // ignore
    };
    Security.removeProvider("XSECT");
    super.destroy();
  }     
  
  /* (non-Java-doc)
   * @see javax.servlet.http.HttpServlet#doGet(HttpServletRequest request, HttpServletResponse response)
   */
  protected void doGet(HttpServletRequest request, HttpServletResponse responsethrows ServletException, IOException {
    try {
      String dataUrlStr = request.getParameter("dataurl");
      if ((dataUrlStr == null|| dataUrlStr.trim().equals("")) {
        respondError(response, "Please enter a valid HTTP URI.");
      else {
        URI dataUri = new URI(dataUrlStr);
        if ((dataUri.getScheme() == null|| !dataUri.getScheme().equalsIgnoreCase("http")) {
          respondError(response, "Only HTTP URIs are allowed.");
        else {
          selectSignatureKey();
          
          ServletOutputStream responseStream = response.getOutputStream();
          createXmlSignature(dataUrlStr, responseStream);
          responseStream.flush();
        }
      }
    catch (IOException ex) {
      throw ex;
    catch (Exception ex) {
      throw new ServletException(ex);
    }
  }
    
  /**
   * Respond with an error message to the client.
   
   @param response The response object.
   @param msg The error message.
   @throws IOException If writing to the response stream fails.
   */
  private void respondError(HttpServletResponse response, String msgthrows IOException {
    response.setContentType("text/html");
    ServletOutputStream responseStream = response.getOutputStream();
    responseStream.println("<html>");
    responseStream.println("<head>");
    responseStream.println("<title>Error</title>");
    responseStream.println("</head>");
    responseStream.println("<body>");
    responseStream.println(msg);
    responseStream.println("</body>");
    responseStream.println("</html>");
    responseStream.flush();  
  }
  
  /**
   * This method gets the first key-entry that is a RSA key with a signature
   * certificate. It stores the key in {@link #signatureKey_} and the
   * certificate in {@link #signingCertificate_}.
   
   @exception GeneralSecurityException If anything with the key store fails.
   */
  private void selectSignatureKey()
      throws GeneralSecurityException
  {
    KeyStore tokenKeyStore = provider_.getTokenManager().getKeyStore();

    // we take the first signature (private) key for simplicity
    Enumeration aliases = tokenKeyStore.aliases();
    while (aliases.hasMoreElements()) {
      String keyAlias = aliases.nextElement().toString();
      Key key = tokenKeyStore.getKey(keyAlias, null);
      if (key instanceof RSAPrivateKey) {
        Certificate[] certificateChain = tokenKeyStore.getCertificateChain(keyAlias);
        X509Certificate signerCertificate = (X509CertificatecertificateChain[0];
        boolean[] keyUsage = signerCertificate.getKeyUsage();
        // check for digital signature or non-repudiation,
        // but also accept if none is set
        if ((keyUsage == null|| keyUsage[0|| keyUsage[1]) {
          signatureKey_ = (PrivateKeykey;
          signingCertificate_ = signerCertificate;
          break;
        }
      }
    }

    if (signatureKey_ == null) {
      throw new GeneralSecurityException(
          "Found no signature key. Ensure that a valid card/HSM is available.");
    }
  }

  /**
   * Create the actual XML signature. This method works like a usual XML
   * signature creation.
   
   @param dataURL The URL of the signed data.
   @param out The stream for writing the resulting XML signature.
   @throws GeneralSecurityException If any crypto operation fails.
   @throws ParserConfigurationException If the DOM-parser is configured
   *           incorrectly.
   @throws MarshalException If the XML signature cannot be marshaled.
   @throws XMLSignatureException If creating the XML signature fails.
   @throws TransformerException If making a string out out the result fails.
   */
  private void createXmlSignature(String dataURL, OutputStream out)
      throws GeneralSecurityException, ParserConfigurationException, XMLSignatureException,
      MarshalException, TransformerException
  {
    // First, create a DOM XMLSignatureFactory that will be used to
    // generate the XMLSignature and marshal it to DOM.
    XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

    // Create a Reference to an external URI that will be digested
    // using the SHA1 digest algorithm
    Reference ref = fac.newReference(dataURL, fac.newDigestMethod(DigestMethod.SHA1, null));
    
    // Create the SignedInfo
    CanonicalizationMethod canonicalizationMethod = fac.newCanonicalizationMethod(
        CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS, (C14NMethodParameterSpecnull);
    SignatureMethod signatureMethod = fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null);
    SignedInfo si =
        fac.newSignedInfo(canonicalizationMethod, signatureMethod, Collections.nCopies(1, ref));

    // Create a KeyValue containing the DSA PublicKey that was generated
    KeyInfoFactory kif = fac.getKeyInfoFactory();
    X509Data x509data = kif.newX509Data(Collections.nCopies(1, signingCertificate_));

    // Create a KeyInfo and add the KeyValue to it
    KeyInfo ki = kif.newKeyInfo(Collections.nCopies(1, x509data));

    // Create the XMLSignature (but don't sign it yet)
    XMLSignature signature = fac.newXMLSignature(si, ki);

    // Create the Document that will hold the resulting XMLSignature
    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
    dbf.setNamespaceAware(true)// must be set
    Document doc = dbf.newDocumentBuilder().newDocument();

    // Create a DOMSignContext and set the signing Key to the DSA
    // PrivateKey and specify where the XMLSignature should be inserted
    // in the target document (in this case, the document root)
    DOMSignContext signContext = new DOMSignContext((KeysignatureKey_, doc);

    // Marshal, generate (and sign) the detached XMLSignature. The DOM
    // Document will contain the XML Signature if this method returns
    // successfully.
    signature.sign(signContext);
    
    // output the resulting document
    TransformerFactory tf = TransformerFactory.newInstance();
    Transformer trans = tf.newTransformer();
    trans.transform(new DOMSource(doc)new StreamResult(out));
  }

  
}