package com.fsenablers.saml;


import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.StringWriter;
import java.io.StringReader;
import java.io.PrintWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.*;

import javax.xml.XMLConstants;
import javax.xml.crypto.*;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dom.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.*;
import javax.xml.crypto.dsig.spec.*;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.DocumentBuilder;

import javax.xml.transform.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;
import javax.xml.xpath.*;
import org.xml.sax.InputSource;


import org.w3c.dom.*;

/**
 * Sign XML file.
 */
public class XMLSigner {
	
	
	private static final String STYLESHEET = 
	"<xsl:stylesheet version='1.0' xmlns:xsl='http://www.w3.org/1999/XSL/Transform' xmlns:ds='http://www.w3.org/2000/09/xmldsig#'>"+
"<xsl:template match='ds:Signature'>"+
"<xsl:element name='{local-name()}' namespace='http://www.w3.org/2000/09/xmldsig#'>"+
   "<xsl:apply-templates select='@*|node()'/>"+
"</xsl:element>"+
"</xsl:template>"+
"<xsl:template match='ds:Signature//*'>"+
"<xsl:element name='{local-name()}' namespace='http://www.w3.org/2000/09/xmldsig#'>"+
   "<xsl:apply-templates select='@*|node()'/>"+
"</xsl:element>"+
"</xsl:template>"+
"<xsl:template match='@*|node()'>"+
  "<xsl:copy>"+
    "<xsl:apply-templates select='@*|node()'/>"+
  "</xsl:copy>"+
"</xsl:template>"+
"</xsl:stylesheet>";
		
	
	
    private static final String KEY_STORE_TYPE = "JKS";
    private static final String KEY_STORE_NAME = "/devtools/data/keystore.jks";
    private static final String KEY_STORE_PASS = "password";
    private static final String PRIVATE_KEY_PASS = "password";
    private static final String KEY_ALIAS = "mykey";

    private static final String PATH = "/PatientRecord/Account";
    private static final String ID = "acct";

