package net.deterlab.abac;

import java.io.*;
import java.math.*;
import java.text.*;

import java.util.*;
import java.security.*;
import java.security.cert.*;

import javax.security.auth.x500.*;

import javax.xml.parsers.*;
import javax.xml.transform.*;
import javax.xml.transform.dom.*;
import javax.xml.transform.stream.*;

import javax.xml.crypto.*;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.*;
import javax.xml.crypto.dsig.keyinfo.*;
import javax.xml.crypto.dsig.spec.*;

import org.xml.sax.*;
import org.w3c.dom.*;

import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.x509.*;
import org.bouncycastle.x509.util.*;
import org.bouncycastle.openssl.*;

/**
 * An ABAC credential, with or without an underlying certificate that
 * represents it.  These are edges in proof garphs and can be constructed from
 * their constituent Roles.
 * @author <a href="http://abac.deterlab.net">ISI ABAC team</a>
 * @version 1.3
 */
public class GENICredential extends Credential implements Comparable {
    /** The role at the head */
    protected Role m_head
    /** The role at the tail */;
    protected Role m_tail;
    /** The signed XML representing this credential */
    protected Document doc;
    /** The identity that issued the certificate */
    protected Identity id;


    /**
     * Keyselector implementation from the web
     */
    protected class X509KeySelector extends KeySelector {
	public KeySelectorResult select(KeyInfo keyInfo,
					KeySelector.Purpose purpose,
					AlgorithmMethod method,
					XMLCryptoContext context)
	    throws KeySelectorException {
	    Iterator ki = keyInfo.getContent().iterator();
	    while (ki.hasNext()) {
		XMLStructure info = (XMLStructure) ki.next();
		if (!(info instanceof X509Data))
		    continue;
		X509Data x509Data = (X509Data) info;
		Iterator xi = x509Data.getContent().iterator();
		while (xi.hasNext()) {
		    Object o = xi.next();
		    if (!(o instanceof X509Certificate))
			continue;
		    final PublicKey key = ((X509Certificate)o).getPublicKey();
		    // Make sure the algorithm is compatible
		    // with the method.
		    if (algEquals(method.getAlgorithm(), key.getAlgorithm())) {
			return new KeySelectorResult() {
			    public Key getKey() { return key; }
			};
		    }
		}
	    }
	    throw new KeySelectorException("No key found!");
	}

	boolean algEquals(String algURI, String algName) {
	    if ((algName.equalsIgnoreCase("DSA") &&
		algURI.equalsIgnoreCase(SignatureMethod.DSA_SHA1)) ||
		(algName.equalsIgnoreCase("RSA") &&

		algURI.equalsIgnoreCase(SignatureMethod.RSA_SHA1))) {
		return true;
	    } else {
		return false;
	    }
	}
    }

    /**
     * Create an empty Credential.
     */
    public GENICredential() {
	m_head = m_tail = null;
	doc = null;
	id = null;
    }
    /**
     * Create a credential from a head and tail role.  This credential has no
     * underlying certificate, and cannot be exported or used in real proofs.
     * make_cert can create a certificate for a credential initialized this
     * way.
     * @param head the Role at the head of the credential
     * @param tail the Role at the tail of the credential
     */
    public GENICredential(Role head, Role tail) {
        m_head = head;
        m_tail = tail;
	doc = null; 
	id = null;
    }

    /**
     * Do the credential extraction from an input stream.  This parses the
     * certificate from the input stream and saves it. The contents must be
     * validated and parsed later.
     * @param stream the InputStream to read the certificate from.
     * @throws IOException if the stream is unparsable
     */
    protected void read_certificate(InputStream stream) 
	    throws IOException {
	DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
	DocumentBuilder db = null;

	if ( dbf == null ) 
	    throw new IOException("Cannot get DocumentBuilderFactory!?");
	try {
	    dbf.setNamespaceAware(true);
	    if ( (db = dbf.newDocumentBuilder()) == null ) 
		throw new IOException("Cannot get DocumentBuilder!?");
	    doc = db.parse(stream);
	    doc.normalizeDocument();
	}
	catch (IllegalArgumentException ie) {
	    throw new IOException("null stream");
	}
	catch (SAXException se) {
	    throw new IOException(se.getMessage());
	}
	catch (Exception e) {
	    throw new IOException(e.getMessage());
	}

    }

