package net.deterlab.abac;

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

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

import net.deterlab.abac.Identity;

import org.bouncycastle.asn1.*;
import org.bouncycastle.x509.*;
import org.bouncycastle.jce.X509Principal;
import org.bouncycastle.jce.provider.X509AttrCertParser;
import org.bouncycastle.jce.provider.X509CertificateObject;
import org.bouncycastle.openssl.PEMReader;

import org.bouncycastle.asn1.util.ASN1Dump;

import java.security.PrivateKey;

public class Credential implements Comparable {
    protected static Vector<Identity> s_ids = new Vector<Identity>();
    protected static String attrOID = "1.3.6.1.5.5.7.10.4";

    /**
     * A dummy credential.
     */
    public Credential() {
	m_head = m_tail = null;
	m_ac = null;
	m_id = null;
    }
    /**
     * Create a credential from a head and tail role. This is only for testing.
     * In a real implementation the Credential must be loaded from an X.509
     * attribute cert.
     */
    public Credential(Role head, Role tail) {
        m_head = head;
        m_tail = tail;
	m_ac = null; 
	m_id = null;
    }

    /**
     * Do the credential initialization from a filename.
     */
    protected void init(InputStream stream) throws Exception {
        X509AttrCertParser parser = new X509AttrCertParser();
        parser.engineInit(stream);
        m_ac = (X509V2AttributeCertificate)parser.engineRead();
	m_id = null;

	if ( m_ac == null ) throw new IOException("Invalid Credential Format");

	for (Identity id: s_ids) {
	    try {
		m_ac.verify(id.getCertificate().getPublicKey(), "BC");
		m_id = id;
		break;
	    }
	    catch (InvalidKeyException e) { }
	}
	if (m_id == null) throw new InvalidKeyException("Unknown identity");

        load_roles();

	if (!m_id.getKeyID().equals(m_head.issuer_part()))
	    throw new InvalidKeyException("Unknown identity");
    }

    /**
     * Create a credential from an attribute cert. Throws an exception if the
     * cert file can't be opened or if there's a format problem with the cert.
     */
    public Credential(String filename) throws Exception {
	init(new FileInputStream(filename));
    }

    /**
     * Create a credential from an attribute cert. Throws an exception if the
     * cert file can't be opened or if there's a format problem with the cert.
     */
    public Credential(File file) throws Exception {
	init(new FileInputStream(file));
    }

    /**
     * Create a credential from an InputStream.
     */
    public Credential(InputStream s) throws Exception { 
	init(s);
    }

    public void make_cert(PrivateKey key) {
	X509V2AttributeCertificateGenerator gen = 
	    new X509V2AttributeCertificateGenerator();

	gen.setIssuer(new AttributeCertificateIssuer(
		    new X509Principal("CN="+m_head.issuer_part())));
	gen.setHolder(new AttributeCertificateHolder(
		    new X509Principal("CN="+m_head.issuer_part())));
	gen.setNotAfter(new Date(System.currentTimeMillis() 
		    + 3600 * 1000 * 24 * 365));
	gen.setNotBefore(new Date(System.currentTimeMillis()));
	gen.setSerialNumber(BigInteger.valueOf(System.currentTimeMillis()));
	gen.addAttribute(new X509Attribute(attrOID, 
		    new DERSequence(
			new DERSequence(
			    new DERUTF8String(toString())))));
	gen.setSignatureAlgorithm("SHA256WithRSAEncryption");

	try { 
	    m_ac = (X509V2AttributeCertificate) gen.generate(key, "BC");
	}
	catch (Exception e) { 
	    System.err.println(e);
	}
    }

    /**
     * Load the roles off the attribute cert. Throws a RuntimeException if
     * there's something wrong with the cert.
     */
    private void load_roles() throws RuntimeException {
        String roles = null;
        try {
            X509Attribute attr = m_ac.getAttributes()[0];

            DERSequence    java     = (DERSequence)attr.getValues()[0];
            DERSequence    fucking  = (DERSequence)java.getObjectAt(0);
            DERUTF8String  sucks    = (DERUTF8String)fucking.getObjectAt(0);

            roles = sucks.getString();
        }
        catch (Exception e) {
            throw new RuntimeException("Your attribute certificate is funky and I'm not gonna debug it", e);
        }

        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.
     */
    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;
    }

    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.
     */
    public Role head() {
        return m_head;
    }

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