    private static enum SignatureType {
	SIGN_BY_ID,
	SIGN_BY_PATH,
	SIGN_WHOLE_DOCUMENT
    };

    
    public static String sign(
    		String keyStoreName,
    		String keyStorePass,
    		String keyAlias,
    		String privateKeyPass,
    		String unsignedXMLSrc, 
    		String signBy, 
    		String signByTarget ) throws Exception {

    	
    	SignatureType sigType = SignatureType.SIGN_WHOLE_DOCUMENT;
	    if ("id".equals(signBy)) {
		sigType = SignatureType.SIGN_BY_ID;
	    }
	    else if ("path".equals(signBy)) {
		sigType = SignatureType.SIGN_BY_PATH;
	    }

	    
		// Instantiate the document to be signed
		DocumentBuilderFactory dbFactory = DocumentBuilderFactory.newInstance();
		dbFactory.setNamespaceAware(true);
		Document doc = dbFactory
			       .newDocumentBuilder()
			       .parse( new InputSource( new StringReader( unsignedXMLSrc )));

		// prepare signature factory
        String providerName = System.getProperty(
			          "jsr105Provider", 
				  "org.jcp.xml.dsig.internal.dom.XMLDSigRI"
			      );

		final XMLSignatureFactory sigFactory = XMLSignatureFactory.getInstance(
						    "DOM",
						    (Provider) Class.forName(providerName).newInstance()
						 );
	
		Node nodeToSign = null;
		Node sigParent = null;
		String referenceURI = null;
		String xPathStr = null;
		XPathExpression expr = null; 
		NodeList nodes;
		// List transforms = null;
        List<Transform> transforms = new ArrayList<Transform>();
	
		XPathFactory factory = XPathFactory.newInstance();
		XPath xpath = factory.newXPath();
		switch (sigType) {
		    case SIGN_BY_ID:
		    	
		    	nodes = null;
		    	String[] attrNames = { "ID", "id", "Id" };
		    	boolean bFound = false;
		    	for( int x = 0; x < attrNames.length; x++ ) {
		    		
			    	xPathStr = String.format("//*[@%s='%s']", attrNames[x], signByTarget );
			    	expr = xpath.compile( xPathStr );
			    	
			    	nodes = (NodeList) expr.evaluate(doc, XPathConstants.NODESET);
			    	if( nodes.getLength() > 0 ) {
			    		bFound = true;
			    		break;
			    	}
		    		
		    	}
		    	/*
		    	xPathStr = String.format("//*[@ID='%s']", signByTarget );
		    	expr = xpath.compile( xPathStr );
		    	
		    	nodes = (NodeList) expr.evaluate(doc, XPathConstants.NODESET);
				if (nodes.getLength() == 0) {
					
					
					
				    System.out.println("Can't find node with id: " + signByTarget);
				    return null;
				}
				*/
				
				if( ! bFound ) {
				    System.out.println("Can't find node with id: " + signByTarget);
				    return null;
					
				}
	
				nodeToSign = nodes.item(0);
				sigParent = nodeToSign; // sigParent = nodeToSign.getParentNode();
				referenceURI = "#" + signByTarget;
	
	
				Transform envTransform = sigFactory.newTransform( Transform.ENVELOPED, (TransformParameterSpec) null );
				transforms.add( envTransform );
	
				Transform c14nTransform = sigFactory.newTransform( CanonicalizationMethod.EXCLUSIVE, (TransformParameterSpec) null );
				transforms.add( c14nTransform );
			break;
			
		    case SIGN_BY_PATH:
			// Find the node to be signed by PATH
			expr = xpath.compile(PATH);
			nodes = (NodeList) expr.evaluate(doc, XPathConstants.NODESET);
			if (nodes.getLength() < 1) {
			    System.out.println("Invalid document, can't find node by PATH: " + PATH);
			    return null;
			}
	
			nodeToSign = nodes.item(0);
			sigParent = nodeToSign.getParentNode();
			referenceURI = ""; // Empty string means whole document
			transforms = new ArrayList<Transform>() {{
			    add(sigFactory.newTransform(
				    Transform.XPATH,
				    new XPathFilterParameterSpec(PATH)
				)
			    );
			    add(sigFactory.newTransform(
				    Transform.ENVELOPED,
				    (TransformParameterSpec) null
				)
			    );
			}};
	
			break;
		    default:
			sigParent = doc.getDocumentElement();
			referenceURI = ""; // Empty string means whole document
			transforms = Collections.singletonList(
				    	sigFactory.newTransform(
					    Transform.ENVELOPED, 
					    (TransformParameterSpec) null
					)
				    );
			break;
		}
	
	        // Retrieve signing key
		KeyStore keyStore = KeyStore.getInstance(KEY_STORE_TYPE);
		keyStore.load(
		    new FileInputStream(keyStoreName),  
		    keyStorePass.toCharArray()
		);
	
		PrivateKey privateKey = (PrivateKey) keyStore.getKey(
					    keyAlias,
					    privateKeyPass.toCharArray()
					);
	
		X509Certificate cert = (X509Certificate) keyStore.getCertificate(keyAlias);
		System.out.println( "CERT - START" );
		System.out.println( cert );
		System.out.println( "CERT - END" );

		
		// System.out.println( cert.getSubjectDN() );
		PublicKey publicKey = cert.getPublicKey();
		System.out.println( "PUBLIC KEY - START" );
		System.out.println( publicKey );
		System.out.println( "PUBLIC KEY - END" );
	
	        // Create a Reference to the enveloped document
		Reference ref = sigFactory.newReference(
				    referenceURI,
				    sigFactory.newDigestMethod(DigestMethod.SHA1, null),
				    transforms,
				    null, 
				    null
				);
	
		
		
		
		// Create the SignedInfo
		SignedInfo signedInfo = sigFactory.newSignedInfo(
					    sigFactory.newCanonicalizationMethod(
						CanonicalizationMethod.EXCLUSIVE,
						(C14NMethodParameterSpec) null
					    ), 
					    sigFactory.newSignatureMethod(
						SignatureMethod.RSA_SHA1,
						null
					    ),
					    Collections.singletonList(ref)
					);
	
	        // Create a KeyValue containing the RSA PublicKey 
		KeyInfoFactory keyInfoFactory = sigFactory.getKeyInfoFactory();
	        // KeyValue keyValue = keyInfoFactory.newKeyValue(publicKey);
	        X509Data x509Data = keyInfoFactory.newX509Data( Collections.singletonList(cert));
	
		// Create a KeyInfo and add the KeyValue to it

	        KeyInfo keyInfo = keyInfoFactory.newKeyInfo(Collections.singletonList(x509Data)); // Collections.singletonList(keyValue)
	
	        // Create a DOMSignContext and specify the RSA PrivateKey and
	        // location of the resulting XMLSignature's parent element
		DOMSignContext dsc = new DOMSignContext(
					 privateKey, 
					 sigParent
				     );
	
		// dsc.setBaseURI( "http://www.fsenablers.com/dsig#" );
		// dsc.putNamespacePrefix( "fse", "http://www.w3.org/2000/09/xmldsig#" );
		// dsc.putNamespacePrefix( "ds", null );
		dsc.setDefaultNamespacePrefix( "ds" );
		
		// Create the XMLSignature (but don't sign it yet)
		XMLSignature signature = sigFactory.newXMLSignature(signedInfo, keyInfo);
		System.out.println( signature.getClass() );
	
	        // Marshal, generate (and sign) the enveloped signature
	        signature.sign(dsc);
	
		// output the resulting document
	    StringWriter signedXMLSrcWriter = new StringWriter();
	    
	    TransformerFactory transformerFactory = TransformerFactory.newInstance();
	    // transformerFactory.setFeature(  XMLConstants.FEATURE_SECURE_PROCESSING, true );
	    // transformerFactory.setAttribute( "http://xml.apache.org/xalan/features/optimize", Boolean.FALSE );
	    System.out.println( transformerFactory.getClass() );
		
	    
        // DocumentBuilderFactory xsf = DocumentBuilderFactory.newInstance();
        // xsf.setNamespaceAware(true);
        //  DocumentBuilder xsBuilder = xsf.newDocumentBuilder();
        // InputSource stylesheet = new InputSource( new StreamSource ( new StringReader( STYLESHEET )));
        
	    // FileInputStream fis = new FileInputStream( "/devtools/stylesheet.xsl" );
	    // StringReader reader = new StringReader( STYLESHEET );
	    // Transformer trans = transformerFactory.newTransformer( new StreamSource( reader ));
	    Transformer trans = transformerFactory.newTransformer();
	    System.out.println( trans.getClass() );
		trans.setOutputProperty( OutputKeys.METHOD, "xml" );
		trans.setOutputProperty( OutputKeys.OMIT_XML_DECLARATION, "yes" );
		
		
		
		trans.transform(new DOMSource(doc), new StreamResult(signedXMLSrcWriter));
		
		// System.out.println( trans.getClass().toString() );
		
		return signedXMLSrcWriter.toString();
	    
    }

    
    