    /**
     * Walk the document and set all instances of an xml:id attribute to be ID
     * attributes in terms of the java libraries.  This is needed for the
     * signature code to find signed subsections.  The "xml:id" is treated as
     * both a namespace free ID with a colon and as an "id" identifier in the
     * "xml" namespace so we don't miss it.
     * @param n the root of the document to mark
     */
    static protected void setIDAttrs(Node n) {
	if ( n.getNodeType() == Node.ELEMENT_NODE) {
	    Element e = (Element) n;
	    String id = e.getAttribute("xml:id");

	    if ( id != null && id.length() > 0 )
		e.setIdAttribute("xml:id", true);

	    id = e.getAttributeNS("xml", "id");
	    if ( id != null && id.length() > 0 )
		e.setIdAttributeNS("xml", "id", true);
	}

	for (Node nn = n.getFirstChild(); nn != null; nn = nn.getNextSibling())
	    setIDAttrs(nn);
    }

    /**
     * Return the child of Node n  that has the given name, if there is one.
     * @param n a Node to search
     * @param name the name to search for
     * @return a Node with the name, may be null
     */
    static protected Node getChildByName(Node n, String name) {
	if (name == null ) return null;

	for (Node nn = n.getFirstChild(); nn != null; nn = nn.getNextSibling())
	    if ( nn.getNodeType() == Node.ELEMENT_NODE && 
		    name.equals(nn.getNodeName())) return nn;
	return null;
    }

    /**
     * Find the X509Certificate in the Signature element and convert it to an
     * ABAC identity.  This assumes a KeyInfo in that section holding an
     * X509Data section containing the certificate in an X509Certificate
     * section.  Any of that not being found will cause a failure.  If the
     * Identity cannot be created a Gignature exception is thrown.
     * @param n a Node pointing to a Signature section
     * @return an Identity constructed frm the X509 Certificate
     * @throws SignatureException if the Identity cannot be created from the
     * certificate.
     */
    static protected Identity getIdentity(Node n) throws SignatureException {
	Identity rv = null;
	Node nn = getChildByName(n, "KeyInfo");
	String certStr = null;

	if ( nn == null ) return null;
	if ( ( nn = getChildByName(nn, "X509Data")) == null ) return null;
	if ( ( nn = getChildByName(nn, "X509Certificate")) == null ) return null;
	if ( ( certStr = nn.getTextContent()) == null ) return null;
	try {
	    certStr = "-----BEGIN CERTIFICATE-----\n" + 
		certStr + 
		"\n-----END CERTIFICATE-----";
	    return new Identity(new StringReader(certStr));
	} 
	catch (Exception e) {
	    throw new SignatureException(e.getMessage());
	}
    }

    /**
     * Initialize a credential from parsed certificate.  Validiate it against
     * the given identities and parse out the roles.  Note that catching
     * java.security.GeneralSecurityException catches most of the exceptions
     * this throws.
     * @param ids a Collection of Identities to use in validating the cert
     * @throws CertificateException if the certificate is badly formatted
     * @throws InvalidKeyException if none of the Identities can validate the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature check fails
     */
    protected void init(Collection<Identity> ids) 
	    throws CertificateException, InvalidKeyException, 
		NoSuchAlgorithmException, NoSuchProviderException,
		SignatureException {

	XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
	NodeList nl = doc.getElementsByTagNameNS(XMLSignature.XMLNS, 
	 	"Signature");
	DOMValidateContext valContext = null;
	XMLSignature signature = null;

	try {

	    if (nl.getLength() == 0) 
		throw new SignatureException("Cannot find Signature element");

	    setIDAttrs(doc);
	    valContext = new DOMValidateContext(new X509KeySelector(),
		    nl.item(0));
	    if ( valContext == null )
		throw new SignatureException("No context");

	    signature = fac.unmarshalXMLSignature(valContext);
	    if (signature == null) 
		throw new SignatureException("No sig");

	    if (!signature.validate(valContext))
		throw new SignatureException("bad signature");

	    load_roles();

	    id = getIdentity(nl.item(0));

	    if ( !ids.contains(id) ) ids.add(id);

	    if (!id.getKeyID().equals(m_head.principal()))
		throw new InvalidKeyException("Unknown identity");
	}
	catch (Exception e) {
	    throw new SignatureException(e.getMessage());
	}
    }

