package org.bouncycastle.tools.openpgp.rampage;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.util.Collection;
import java.util.Iterator;

import org.bouncycastle.bcpg.ArmoredInputStream;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.openpgp.PGPCompressedData;
import org.bouncycastle.openpgp.PGPCompressedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedData;
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedDataList;
import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPKeyValidationException;
import org.bouncycastle.openpgp.PGPLiteralData;
import org.bouncycastle.openpgp.PGPObjectFactory;
import org.bouncycastle.openpgp.PGPOnePassSignature;
import org.bouncycastle.openpgp.PGPOnePassSignatureList;
import org.bouncycastle.openpgp.PGPPBEEncryptedData;
import org.bouncycastle.openpgp.PGPPublicKey;
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData;
import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection;
import org.bouncycastle.openpgp.PGPSecretKey;
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection;
import org.bouncycastle.openpgp.PGPSignature;
import org.bouncycastle.openpgp.PGPSignatureList;
import org.bouncycastle.openpgp.PGPUtil;
import org.bouncycastle.tools.openpgp.*;
import org.bouncycastle.tools.openpgp.util.*;

/**
 * A general tool for manipulating PGP objects.
 */
public class PGPRampageEngine implements ProcessingEngine
{
    private boolean   _verbose = true;
    private PGPParams _params  = null;

    public PGPRampageEngine(PGPParams params)
    {
        _params = params;
    }

    public void process()
    {
        boolean error = false;
        if (_params == null)
        {
            System.out.println("null parameters, much bad");
            return;
        }

        try
        {
            if (_params.isDecrypting())
            {
                if (((_params.getSecretKeyRingFile() != null) || (_params.getPassPhrase() != null))
                        && (_params.getPublicKeyRingFile() != null))
                {
                    if (_params.getKeyPassPhrase() != null)
                    {
                        decryptKeyBasedFile(new FileInputStream(_params.getInputFile()),
                            new FileInputStream(_params.getPublicKeyRingFile()),
                            new FileInputStream(_params.getSecretKeyRingFile()),
                            _params.getKeyPassPhrase().toCharArray(), _params.isMDCRequired());
                    }
                    else
                    {
                        decryptPBEBasedFile(new FileInputStream(_params.getInputFile()),
                            _params.getPassPhrase().toCharArray(), _params.isMDCRequired());
                    }
                }
                else
                {
                    System.out.println("Decryption could not be completed due to lack of information");
                }
            }
            else if (_params.isEncrypting())
            {
                if ((_params.getPublicKeyRingFile() != null)
                        && (_params.getRecipient() != null))
                {
                    FileInputStream publicRing = new FileInputStream(
                            _params.getPublicKeyRingFile());
                    String inputFileName = _params.getInputFile().getAbsolutePath();
                    String fileSuffix = _params.isAsciiArmor()
                            ? PGPParams.ASCII_SUFFIX
                            : PGPParams.BINARY_SUFFIX;
                    FileOutputStream outFile = new FileOutputStream(inputFileName + fileSuffix);

                    encryptFile(outFile, inputFileName, readPublicKey(publicRing,
                        _params.getRecipient()), _params.isAsciiArmor(), true);
                }
                else
                {
                    System.out.println("Encryption could not be completed due to lack of information");
                }
            }
            else if (_params.isVerify())
            {
                if (_params.getPublicKeyRingFile() != null)
                {
                    try
                    {
                        verifyFile(new FileInputStream(_params.getInputFile()),
                            new FileInputStream(_params.getPublicKeyRingFile()));
                    }
                    catch (SignatureException e)
                    {
                        error = true;
                        System.err.println("exception processing signature: " + e);
                    }
                }
                else
                {
                    System.out.println("PublicKeyRingFile is required for signature verification");
                }
            }
            else if (_params.isSigning())
            {
                System.out.println("Signing Not implemented");
            }
            else
            {
                System.out.println("Operation not implemented - please wait");
            }

        }

        catch (PGPException e)
        {
            error = true;
            System.err.println(e.toString());
        }
        catch (NoSuchProviderException e)
        {
            error = true;
            System.err.println(e.toString());
        }
        catch (IOException e)
        {
            error = true;
            System.err.println(e.toString());
        }
        catch (SignatureException unexpected)
        {
            error = true;
            unexpected.printStackTrace();
        }

        if (error)
        {
            System.out.println("error!");
            System.exit(1);
        }
    }

