package org.bouncycastle.openpgp.examples;

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.Security;

import org.bouncycastle.bcpg.ArmoredInputStream;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedDataList;
import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPLiteralData;
import org.bouncycastle.openpgp.PGPCompressedDataGenerator;
import org.bouncycastle.openpgp.PGPCompressedData;
import org.bouncycastle.openpgp.PGPObjectFactory;
import org.bouncycastle.openpgp.PGPPBEEncryptedData;
import org.bouncycastle.openpgp.PGPUtil;

/**
 * A simple utility class that encrypts/decrypts password based
 * encryption files.
 * <p>
 * To encrypt a file: PBEFileProcessor -e [-a] fileName passPhrase.<br>
 * If -a is specified the output file will be "ascii-armored".
 * <p>
 * To decrypt: PBEFileProcessor -d fileName passPhrase.
 * <p>
 * Note: this example will silently overwrite files, nor does it pay any attention to
 * the specification of "_CONSOLE" in the filename. It also expects that a single pass phrase
 * will have been used.
 */
public class PBEFileProcessor
{
	/**
	 * decrypt the passed in message stream
	 */
	private static void decryptFile(
		InputStream	in,
		char[]			passPhrase,
		boolean		armored)
		throws Exception
	{
		if (armored)
		{
			in = new ArmoredInputStream(in);
		}
		
        PGPObjectFactory pgpF = new PGPObjectFactory(in);

		PGPEncryptedDataList	enc = (PGPEncryptedDataList)pgpF.nextObject();
		
        PGPPBEEncryptedData	pbe = (PGPPBEEncryptedData)enc.get(0);

        InputStream clear = pbe.getDataStream(passPhrase, "BC");
        
        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);
		}
	}

	private static void encryptFile(
		OutputStream	out,
		String				fileName,
		char[]				passPhrase,
		boolean			armor)
		throws IOException, NoSuchProviderException, PGPException
	{	
		if (armor)
		{
			out = new ArmoredOutputStream(out);
		}
		
		ByteArrayOutputStream	bOut = new ByteArrayOutputStream();
		

		PGPCompressedDataGenerator comData = new PGPCompressedDataGenerator(
																PGPCompressedDataGenerator.ZIP);
																
		PGPUtil.writeFileToLiteralData(comData.open(bOut), PGPLiteralData.BINARY, new File(fileName));
		
		comData.close();
		
		PGPEncryptedDataGenerator	cPk = new PGPEncryptedDataGenerator(PGPEncryptedDataGenerator.CAST5, new SecureRandom(), "BC");
			
		cPk.addMethod(passPhrase);
		
		byte[]				bytes = bOut.toByteArray();
		
		OutputStream	cOut = cPk.open(out, bytes.length);

		cOut.write(bytes);
		
		cOut.close();
    }

    public static void main(
    	String[] args)
    	throws Exception
    {
        Security.addProvider(new BouncyCastleProvider());

		if (args[0].equals("-e"))
		{
			if (args[1].equals("-a"))
			{
				FileOutputStream	out = new FileOutputStream(args[2] + ".asc");
				encryptFile(out, args[2], args[3].toCharArray(), true);
			}
			else
			{
				FileOutputStream	out = new FileOutputStream(args[1] + ".bpg");
				encryptFile(out, args[1], args[2].toCharArray(), false);
			}
		}
		else if (args[0].equals("-d"))
		{
			if (args[1].equals("-a"))
			{
				FileInputStream	in = new FileInputStream(args[2]);
				decryptFile(in, args[3].toCharArray(), true);
			}
			else
			{
				FileInputStream	in = new FileInputStream(args[1]);
				decryptFile(in, args[2].toCharArray(), false);
			}
		}
		else
		{
			System.err.println("usage: PBEFileProcessor -e|-d [-a] file passPhrase");
		}
    }
}