    /**
     * Parse a credential from an InputStream and initialize the role from it.
     * Combine read_credential(stream) and init(ids).  Note that catching
     * java.security.GeneralSecurityException catches most of the exceptions
     * this throws.
     * @param stream the InputStream to read the certificate from.
     * @param ids a Collection of Identities to use in validating the cert
     * @throws CertificateException if the certificate is badly formatted
     * @throws InvalidKeyException if none of the Identities can validate the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature check fails
     * @throws IOException if the certificate is unparsable.
     */
    protected void init(InputStream stream, Collection<Identity> ids) 
	    throws CertificateException, InvalidKeyException, 
		NoSuchAlgorithmException, NoSuchProviderException,
		SignatureException, IOException {
	read_certificate(stream);
	if (doc == null) throw new IOException("Unknown Format");
	init(ids);
    }

    /**
     * Create a credential from an attribute cert in a file. Throws an
     * exception if the cert file can't be opened or if there's a format
     * problem with the cert.  Note that catching
     * java.security.GeneralSecurityException catches most of the exceptions
     * this throws.
     * @param filename a String containing the filename to read
     * @param ids a Collection of Identities to use in validating the cert
     * @throws StreamParsingException if the stream is unparsable
     * @throws CertificateException if the certificate is badly formatted
     * @throws InvalidKeyException if none of the Identities can validate the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature check fails
     * @throws IOException if the certificate is unparsable.
     */
    public GENICredential(String filename, Collection<Identity> ids) 
	throws Exception { init(new FileInputStream(filename), ids); }

    /**
     * Create a credential from an attribute cert in a file. Throws an
     * exception if the cert file can't be opened or if there's a format
     * problem with the cert.  Note that catching
     * java.security.GeneralSecurityException catches most of the exceptions
     * this throws.
     * @param file the File to read
     * @param ids a Collection of Identities to use in validating the cert
     * @throws StreamParsingException if the stream is unparsable
     * @throws CertificateException if the certificate is badly formatted
     * @throws InvalidKeyException if none of the Identities can validate the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature check fails
     * @throws IOException if the certificate is unparsable.
     */
    public GENICredential(File file, Collection<Identity> ids) 
	    throws CertificateException, InvalidKeyException, 
		NoSuchAlgorithmException, NoSuchProviderException,
		SignatureException, StreamParsingException, IOException {
	init(new FileInputStream(file), ids);
    }

    /**
     * Create a credential from an InputStream.  Throws an exception if the
     * stream can't be parsed or if there's a format problem with the cert.
     * Note that catching java.security.GeneralSecurityException catches most
     * of the exceptions this throws.
     * @param s the InputStream to read
     * @param ids a Collection of Identities to use in validating the cert
     * @throws StreamParsingException if the stream is unparsable
     * @throws CertificateException if the certificate is badly formatted
     * @throws InvalidKeyException if none of the Identities can validate the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature check fails
     * @throws IOException if the certificate is unparsable.
     */
    public GENICredential(InputStream s, Collection<Identity> ids) 
	    throws CertificateException, InvalidKeyException, 
		NoSuchAlgorithmException, NoSuchProviderException,
		SignatureException, StreamParsingException, IOException {
	init(s, ids);
    }