    private PGPPublicKey readPublicKey(InputStream in, String recipient) throws IOException,
            PGPException
    {
        in = PGPUtil.getDecoderStream(in);

        PGPPublicKeyRingCollection pgpPub = new PGPPublicKeyRingCollection(in);

        //
        // we just loop through the collection till we find a key suitable for encryption, in the real
        // world you would probably want to be a bit smarter about this.
        //
        PGPPublicKey key = null;

        //
        // iterate through the key rings.
        //
        Iterator rIt = pgpPub.getKeyRings();

        //System.out.println("processing public key ring, looking for : "+recipient);
        while (key == null && rIt.hasNext())
        {
            boolean nextEncryption = false;
            PGPPublicKeyRing kRing = (PGPPublicKeyRing) rIt.next();
            //System.out.println("Found a ring with keys ");
            Iterator kIt = kRing.getPublicKeys();

            while (key == null && kIt.hasNext())
            {
                PGPPublicKey k = (PGPPublicKey) kIt.next();
                Iterator userIDs = k.getUserIDs();
                String name = "<not specified>";
                if (userIDs.hasNext())
                {
                    name = (String) userIDs.next();
                }
                //System.out.println("found a key with name "+name);

                if (name.indexOf(recipient) >= 0)
                {
                    nextEncryption = true;
                }

                if (k.isEncryptionKey() && nextEncryption)
                {
                    //System.out.println("Found the key I'm looking for");
                    key = k;
                }
            }
        }

        if (key == null)
        {
            throw new IllegalArgumentException("Can't find encryption key in key ring.");
        }

        return key;
    }

    private void encryptFile(OutputStream out, String fileName, PGPPublicKey encKey,
            boolean armor, boolean withIntegrityCheck) throws IOException,
            NoSuchProviderException
    {
        if (armor)
        {
            out = new ArmoredOutputStream(out);
        }

        try
        {
            ByteArrayOutputStream bOut = new ByteArrayOutputStream();

            PGPCompressedDataGenerator comData = new PGPCompressedDataGenerator(
                    PGPCompressedData.ZIP);

            PGPUtil.writeFileToLiteralData(comData.open(bOut), PGPLiteralData.BINARY, new File(
                    fileName));

            comData.close();

            PGPEncryptedDataGenerator cPk = new PGPEncryptedDataGenerator(
                    PGPEncryptedData.CAST5, withIntegrityCheck, new SecureRandom(), "BC");

            cPk.addMethod(encKey);

            byte[] bytes = bOut.toByteArray();

            OutputStream cOut = cPk.open(out, bytes.length);

            cOut.write(bytes);

            cPk.close();

            out.close();
        }
        catch (PGPException e)
        {
            System.err.println(e);
            if (e.getUnderlyingException() != null)
            {
                e.getUnderlyingException().printStackTrace();
            }
        }
    }