    /**
     * Gets the cert associated with this credential (if any).
     */
    public X509V2AttributeCertificate cert() {
        return m_ac;
    }

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

    public String simpleString() {
	return m_head.simpleString() + " <- " + m_tail.simpleString();
    }

    public void write(OutputStream s) throws IOException {
	s.write(m_ac.getEncoded());
    }

    public void write(String fn) throws IOException, FileNotFoundException {
	write(new FileOutputStream(fn));
    }

    public boolean hasCertificate() { return m_ac != null; }

    public Identity getID() { return m_id; }

    /**
     * Import a zip file.  First import all the identities
     * (pem), then the credentials (der) into the credential graph then any
     * alias files into the two maps.  If keys is not null, any key pairs in
     * PEM files are put in there.  If errors is not null, errors reading files
     * are added indexed by filename.
     */
    static public Collection<Credential> readZipFile(File zf, 
	    Collection<KeyPair> keys, Map<String, Exception> errors) 
		throws IOException {
	Vector<Credential> creds = new Vector<Credential>();
	Vector<ZipEntry> derEntries = new Vector<ZipEntry>();
	Map<String, Identity> ids = new TreeMap<String, Identity>();
	Map<String, KeyPair> kps = new TreeMap<String, KeyPair>();

	ZipFile z = new ZipFile(zf);

	for (Enumeration<? extends ZipEntry> ze = z.entries(); 
		ze.hasMoreElements();) {
	    ZipEntry  f = ze.nextElement();
	    try {
		PEMReader r = new PEMReader(
			new InputStreamReader(z.getInputStream(f)));
		Object o = readPEM(r);

		if ( o != null ) {
		    if (o instanceof Identity) {
			Identity i = (Identity) o;
			String kid = i.getKeyID();

			if (kps.containsKey(kid) ) {
			    i.setKeyPair(kps.get(kid));
			    kps.remove(kid);
			}
			else if (i.getKeyPair() == null ) 
			    ids.put(i.getKeyID(), i);

			Credential.addIdentity(i);
		    }
		    else if (o instanceof KeyPair ) {
			KeyPair kp = (KeyPair) o;
			String kid = Identity.extractKeyID(kp.getPublic());

			if (ids.containsKey(kid)) {
			    Identity i = ids.get(kid);

			    i.setKeyPair(kp);
			    ids.remove(kid);
			}
			else {
			    kps.put(kid, kp);
			}
		    }
		}
		else {
		    // Not a PEM file
		    derEntries.add(f);
		    continue;
		}
	    }
	    catch (Exception e ) {
		if (errors != null ) errors.put(f.getName(), e);
	    }
	}

	for ( ZipEntry f : derEntries ) {
	    try {
		creds.add(new Credential(z.getInputStream(f)));
	    }
	    catch (Exception e ) {
		if (errors != null ) errors.put(f.getName(), e);
	    }
	}
	return creds;
    }

    static public Collection<Credential> readZipFile(File d) 
	    throws IOException {
	return readZipFile(d, null, null);
    }
    static public Collection<Credential> readZipFile(File d, 
	    Map<String, Exception> errors) throws IOException {
	return readZipFile(d, null, errors);
    }
    static public Collection<Credential> readZipFile(File d, 
	    Collection<KeyPair> keys) throws IOException {
	return readZipFile(d, keys, null);
    }

    protected static Object readPEM(PEMReader r) throws IOException {
	Identity i = null;
	KeyPair keys = null;
	Object o = null;

	while ( (o = r.readObject()) != null ) {
	    if (o instanceof X509CertificateObject) {
		if ( i == null ) {
		    try {
			i = new Identity((X509CertificateObject)o);
		    }
		    catch (Exception e) {
			// Translate Idenitiy exceptions to IOException
			throw new IOException(e);
		    }
		    if (keys != null ) {
			i.setKeyPair(keys);
			keys = null;
		    }
		}
		else throw new IOException("Two certificates");
	    }
	    else if (o instanceof KeyPair ) {
		if ( i != null ) i.setKeyPair((KeyPair) o);
		else keys = (KeyPair) o;
	    }
	    else {
		throw new IOException("Unexpected PEM object: " + 
			o.getClass().getName());
	    }
	}

	if ( i != null ) return i;
	else if ( keys != null) return keys;
	else return null;
    }