    /**
     * Create a certificate from this credential issued by the given identity.
     * Note that catching java.security.GeneralSecurityException catches most
     * of the exceptions this throws.
     * @param i the Identity that will issue the certificate
     * @throws IOException reading or writing problems
     * @throws CertificateEncodingException Problem creating certificate
     * @throws InvalidKeyException if none of the Identities can sign the
     *				    certificate
     * @throws NoSuchAlgorithmException if the credential uses an unknown
     *				    signature algorithm
     * @throws NoSuchProviderException if the provider of the signature
     *				    algorithm is unavailable
     * @throws SignatureException if the signature creation fails
     */
    public void make_cert(Identity i) 
	    throws IOException, CertificateEncodingException,
	       NoSuchProviderException, NoSuchAlgorithmException,
	       SignatureException, InvalidKeyException {
	try {
	    SimpleDateFormat df = new SimpleDateFormat("yyy-MM-dd'T'HH:mm:ss'Z'");
	    StringBuffer expBuf = new StringBuffer();
	    X509Certificate cert = i.getCertificate();
	    KeyPair kp = i.getKeyPair();
	    if ( cert == null ) 
		throw new IOException("No credential in identity?");
	    if ( kp == null ) 
		throw new IOException("No keypair in identity?");

	    XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
	    Reference ref = fac.newReference("#ref0", 
		    fac.newDigestMethod(DigestMethod.SHA1, null),
		    Collections.singletonList(
			fac.newTransform(Transform.ENVELOPED, 
			    (TransformParameterSpec) null)),
		    null, null);

	    SignedInfo si = fac.newSignedInfo(
		    fac.newCanonicalizationMethod(
			CanonicalizationMethod.INCLUSIVE,
			(C14NMethodParameterSpec) null),
		    fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
		    Collections.singletonList(ref));
	    KeyInfoFactory kif = fac.getKeyInfoFactory();
	    List x509Content = new ArrayList();
	    x509Content.add(cert.getSubjectX500Principal().getName());
	    x509Content.add(cert);
	    X509Data xd = kif.newX509Data(x509Content);
	    KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));


	    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
	    DocumentBuilder db = null;

	    if ( dbf == null ) 
		throw new IOException("Cannot get DocumentBuilderFactory!?");
	    if ( (db = dbf.newDocumentBuilder()) == null ) 
		throw new IOException("Cannot get DocumentBuilder!?");
	    doc = db.newDocument();
	    if ( doc == null ) 
		throw new IOException("No Document");
	    Element root = doc.createElement("signed-credential"); 
	    doc.appendChild(root);
	    Element cred = doc.createElement("credential");
	    Element sig = doc.createElement("credential");
	    cred.setAttributeNS("xml", "id", "ref0");
	    root.appendChild(cred);
	    cred.appendChild(doc.createElement("type").appendChild(
			doc.createTextNode("abac")));
	    df.format(new Date(System.currentTimeMillis() 
			+ (3600L * 1000L * 24L * 365L)), expBuf, 
		    new FieldPosition(0));
	    cred.appendChild(doc.createElement("expires").appendChild(
			doc.createTextNode(expBuf.toString())));
	    cred.appendChild(doc.createElement("version").appendChild(
			doc.createTextNode("1.0")));
	    cred.appendChild(doc.createElement("rt0").appendChild(
			doc.createTextNode(m_head + "<-" + m_tail)));
	    root.appendChild(sig);

