package src.psis;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.security.interfaces.DSAKey;
import java.security.interfaces.RSAKey;

import javax.security.auth.x500.X500Principal;
import javax.xml.parsers.ParserConfigurationException;

import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Constants;
import org.apache.xmlbeans.XmlException;
import org.apache.xmlbeans.XmlObject;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.xmlsoap.schemas.soap.envelope.Body;
import org.xmlsoap.schemas.soap.envelope.EnvelopeDocument;

import src.psis.utils.AnyTypeHelper;
import src.psis.utils.X500NameUtils;
import src.psis.utils.X500NameUtils.Format;
import x0Assertion.oasisNamesTcSAML1.NameIdentifierType;
import x0CoreSchema.oasisNamesTcDss1.AnyType;
import x0CoreSchema.oasisNamesTcDss1.ClaimedIdentityDocument;
import x0CoreSchema.oasisNamesTcDss1.RequestBaseType;
import x0CoreSchema.oasisNamesTcDss1.ClaimedIdentityDocument.ClaimedIdentity;
import x0ProfilesArchive.oasisNamesTcDss1.ArchiveIdentifierRequest;
import x0CoreSchema.oasisNamesTcDss1.OptionalInputsDocument.OptionalInputs;

/**
 * Sign dss request messages
 * @author Dave Garcia
 *
 */
public class MessageSigner {

	/**
	 * Signs the given request
	 * @param notSignedRequest
	 * @return
	 * @throws ParserConfigurationException 
	 * @throws IOException 
	 * @throws XmlException 
	 */
	public static InputStream signRequest(InputStream notSignedRequest,X509Certificate signingCertificate,PrivateKey signingKey) throws XmlException, IOException, ParserConfigurationException{
		if(!org.apache.xml.security.Init.isInitialized()) {
			org.apache.xml.security.Init.init();
		}
		
		//Recover the request
		EnvelopeDocument envelope=(EnvelopeDocument) XmlObject.Factory.parse(notSignedRequest);
		Body body = envelope.getEnvelope().getBody();
		
		XmlObject request=AnyTypeHelper.getChildren(body)[0];
		request=AnyTypeHelper.detach(request);
		
		checkProfile(request);
		
		//Subject name
		NameIdentifierType ni= NameIdentifierType.Factory.newInstance();
		ni.setFormat("urn:oasis:names:tc:SAML:1.1:nameid-format:X509SubjectName");
		X500Principal subjectX500Principal = signingCertificate.getSubjectX500Principal();
		
		String name;
		try {
			name = X500NameUtils.serialize(subjectX500Principal.getEncoded(),Format.RFC2253,true);
		} catch (Exception e) {
			throw new IOException(e.toString());
		}
		
		ni.setStringValue(name);
		
		//Subject name value
		ClaimedIdentity ci= ClaimedIdentityDocument.Factory.newInstance().addNewClaimedIdentity();
		ci.setName(ni);
		ci.addNewSupportingInfo();
		
		OptionalInputs optionalInputs=prepareOptionalInputs(request);
		
		Element adoptedChild=(Element) optionalInputs.getDomNode().getOwnerDocument().importNode(ci.getDomNode(),true);
		optionalInputs.getDomNode().appendChild(adoptedChild);
		
		//Namespace flushing
		AnyTypeHelper.flush(request);

		Document doc= request.getDomNode().getOwnerDocument();
		
		//We recover the element where the signature will be placed
		Element signatureElement=recoverSupportingInfoElement(request);
		
		sign(doc, signatureElement, signingCertificate,signingKey);
		
		EnvelopeDocument result=EnvelopeDocument.Factory.newInstance();
		Body resultBody = result.addNewEnvelope().addNewBody();
		Element documentElement = doc.getDocumentElement();
		Node importedRequest = ((Document)result.getDomNode()).importNode(documentElement,true);
		resultBody.getDomNode().appendChild(importedRequest);
		
		ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
		result.save(byteArrayOutputStream);
		
		return new ByteArrayInputStream(byteArrayOutputStream.toByteArray());
	}
	
