package in.gov.digilocker.uidai.services.docsigner;

import java.io.StringReader;
import java.io.StringWriter;
import java.security.KeyStore;
import java.security.Provider;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

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.Transform;
import javax.xml.crypto.dsig.XMLSignature;
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.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.xml.sax.InputSource;

public class DocSigner {
	private static final String prov = System.getProperty("jsr105Provider", "org.jcp.xml.dsig.internal.dom.XMLDSigRI");
	
	public static String signXML(String xmlDocument, boolean includeKeyInfo, KeyStore.PrivateKeyEntry keyEntry)throws Exception {		
		DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
		dbf.setNamespaceAware(true);
		Document inputDocument = dbf.newDocumentBuilder().parse(new InputSource(new StringReader(xmlDocument)));
		Document signedDocument = sign(inputDocument, includeKeyInfo, keyEntry);		
		StringWriter stringWriter = new StringWriter();
		TransformerFactory tf = TransformerFactory.newInstance();
		Transformer trans = tf.newTransformer();
		trans.transform(new DOMSource(signedDocument), new StreamResult(stringWriter));
		return stringWriter.getBuffer().toString();
	}
	
	private static Document sign(Document xmlDoc, boolean includeKeyInfo, KeyStore.PrivateKeyEntry keyEntry)throws Exception {
		if (System.getenv("SKIP_DIGITAL_SIGNATURE") != null) {
			return xmlDoc;
		}		
		XMLSignatureFactory factory = XMLSignatureFactory.getInstance("DOM",(Provider) Class.forName(prov).newInstance());
		DigestMethod digestMethod = factory.newDigestMethod(DigestMethod.SHA1, null);
		Transform transform = factory.newTransform(Transform.ENVELOPED, (TransformParameterSpec) null);
		Reference reference = factory.newReference("", digestMethod, Collections.singletonList(transform), null, null);
		CanonicalizationMethod canonicalizationMethod = factory.newCanonicalizationMethod(CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS, (C14NMethodParameterSpec) null);
		SignatureMethod signatureMethod = factory.newSignatureMethod(SignatureMethod.RSA_SHA1, null);
		SignedInfo sInfo = factory.newSignedInfo(canonicalizationMethod, signatureMethod,Collections.singletonList(reference));	
		if (keyEntry == null) {
			throw new RuntimeException("Key could not be read for digital signature. Please check value of signature alias and signature password, and restart the Auth Client");
		}
		X509Certificate x509Cert = (X509Certificate) keyEntry.getCertificate();
		KeyInfo kInfo = getKeyInfo(x509Cert, factory);
		DOMSignContext dsc = new DOMSignContext(keyEntry.getPrivateKey(), xmlDoc.getDocumentElement());
		XMLSignature signature = factory.newXMLSignature(sInfo, includeKeyInfo ? kInfo : null);
		signature.sign(dsc);
		Node node = dsc.getParent();
		return node.getOwnerDocument();
	}	
	
	private static KeyInfo getKeyInfo(X509Certificate cert, XMLSignatureFactory fac) {		
		KeyInfoFactory kif = fac.getKeyInfoFactory();
		List<Object> x509Content = new ArrayList<>();
		x509Content.add(cert.getSubjectX500Principal().getName());
		x509Content.add(cert);
		X509Data xd = kif.newX509Data(x509Content);
		return kif.newKeyInfo(Collections.singletonList(xd));
	}
}