    /**
     * Import a directory full of files.  First import all the identities
     * (pem), then the credentials (der) into the credential graph then any
     * alias files into the two maps.  If keys is not null, any key pairs in
     * PEM files are put in there.  If errors is not null, errors reading files
     * are added indexed by filename.
     */
    static public Collection<Credential> readDirectory(File d, 
	    Collection<KeyPair> keys, Map<String, Exception> errors) {
	Vector<Credential> creds = new Vector<Credential>();
	Vector<File> derFiles = new Vector<File>();
	Collection<File> files = new Vector<File>();
	Map<String, Identity> ids = new TreeMap<String, Identity>();
	Map<String, KeyPair> kps = new TreeMap<String, KeyPair>();

	if (d.isDirectory() ) 
	    for (File f : d.listFiles()) 
		files.add(f);
	else files.add(d);

	for (File f: files ) {
	    try {
		PEMReader r = new PEMReader(new FileReader(f));
		Object o = readPEM(r);

		if ( o != null ) {
		    if (o instanceof Identity) {
			Identity i = (Identity) o;
			String kid = i.getKeyID();

			if (kps.containsKey(kid) ) {
			    i.setKeyPair(kps.get(kid));
			    kps.remove(kid);
			}
			else if (i.getKeyPair() == null ) 
			    ids.put(i.getKeyID(), i);

			Credential.addIdentity(i);
		    }
		    else if (o instanceof KeyPair ) {
			KeyPair kp = (KeyPair) o;
			String kid = Identity.extractKeyID(kp.getPublic());

			if (ids.containsKey(kid)) {
			    Identity i = ids.get(kid);

			    i.setKeyPair(kp);
			    ids.remove(kid);
			}
			else {
			    kps.put(kid, kp);
			}
		    }
		}
		else {
		    // Not a PEM file
		    derFiles.add(f);
		    continue;
		}
	    }
	    catch (Exception e ) {
		if (errors != null ) errors.put(f.getName(), e);
	    }
	}

	for ( File f : derFiles ) {
	    try {
		creds.add(new Credential(f));
	    }
	    catch (Exception e ) {
		if (errors != null ) errors.put(f.getName(), e);
	    }
	}
	return creds;
    }

    static public Collection<Credential> readDirectory(File d) {
	return readDirectory(d, null, null);
    }
    static public Collection<Credential> readDirectory(File d, 
	    Map<String, Exception> errors) {
	return readDirectory(d, null, errors);
    }
    static public Collection<Credential> readDirectory(File d, 
	    Collection<KeyPair> keys) {
	return readDirectory(d, keys, null);
    }

    static public void writeZipFile(Collection<Credential> creds, File f,
	    boolean allIDs) 
	    throws IOException {
	ZipOutputStream z = new ZipOutputStream(new FileOutputStream(f));
	Set<Identity> ids = allIDs ? 
	    new TreeSet<Identity>(s_ids) : new TreeSet<Identity>();

	int n = 0;
	for (Credential c: creds) {
	    z.putNextEntry(new ZipEntry("attr" + n++  + ".der"));
	    c.write(z);
	    z.closeEntry();
	    if ( c.getID() != null && !allIDs) ids.add(c.getID());
	}
	for (Identity i: ids) {
	    z.putNextEntry(new ZipEntry(i.getName() + ".pem"));
	    i.write(z);
	    z.closeEntry();
	}
	z.close();
    }


private Role m_head, m_tail;

private X509V2AttributeCertificate m_ac;
private Identity m_id;

    /**
     * Put the Identity into the set of ids used to validate certificates.
     * Also put the keyID and name into the translation mappings used by Roles
     * to pretty print.  In the role mapping, if multiple ids use the same
     * common name they are disambiguated.  Only one entry for keyid is
     * allowed.
     */
    public static void addIdentity(Identity id) { 
	s_ids.add(id);
	if (id.getName() != null && id.getKeyID() != null) {
	    if ( !Role.key_in_mapping(id.getKeyID()) ) {
		String name = id.getName();
		int n= 1;

		while (Role.name_in_mapping(name)) {
		    name = id.getName() + n++;
		}
		Role.add_mapping(name, id.getKeyID());
	    }
	}
    }
    public static Collection<Identity> identities() { return s_ids; }
    public static void clearIdentities() {
	s_ids.clear(); Role.clear_mapping();
    }
}