	    DOMSignContext dsc = new DOMSignContext(kp.getPrivate(), sig);
	    XMLSignature signature = fac.newXMLSignature(si, ki);
	    signature.sign(dsc);
	}
	catch (Exception e) {
	    throw new SignatureException(e.getMessage());
	}

    }

    /**
     * Load the roles off the attribute cert. Throws a RuntimeException if
     * there's something wrong with the cert.
     */
    private void load_roles() {
	if ( doc == null ) 
            throw new RuntimeException("No credential");

        String roles = null;
	NodeList nodes = doc.getElementsByTagName("rt0");

	if (nodes == null || nodes.getLength() != 1) 
            throw new RuntimeException("More than one rt0 element?");

	Node node = nodes.item(0);

	if ( node == null ) 
            throw new RuntimeException("bad rt0 element?");

	if ( (roles = node.getTextContent()) == null ) 
            throw new RuntimeException("bad rt0 element (no text)?");

        String[] parts = roles.split("\\s*<--?\\s*");
        if (parts.length != 2)
            throw new RuntimeException("Invalid attribute: " + roles);

        m_head = new Role(parts[0]);
        m_tail = new Role(parts[1]);
    }

    /**
     * Two credentials are the same if their roles are the same.
     * @param o an Object to compare
     * @return true if the Credentials have the Roles 
     */
    public boolean equals(Object o) {
	if ( o instanceof Credential ) {
	    Credential c = (Credential) o;

	    if (m_head == null || m_tail == null ) return false;
	    else return (m_head.equals(c.head()) && m_tail.equals(c.tail()));
	}
	else return false;
    }

    /**
     * Allow credentials to be compared.  They are ordered by their Roles, head
     * then tail.
     * @param o an Object to compare
     * @return -1 if this Credential is before, 0 if they are the same, and 1
     *		    if this Credential is after the given object.
     */
    public int compareTo(Object o) {
	if (o instanceof Credential) {
	    Credential c = (Credential) o;

	    if (head().equals(c.head())) return tail().compareTo(c.tail());
	    else return head().compareTo(c.head());
	}
	else return 1;
    }


    /**
     * Get the head role from the credential.
     * @return the Role in the head
     */
    public Role head() { return m_head; }

    /**
     * Get the tail role from the credential
     * @return the Role in the tail
     */
    public Role tail() { return m_tail; }

    /**
     * Gets the cert associated with this credential (if any).
     * @return the attribute cert associated with this credential (if any).
     */
    public Document GENIcert() { return doc; }

    /**
     * Turn the credential into string form. The format is head &lt;- tail. For
     * example: A.r1 &lt;- B.r2.r3.  Principal names are key identifiers.
     * @return the string form
     */
    public String toString() {
        return m_head + " <- " + m_tail;
    }

    /**
     * Turn the credential into string form. The format is head &lt;- tail. For
     * example: A.r1 &lt;- B.r2.r3.  Principal names are shortened to menmonics
     * if the Context knows the identity.
     * @param c the Context to translate names in
     * @return the string form
     */
    public String simpleString(Context c) {
	return m_head.simpleString(c) + " <- " + m_tail.simpleString(c);
    }

    /**
     * Output the signed GENI ABAC certificate associated with this
     * Credential to the OutputStream.
     * @param s the OutputStream on which to write
     * @throws IOException if there is an error writing.
     */
    public void write(OutputStream s) throws IOException {
	if ( doc == null ) 
	    return;
	try {
	    TransformerFactory tf = TransformerFactory.newInstance();
	    Transformer t = tf.newTransformer();
	    DOMSource source = new DOMSource(doc);
	    StreamResult result = new StreamResult(s);

	    t.transform(source, result);
	    s.flush();
	} 
	catch (Exception e) {
	    throw new IOException(e.getMessage());
	}
    }

    /**
     * Output the signed GENI ABAC certificate associated with this
     * Credential to the filename given.
     * @param fn a String containing the output filename
     * @throws IOException if there is an error writing.
     */
    public void write(String fn) throws IOException, FileNotFoundException {
	write(new FileOutputStream(fn));
    }

    /**
     * Return true if this Credential has a certificate associated.  A jabac
     * extension.
     * @return true if this Credential has a certificate associated.
     */
    public boolean hasCertificate() { return doc != null; }

    /**
     * Return the Identity that issued the underlying certificate.  A jabac
     * extension.
     * @return the Identity that issued the underlying certificate.
     */
    public Identity issuer() { return id; }
    /**
     * Return the X509Certificate that issued the underlying certificate.
     * @return the X509Certificate that issued the underlying certificate.
     */
    public X509Certificate issuerCert() { return id.getCertificate(); }
    /**
     * Return the X509V2AttributeCertificate that underlys the Credential
     * @return the X509V2AttributeCertificate that underlys the Credential.
    public X509V2AttributeCertificate attributeCert() { return ac; }
     */


}