package org.bouncycastle.ocsp;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.Iterator;

import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1OutputStream;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERBitString;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.DERObject;
import org.bouncycastle.asn1.DERObjectIdentifier;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.ocsp.OCSPRequest;
import org.bouncycastle.asn1.ocsp.Request;
import org.bouncycastle.asn1.ocsp.Signature;
import org.bouncycastle.asn1.ocsp.TBSRequest;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.GeneralName;
import org.bouncycastle.asn1.x509.X509CertificateStructure;
import org.bouncycastle.asn1.x509.X509Extensions;

public class OCSPReqGenerator
{
    private ArrayList       list = new ArrayList();
    private GeneralName     requestorName = null;
    private X509Extensions  requestExtensions = null;

    private static Hashtable            algorithms = new Hashtable();

    static
    {
        algorithms.put("MD2WITHRSAENCRYPTION", new DERObjectIdentifier("1.2.840.113549.1.1.2"));
        algorithms.put("MD2WITHRSA", new DERObjectIdentifier("1.2.840.113549.1.1.2"));
        algorithms.put("MD5WITHRSAENCRYPTION", new DERObjectIdentifier("1.2.840.113549.1.1.4"));
        algorithms.put("MD5WITHRSA", new DERObjectIdentifier("1.2.840.113549.1.1.4"));
        algorithms.put("SHA1WITHRSAENCRYPTION", new DERObjectIdentifier("1.2.840.113549.1.1.5"));
        algorithms.put("SHA1WITHRSA", new DERObjectIdentifier("1.2.840.113549.1.1.5"));
        algorithms.put("RIPEMD160WITHRSAENCRYPTION", new DERObjectIdentifier("1.3.36.3.3.1.2"));
        algorithms.put("RIPEMD160WITHRSA", new DERObjectIdentifier("1.3.36.3.3.1.2"));
        algorithms.put("SHA1WITHDSA", new DERObjectIdentifier("1.2.840.10040.4.3"));
        algorithms.put("DSAWITHSHA1", new DERObjectIdentifier("1.2.840.10040.4.3"));
        algorithms.put("SHA1WITHECDSA", new DERObjectIdentifier("1.2.840.10045.4.1"));
        algorithms.put("ECDSAWITHSHA1", new DERObjectIdentifier("1.2.840.10045.4.1"));
    }

    private class RequestObject
    {
        CertificateID   certId;
        X509Extensions  extensions;

        public RequestObject(
            CertificateID   certId,
            X509Extensions  extensions)
        {
            this.certId = certId;
            this.extensions = extensions;
        }

        public Request toRequest()
            throws Exception
        {
            return new Request(certId.toASN1Object(), extensions);
        }
    }

    private DERObject makeObj(
        byte[]  encoding)
        throws IOException
    {
        if (encoding == null)
        {
            return null;
        }

        ByteArrayInputStream    bIn = new ByteArrayInputStream(encoding);
        ASN1InputStream         aIn = new ASN1InputStream(bIn);

        return aIn.readObject();
    }

    /**
     * add a request for the given issuerCert using the given hash algorithm.
     */
    public void addRequest(
        CertificateID   certId)
    {
        list.add(new RequestObject(certId, null));
    }

    public void addRequest(
        CertificateID   certId,
        X509Extensions  requestExtensions)
    {
        list.add(new RequestObject(certId, requestExtensions));
    }

    public void setRequestorName(
        GeneralName         requestorName)
    {
        this.requestorName = requestorName;
    }

    public void setRequestExtensions(
        X509Extensions      requestExtensions)
    {
        this.requestExtensions = requestExtensions;
    }

    private OCSPReq generateRequest(
        DERObjectIdentifier signingAlgorithm,
        PrivateKey          key,
        X509Certificate[]   chain,
        String              provider,
        SecureRandom        random)
        throws OCSPException, NoSuchProviderException
    {
        Iterator    it = list.iterator();

        ASN1EncodableVector requests = new ASN1EncodableVector();

        while (it.hasNext())
        {
            try
            {
                requests.add(((RequestObject)it.next()).toRequest());
            }
            catch (Exception e)
            {
                throw new OCSPException("exception creating Request", e);
            }
        }

        TBSRequest  tbsReq = new TBSRequest(requestorName, new DERSequence(requests), requestExtensions);

        java.security.Signature sig = null;
        Signature               signature = null;

        if (signingAlgorithm != null)
        {
            try
            {
                sig = java.security.Signature.getInstance(signingAlgorithm.getId(), provider);
                if (random != null)
                {
                    sig.initSign(key, random);
                }
                else
                {
                    sig.initSign(key);
                }

            }
            catch (NoSuchAlgorithmException e)
            {
                throw new OCSPException("exception creating signature: " + e, e);
            }
            catch (InvalidKeyException e)
            {
                throw new OCSPException("exception creating signature: " + e, e);
            }

            DERBitString    bitSig = null;

            try
            {
                ByteArrayOutputStream   bOut = new ByteArrayOutputStream();
                ASN1OutputStream        aOut = new ASN1OutputStream(bOut);

                aOut.writeObject(tbsReq);

                sig.update(bOut.toByteArray());

                bitSig = new DERBitString(sig.sign());
            }
            catch (Exception e)
            {
                throw new OCSPException("exception processing TBSRequest: " + e, e);
            }

            AlgorithmIdentifier sigAlgId = new AlgorithmIdentifier(signingAlgorithm, new DERNull());

            if (chain != null && chain.length > 0)
            {
                ASN1EncodableVector v = new ASN1EncodableVector();
                try
                {
                    for (int i = 0; i != chain.length; i++)
                    {
                        v.add(new X509CertificateStructure(
                                (ASN1Sequence)makeObj(chain[i].getEncoded())));
                    }
                }
                catch (IOException e)
                {
                    throw new OCSPException("error processing certs", e);
                }
                catch (CertificateEncodingException e)
                {
                    throw new OCSPException("error encoding certs", e);
                }

                signature = new Signature(sigAlgId, bitSig, new DERSequence(v));
            }
            else
            {
                signature = new Signature(sigAlgId, bitSig);
            }
        }

        return new OCSPReq(new OCSPRequest(tbsReq, signature));
    }
    
    public OCSPReq generate()
        throws OCSPException
    {
        try
        {
            return generateRequest(null, null, null, null, null);
        }
        catch (NoSuchProviderException e)
        {
            //
            // this shouldn't happen but...
            //
            throw new OCSPException("no provider! - " + e, e);
        }
    }

    public OCSPReq generate(
        String              signingAlgorithm,
        PrivateKey          key,
        X509Certificate[]   chain,
        String              provider)
        throws OCSPException, NoSuchProviderException, IllegalArgumentException
    {
        return generate(signingAlgorithm, key, chain, provider, null);
    }

    public OCSPReq generate(
        String              signingAlgorithm,
        PrivateKey          key,
        X509Certificate[]   chain,
        String              provider,
        SecureRandom        random)
        throws OCSPException, NoSuchProviderException, IllegalArgumentException
    {
        if (signingAlgorithm == null)
        {
            throw new IllegalArgumentException("no signing algorithm specified");
        }

        DERObjectIdentifier oid = (DERObjectIdentifier)algorithms.get(signingAlgorithm.toUpperCase());
        if (oid == null)
        {
            throw new IllegalArgumentException("unknown signing algorithm specified");
        }

        return generateRequest(oid, key, chain, provider, random);
    }
}