    /**
     * decrypt the passed in message stream
     */
    public void decryptKeyBasedFile(InputStream fileToDecrypt,
            InputStream publicKeyInputStream, InputStream secretKeyInputStream, char[] passwd,
            boolean mdcRequired) throws IOException, NoSuchProviderException, PGPException,
            SignatureException
    {
        boolean error = false;

        fileToDecrypt = PGPUtil.getDecoderStream(fileToDecrypt);

        PGPObjectFactory outerWrapper = new PGPObjectFactory(fileToDecrypt);
        PGPEncryptedDataList enc;

        Object o = outerWrapper.nextObject();
        //
        // the first object might be a PGP marker packet.
        //
        if (o instanceof PGPEncryptedDataList)
        {
            enc = (PGPEncryptedDataList) o;
        }
        else
        {
            enc = (PGPEncryptedDataList) outerWrapper.nextObject();
        }

        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(
                PGPUtil.getDecoderStream(secretKeyInputStream));

        PGPSecretKey pgpSecKey = null;
        PGPPublicKeyEncryptedData pked = null;
        int count = 0;

        // find the secret key that is needed
        while (count != enc.size())
        {
            if (enc.get(count) instanceof PGPPublicKeyEncryptedData)
            {
                pked = (PGPPublicKeyEncryptedData) enc.get(count);
                pgpSecKey = pgpSec.getSecretKey(pked.getKeyID());
                if (pgpSecKey != null)
                {
                    break;
                }
            }

            count++;
        }

        if (pgpSecKey == null)
        {
            throw new PGPException("corresponding secret key not found.");
        }

        InputStream clear = pked.getDataStream(pgpSecKey.extractPrivateKey(passwd, "BC"), "BC");

        PGPObjectFactory plainFact = new PGPObjectFactory(clear);

        PGPCompressedData cData = (PGPCompressedData) plainFact.nextObject();

        PGPObjectFactory pgpFact = new PGPObjectFactory(cData.getDataStream());

        Object message = pgpFact.nextObject();

        // 
        if (message instanceof PGPLiteralData)
        {
            PGPLiteralData ld = (PGPLiteralData) message;

            FileOutputStream fOut = new FileOutputStream(ld.getFileName());

            InputStream unc = ld.getInputStream();
            int ch;

            while ((ch = unc.read()) >= 0)
            {
                fOut.write(ch);
            }
        }
        else
        // onepass signature
        {
            PGPOnePassSignatureList onePassSigList = (PGPOnePassSignatureList) message;
            PGPOnePassSignature onePassSig = onePassSigList.get(0);

            PGPLiteralData ld = (PGPLiteralData) pgpFact.nextObject();

            InputStream dataIn = ld.getInputStream();

            PGPPublicKeyRingCollection pgpRing = new PGPPublicKeyRingCollection(
                    PGPUtil.getDecoderStream(publicKeyInputStream));

            PGPPublicKey publicKey = pgpRing.getPublicKey(onePassSig.getKeyID());
            FileOutputStream out = new FileOutputStream(ld.getFileName());

            onePassSig.initVerify(publicKey, "BC");

            int ch;
            while ((ch = dataIn.read()) >= 0)
            {
                onePassSig.update((byte) ch);
                out.write(ch);
            }
            out.close();

            PGPSignatureList sigList = (PGPSignatureList) pgpFact.nextObject();
            if (onePassSig.verify(sigList.get(0)))
            {
                System.out.println("Signature verified");
            }
            else
            {
                System.out.println("Signature verification failed");
            }

        }

        if (pked.isIntegrityProtected())
        {
            if (!pked.verify())
            {
                if (_verbose)
                {
                    System.out.println("message failed integrity check");
                }
                error = true;
            }
            else
            {
                if (_verbose)
                {
                    System.out.println("message integrity check passed");
                }
            }
        }
        else
        {
            if (_verbose)
            {
                System.out.println("no message integrity check");
            }

            if (mdcRequired)
            {
                error = true;
            }
        }
    }

    /**
     * decrypt the passed in message stream
     */
    public void decryptPBEBasedFile(InputStream in, char[] passPhrase, boolean mdcRequired)
            throws IOException, PGPException, NoSuchProviderException
    {
        //
        // we need to be able to reset the stream if we try a
        // wrong password, we'll assume that all the mechanisms
        // appear in the first 10k for the moment...
        //
        int READ_LIMIT = 10 * 1024;
        boolean error = false;

        in.mark(READ_LIMIT);

        PGPPBEEncryptedData pbe;
        InputStream clear;
        int count = 0;

        for (;;)
        {
            InputStream dIn = PGPUtil.getDecoderStream(in);

            PGPObjectFactory pgpF = new PGPObjectFactory(dIn);
            PGPEncryptedDataList enc;
            Object o = pgpF.nextObject();

            //
            // the first object might be a PGP marker packet.
            //
            if (o instanceof PGPEncryptedDataList)
            {
                enc = (PGPEncryptedDataList) o;
            }
            else
            {
                enc = (PGPEncryptedDataList) pgpF.nextObject();
            }

            while (count < enc.size())
            {
                if (enc.get(count) instanceof PGPPBEEncryptedData)
                {
                    break;
                }

                count++;
            }

            if (count >= enc.size())
            {
                throw new PGPException("password invalid");
            }

            pbe = (PGPPBEEncryptedData) enc.get(count);

            try
            {
                clear = pbe.getDataStream(passPhrase, "BC");
            }
            catch (PGPKeyValidationException e)
            {
                in.reset();
                continue;
            }

            break;
        }

        PGPObjectFactory pgpFact = new PGPObjectFactory(clear);

        PGPCompressedData cData = (PGPCompressedData) pgpFact.nextObject();

        pgpFact = new PGPObjectFactory(cData.getDataStream());

        PGPLiteralData ld = (PGPLiteralData) pgpFact.nextObject();

        FileOutputStream fOut = new FileOutputStream(ld.getFileName());

        InputStream unc = ld.getInputStream();
        int ch;

        while ((ch = unc.read()) >= 0)
        {
            fOut.write(ch);
        }

        if (pbe.isIntegrityProtected())
        {
            if (!pbe.verify())
            {
                if (_verbose)
                {
                    System.out.println("message failed integrity check");
                }
                error = true;
            }
            else
            {
                if (_verbose)
                {
                    System.out.println("message integrity check passed");
                }
            }
        }
        else
        {
            if (_verbose)
            {
                System.out.println("no message integrity check");
            }

            if (mdcRequired)
            {
                error = true;
            }
        }
    }

