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.PublicKey;
import java.security.SecureRandom;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Date;
import java.util.Hashtable;
import java.util.Iterator;

import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.ocsp.BasicOCSPResponse;
import org.bouncycastle.asn1.ocsp.CertStatus;
import org.bouncycastle.asn1.ocsp.ResponseData;
import org.bouncycastle.asn1.ocsp.RevokedInfo;
import org.bouncycastle.asn1.ocsp.SingleResponse;
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 BasicOCSPRespGenerator
{
    private ArrayList       list = new ArrayList();
    private GeneralName     responderName = null;
    private X509Extensions  responseExtensions = null;
    private RespID          responderID;

    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 ResponseObject
    {
        CertificateID       certId;
        CertStatus          certStatus;
        DERGeneralizedTime  nextUpdate;
        X509Extensions      extensions;

        public ResponseObject(
            CertificateID       certId,
            CertificateStatus   certStatus,
            Date                nextUpdate,
            X509Extensions      extensions)
        {
            this.certId = certId;

            if (certStatus == null)
            {
                this.certStatus = new CertStatus();
            }
            else 
            {
                this.certStatus = new CertStatus(new RevokedInfo(
                                    new DERGeneralizedTime(new Date()), null));
            }

            if (nextUpdate != null)
            {
                this.nextUpdate = new DERGeneralizedTime(nextUpdate);
            }
            else
            {
                this.nextUpdate = null;
            }

            this.extensions = extensions;
        }

        public SingleResponse toResponse()
            throws Exception
        {
            return new SingleResponse(certId.toASN1Object(), certStatus, new DERGeneralizedTime(new Date()), nextUpdate, 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();
    }

    /**
     * set the responderID.
     */
    public void setResponderID(
        RespID  responderID)
    {
        this.responderID = responderID;
    }

    /**
     * set the responderID to be the SHA-1 keyHash of the passed in public key.
     */
    public void setResponderID(
        PublicKey       key)
        throws OCSPException
    {
        this.responderID = new RespID(key);
    }

    public void addResponse(
        CertificateID       certID,
        CertificateStatus   certStatus)
    {
        list.add(new ResponseObject(certID, certStatus, null, null));
    }

    public void addResponse(
        CertificateID       certID,
        CertificateStatus   certStatus,
        Date                nextUpdate,
        X509Extensions      singleExtensions)
    {
        list.add(new ResponseObject(certID, certStatus, nextUpdate, singleExtensions));
    }

    public void setResponseExtensions(
        X509Extensions  responseExtensions)
    {
        this.responseExtensions = responseExtensions;
    }

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

        ASN1EncodableVector responses = new ASN1EncodableVector();

        while (it.hasNext())
        {
            try
            {
                responses.add(((ResponseObject)it.next()).toResponse());
            }
            catch (Exception e)
            {
                throw new OCSPException("exception creating Request", e);
            }
        }

        ResponseData  tbsResp = new ResponseData(new DERInteger(0), responderID.toASN1Object(), new DERGeneralizedTime(new Date()), new DERSequence(responses), responseExtensions);

        java.security.Signature sig = 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();
            DEROutputStream         dOut = new DEROutputStream(bOut);

            dOut.writeObject(tbsResp);

            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);
            }

            return new BasicOCSPResp(new BasicOCSPResponse(tbsResp, sigAlgId, bitSig, new DERSequence(v)));
        }
        else
        {
            return new BasicOCSPResp(new BasicOCSPResponse(tbsResp, sigAlgId, bitSig, null));
        }
    }
    
    public BasicOCSPResp generate()
        throws OCSPException
    {
        try
        {
            return generateResponse(null, null, null, null, null);
        }
        catch (NoSuchProviderException e)
        {
            //
            // this shouldn't happen
            //
            throw new OCSPException("provider not found: " + e, e);
        }
    }

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

    public BasicOCSPResp 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 generateResponse(oid, key, chain, provider, random);
    }
}