	/**
	 * Creates the optional inputs element if it wasn't created previously 
	 * @param request
	 */	
	private static AnyType prepareOptionalInputs(XmlObject request){
		OptionalInputs optionalInputs;
		if (request instanceof RequestBaseType){
			RequestBaseType castedRequest=(RequestBaseType)request;
			optionalInputs=castedRequest.getOptionalInputs();
			if (optionalInputs==null){
				optionalInputs=castedRequest.addNewOptionalInputs();
			}
		}else if(request instanceof ArchiveIdentifierRequest){
			ArchiveIdentifierRequest castedRequest=(ArchiveIdentifierRequest)request;
			optionalInputs=castedRequest.getOptionalInputs();
			if (optionalInputs==null){
				optionalInputs=castedRequest.addNewOptionalInputs();
			}
		}else{
			throw new RuntimeException("The provided object type cannot be signed using this tool");
		}
		return optionalInputs;
	}
	
	/**
	 * Checks profile related constraints
	 * @param request
	 */
	private static void checkProfile(XmlObject request){
		if (request instanceof RequestBaseType){
			RequestBaseType castedRequest=(RequestBaseType)request;
			if(castedRequest.getProfile()!=null && !castedRequest.getProfile().equals("urn:oasis:names:tc:dss:1.0:profiles:XSS")){
				System.err.println("Signed request is only available at xss profile ");
			}else{
				castedRequest.setProfile("urn:oasis:names:tc:dss:1.0:profiles:XSS");
			}
		}else if(request instanceof ArchiveIdentifierRequest){
			ArchiveIdentifierRequest castedRequest=(ArchiveIdentifierRequest)request;
			if(castedRequest.getProfile()!=null && !castedRequest.getProfile().equals("urn:oasis:names:tc:dss:1.0:profiles:XSS")){
				System.err.println("Signed request is only available at xss profile ");
			}else{
				castedRequest.setProfile("urn:oasis:names:tc:dss:1.0:profiles:XSS");
			}
		}else{
			throw new RuntimeException("The provided object type cannot be signed using this tool");
		}
	}
	
	/**
	 * Signs the given request
	 * @param doc
	 * @param signatureParent
	 * @param signingCertificate
	 * @param privateKey
	 */
	private static void sign(Document doc, Element signatureParent, X509Certificate signingCertificate,PrivateKey privateKey) {
		try {
			String baseURI= "";
			String sigAlgId= null;
			
			if(privateKey instanceof RSAKey) {
				// == RSA-SHA1
				sigAlgId= XMLSignature.ALGO_ID_SIGNATURE_RSA;
			} else if(privateKey instanceof DSAKey) {
				// == DSA-SHA1
				sigAlgId= XMLSignature.ALGO_ID_SIGNATURE_DSA;
			} else {
				throw new RuntimeException("Unexpected Signing Key: " + privateKey.getClass());
			}
			
			XMLSignature xmlSig= new XMLSignature(doc, baseURI, sigAlgId);
	
			org.w3c.dom.Element sigElement = xmlSig.getElement();
			
			signatureParent.appendChild(sigElement);
			
			Transforms transforms = new Transforms(doc);
			transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
			transforms.addTransform(Transforms.TRANSFORM_C14N_EXCL_WITH_COMMENTS);
			xmlSig.addDocument(baseURI, transforms, Constants.ALGO_ID_DIGEST_SHA1);
			
			X509Certificate signingCert= signingCertificate;
			if(signingCert!=null) {
				xmlSig.addKeyInfo(signingCert);
			}
	
			xmlSig.sign(privateKey);
		} catch (Exception e) {
			throw new RuntimeException("Could not sign the Request", e);
		}
	}
	
	/**
	 * Recovers the //dss:ClaimedIdentity/dss:SupportingInfo element
	 */
	private static Element recoverSupportingInfoElement(XmlObject rootElement){
		String query="declare namespace dss='urn:oasis:names:tc:dss:1.0:core:schema'; //dss:ClaimedIdentity/dss:SupportingInfo";
		return (Element) rootElement.selectPath(query)[0].getDomNode();
	}
}