    /**
     * verify the passed in file as being correctly signed.
     */
    public void verifyFile(InputStream in, InputStream keyIn) throws IOException,
            NoSuchProviderException, PGPException, SignatureException
    {
        boolean error = false;

        in = PGPUtil.getDecoderStream(in);

        //
        // a clear signed file
        //
        if (in instanceof ArmoredInputStream && ((ArmoredInputStream) in).isClearText())
        {
            //
            // read the input, making sure we ingore the last newline.
            //
            int ch;
            ArmoredInputStream aIn = (ArmoredInputStream) in;
            boolean newLine = false;
            ByteArrayOutputStream bOut = new ByteArrayOutputStream();

            while ((ch = aIn.read()) >= 0 && aIn.isClearText())
            {
                if (newLine)
                {
                    bOut.write((byte) '\n');
                    newLine = false;
                }
                if (ch == '\n')
                {
                    newLine = true;
                    continue;
                }

                bOut.write((byte) ch);
            }

            PGPPublicKeyRingCollection pgpRings = new PGPPublicKeyRingCollection(keyIn);

            PGPObjectFactory pgpFact = new PGPObjectFactory(aIn);
            PGPSignatureList p3 = (PGPSignatureList) pgpFact.nextObject();
            PGPSignature sig = null;
            int count = 0;
            PGPPublicKey key = null;

            while (count != p3.size())
            {
                sig = (PGPSignature) p3.get(count);
                key = pgpRings.getPublicKey(sig.getKeyID());
                if (key != null)
                {
                    break;
                }

                count++;
            }

            if (key == null)
            {
                throw new PGPException("corresponding public key not found.");
            }

            sig.initVerify(pgpRings.getPublicKey(sig.getKeyID()), "BC");

            sig.update(bOut.toByteArray());

            if (sig.verify())
            {
                if (_verbose)
                {
                    System.out.println("signature verified.");
                }
            }
            else
            {
                if (_verbose)
                {
                    System.out.println("signature verification failed.");
                }
                error = true;
            }
        }
        else
        {
            PGPObjectFactory pgpFact = new PGPObjectFactory(in);

            PGPCompressedData c1 = (PGPCompressedData) pgpFact.nextObject();

            pgpFact = new PGPObjectFactory(c1.getDataStream());

            PGPOnePassSignatureList p1 = (PGPOnePassSignatureList) pgpFact.nextObject();
            PGPPublicKeyRingCollection pgpRing = new PGPPublicKeyRingCollection(
                    PGPUtil.getDecoderStream(keyIn));
            PGPPublicKey key = null;
            int count = 0;
            PGPOnePassSignature ops = null;

            while (count != p1.size())
            {
                ops = (PGPOnePassSignature) p1.get(count);
                key = pgpRing.getPublicKey(ops.getKeyID());
                if (key != null)
                {
                    break;
                }

                count++;
            }

            if (key == null)
            {
                throw new PGPException("corresponding public key not found.");
            }

            PGPLiteralData p2 = (PGPLiteralData) pgpFact.nextObject();
            InputStream dIn = p2.getInputStream();
            int ch;
            FileOutputStream out = new FileOutputStream(p2.getFileName());

            ops.initVerify(key, "BC");

            while ((ch = dIn.read()) >= 0)
            {
                ops.update((byte) ch);
                out.write(ch);
            }

            out.close();

            PGPSignatureList p3 = (PGPSignatureList) pgpFact.nextObject();

            if (ops.verify(p3.get(0)))
            {
                if (_verbose)
                {
                    System.out.println("signature verified.");
                }
            }
            else
            {
                if (_verbose)
                {
                    System.out.println("signature verification failed.");
                }
                error = true;
            }
        }
    }

    // TODO: Implement these methods so that can be used by the CmdLineProcessors
    // and the test cases.
    public boolean isError()
    {
        return false;
    }

    public Collection errorMessages()
    {
        return null;
    }

}