    public static void main(String[] args) throws Exception {
	if (args.length < 2) {
	    usage();
	    return;
	}

	String inputFile = args[0];
	String outputFile = args[1];
	
	String signByTarget = null;
	
	String signBy = "whole";
	if (args.length >= 3) {
		signBy = args[2];
	}
	if( "id".equals( signBy ) && args.length >= 4 ) {
		signByTarget = args[3];
	}
	
	StringWriter unsignedXMLSrcWriter = new StringWriter();
	PrintWriter writer = new PrintWriter( unsignedXMLSrcWriter );
	BufferedReader inputFileReader = new BufferedReader( new FileReader( inputFile ));
	String lineOfText = inputFileReader.readLine();
	while( lineOfText != null ) {
		writer.println( lineOfText );
		lineOfText = inputFileReader.readLine();
	}
	
	String unsignedXMLSrc = unsignedXMLSrcWriter.toString();
	
	String signedXMLSrc = sign( KEY_STORE_NAME, KEY_STORE_PASS, KEY_ALIAS, PRIVATE_KEY_PASS, unsignedXMLSrc, signBy, signByTarget );
	
	FileWriter fileWriter = new FileWriter( outputFile );
	fileWriter.write( signedXMLSrc );
	fileWriter.close();
	

    }
   
    private static void usage() {
	System.out.println("Usage: java XMLSigner <inputFile> <outputFile> [id|path|whole]");
    }
}