package org.bouncycastle.openpgp;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.KeyFactory;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.interfaces.DSAParams;
import java.security.interfaces.DSAPrivateKey;
import java.security.interfaces.DSAPublicKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.DSAPrivateKeySpec;
import java.security.spec.RSAPrivateCrtKeySpec;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.Date;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

import org.bouncycastle.bcpg.*;
import org.bouncycastle.jce.interfaces.ElGamalPrivateKey;
import org.bouncycastle.jce.interfaces.ElGamalPublicKey;
import org.bouncycastle.jce.spec.ElGamalParameterSpec;
import org.bouncycastle.jce.spec.ElGamalPrivateKeySpec;

/**
 * general class to handle a PGP secret key object.
 */
public class PGPSecretKey
{
	private long 		keyID;
	private byte[]	fingerPrint;
	
    SecretKeyPacket	secret;
    TrustPacket			trust;
	ArrayList				ids;
	ArrayList				idTrusts;
	ArrayList				idSigs;
	PGPPublicKey		pub;
	ArrayList				subSigs = null;
	
	PGPSecretKey(
		SecretKeyPacket	secret,
		TrustPacket			trust,
		MessageDigest		sha,
		ArrayList				ids,
		ArrayList				idTrusts,
		ArrayList				idSigs)
		throws IOException
	{
		this.secret = secret;
		this.trust = trust;
		this.ids = ids;
		this.idTrusts = idTrusts;
		this.idSigs = idSigs;
		this.pub = new PGPPublicKey(secret.getPublicKeyPacket(), trust, sha, ids, null, idSigs);
	}
	
	PGPSecretKey(
		SecretKeyPacket	secret,
		TrustPacket			trust,
		MessageDigest		sha,
		ArrayList				subSigs)
		throws IOException
	{
		this.secret = secret;
		this.trust = trust;
		this.subSigs = subSigs;
		this.pub = new PGPPublicKey(secret.getPublicKeyPacket(), trust, sha, ids, null, idSigs);
	}
	
	public PGPSecretKey(
		int										certificationLevel,
		int										algorithm,
		PublicKey								pubKey,
		PrivateKey								privKey,
		Date										time,
		String										id,
		int											encAlgorithm,
		char[]										passPhrase,
		PGPSignatureSubpacketVector	hashedPcks,
		PGPSignatureSubpacketVector	unhashedPcks,
		SecureRandom						rand,
		String									provider)
		throws PGPException, NoSuchProviderException
	{
		PublicKeyPacket			pubPk;
		BCPGObject				secKey;
		MessageDigest 			sha;
		
		try
		{
			sha = MessageDigest.getInstance("SHA1");
		}
		catch (NoSuchAlgorithmException e)
		{
			throw new PGPException("can't find SHA1 digest");
		}
		
		if (pubKey instanceof RSAPublicKey)
		{
			RSAPublicKey		rK = (RSAPublicKey)pubKey;
			RSAPrivateCrtKey	rsK = (RSAPrivateCrtKey)privKey;
			
			pubPk = new PublicKeyPacket(algorithm, time, new RSAPublicBCPGKey(rK.getModulus(), rK.getPublicExponent()));
			secKey = new RSASecretBCPGKey(rsK.getPrivateExponent(), rsK.getPrimeP(), rsK.getPrimeQ());
		}
		else if (pubKey instanceof DSAPublicKey)
		{
			DSAPublicKey	dK = (DSAPublicKey)pubKey;
			DSAParams		dP = dK.getParams();
			DSAPrivateKey	dsK = (DSAPrivateKey)privKey;
			
			pubPk = new PublicKeyPacket(algorithm, time, new DSAPublicBCPGKey(dP.getP(), dP.getQ(), dP.getG(),  dK.getY()));
			secKey = new DSASecretBCPGKey(dsK.getX());
		}
		else if (pubKey instanceof ElGamalPublicKey)
		{
			ElGamalPublicKey			eK = (ElGamalPublicKey)pubKey;
			ElGamalParameterSpec	eS = eK.getParams();
			ElGamalPrivateKey			esK = (ElGamalPrivateKey)privKey;
			
			pubPk = new PublicKeyPacket(algorithm, time, new ElGamalPublicBCPGKey(eS.getP(), eS.getG(), eK.getY()));
			secKey = new ElGamalSecretBCPGKey(esK.getX());
		}
		else
		{
			throw new PGPException("unknown key class");
		}

		String	cName = PGPUtil.getSymmetricCipherName(encAlgorithm);
		
		if (cName == null)
		{
			throw new PGPException("null cipher specified");
		}
		
		Cipher										c = null;
		
		try
		{
			c = Cipher.getInstance(cName + "/CFB/NoPadding", provider);
		}
		catch (NoSuchProviderException e)
		{
			throw e;
		}
		catch (Exception e)
		{
			throw new PGPException("Exception creating cipher", e);
		}
		
		byte[]										iv = new byte[8];
		
		rand.nextBytes(iv);
		
		S2K											s2k = new S2K(HashAlgorithmTags.SHA1, iv, 0x60);

		try
		{
			ByteArrayOutputStream	bOut = new ByteArrayOutputStream();
			BCPGOutputStream			pOut = new BCPGOutputStream(bOut);
			
			SecretKey	key = PGPUtil.makeKeyFromPassPhrase(encAlgorithm, s2k, passPhrase, provider);

			c.init(Cipher.ENCRYPT_MODE, key, rand);
		
			iv = c.getIV();
			
			pOut.writeObject(secKey);
			
			byte[]	keyData = bOut.toByteArray();
			int		checkSum = 0;
			
			for (int i = 0; i != keyData.length; i++)
			{
				checkSum += keyData[i] & 0xff;
			}
			
			pOut.write(checkSum >> 8);
			pOut.write(checkSum);
			
			byte[]	encData = c.doFinal(bOut.toByteArray());

			this.secret = new SecretKeyPacket(pubPk, encAlgorithm, s2k, iv, encData);
			this.trust = null;
			
			this.ids = new ArrayList();
			ids.add(id);
			
			this.idTrusts = new ArrayList();
			idTrusts.add(null);
			
			this.idSigs = new ArrayList();
			
			this.pub = new PGPPublicKey(pubPk, trust, sha, ids, idTrusts, idSigs);
			
			PGPSignatureGenerator	sGen = new PGPSignatureGenerator(pubPk.getAlgorithm(), HashAlgorithmTags.SHA1, provider);
			
			//
			// generate the certification
			//
			sGen.initSign(certificationLevel, new PGPPrivateKey(privKey, pub.getKeyID()));
			
			sGen.setHashedSubpackets(hashedPcks);
			sGen.setUnhashedSubpackets(unhashedPcks);
			
			ArrayList sigList = new ArrayList();
			
			sigList.add(sGen.generateCertification(id, pub));
			idSigs.add(sigList);
		}
		catch (PGPException e)
		{
			throw e;
		}
		catch (Exception e)
		{
			throw new PGPException("Exception encrypting key", e);
		}
	}

	/**
	 * return true if this key is marked as suitable for signature generation.
	 */
	public boolean isSigningKey()
	{
		int algorithm = secret.getEncAlgorithm();

		return ((algorithm == PGPPublicKey.RSA_GENERAL) || (algorithm == PGPPublicKey.RSA_SIGN)
					|| (algorithm == PGPPublicKey.DSA) || (algorithm == PGPPublicKey.ECDSA));
	}
	
	/**
	 * Return true if this is a master key.
	 * @return true if a master key.
	 */
	public boolean isMasterKey()
	{
		return (subSigs == null);
	}
	
	/**
	 * Return the keyID associated with this key.
	 * 
	 * @return long
	 */
	public long getKeyID()
	{
		return pub.getKeyID();
	}
	
	/**
	 * Return the public key associated with this key.
	 * 
	 * @return PGPPublicKey
	 */
    public PGPPublicKey getPublicKey()
    {
    	return pub;
    }
    
	/**
	 * Return any userIDs associated with the key.
	 * 
	 * @return an iterator of Strings.
	 */
	public Iterator getUserIDs()
	{
		return pub.getUserIDs();
	}
	
	/**
	 * Return any user attribute vectors associated with the key.
	 * 
	 * @return an iterator of Strings.
	 */
	public Iterator getUserAttributes()
	{
		return pub.getUserAttributes();
	}
	
    /**
     * Extract a PGPPrivate key from the SecretKey's encrypted contents.
     * 
     * @param passPhrase
     * @param provider
     * @return PGPPrivateKey
     * @throws PGPException
     * @throws NoSuchProviderException
     */
	public  PGPPrivateKey extractPrivateKey(
		char[]				passPhrase,
		String				provider)
		throws PGPException, NoSuchProviderException
	{
		PublicKeyPacket			pub = secret.getPublicKeyPacket();
		BCPGKey					sKey = null;
	
		String	cName = PGPUtil.getSymmetricCipherName(secret.getEncAlgorithm());
	
		
		Cipher						c = null;
		
		if (cName != null)
		{
			try
			{
				c = Cipher.getInstance(cName + "/CFB/NoPadding", provider);
			}
			catch (NoSuchProviderException e)
			{
				throw e;
			}
			catch (Exception e)
			{
				throw new PGPException("Exception creating cipher", e);
			}
		}
	
		byte[]	encData = secret.getSecretKeyData();
		byte[]	data = null;
	
		try
		{
			if (c != null)
			{
				try
				{
					if (secret.getPublicKeyPacket().getVersion() == 4)
					{
						IvParameterSpec ivSpec = new IvParameterSpec(secret.getIV());
		
						SecretKey	key = PGPUtil.makeKeyFromPassPhrase(secret.getEncAlgorithm(), secret.getS2K(), passPhrase, provider);
		
						c.init(Cipher.DECRYPT_MODE, key, ivSpec);
					
						data = c.doFinal(encData, 0, encData.length);
					}
					else // version 2 or 3, RSA only.
					{
						SecretKey	key = PGPUtil.makeKeyFromPassPhrase(secret.getEncAlgorithm(), secret.getS2K(), passPhrase, provider);
	
						data = new byte[encData.length];
				
						byte[]	iv = new byte[secret.getIV().length];
				
						System.arraycopy(secret.getIV(), 0, iv, 0, iv.length);
				
						//
						// read in the four numbers
						//
						int	pos = 0;
						
						for (int i = 0; i != 4; i++)
						{
							c.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
					
							int encLen = (((encData[pos] << 8) | (encData[pos + 1] & 0xff)) + 7) / 8;

							data[pos] = encData[pos];
							data[pos + 1] = encData[pos + 1];

							c.doFinal(encData, pos + 2, encLen, data, pos + 2);
							pos += 2 + encLen;
				
							if (i != 3)
							{
								System.arraycopy(encData, pos - iv.length, iv, 0, iv.length);
							}
						}

						//
						// verify checksum
						//
						
						int cs = ((encData[pos] << 8) & 0xff00) | (encData[pos + 1] & 0xff);
						int calcCs = 0;
						for (int j=0; j < data.length-2; j++) 
						{
							calcCs += data[j] & 0xff;
						}
			
						calcCs &= 0xffff;
						if (calcCs != cs) 
						{
							throw new PGPException("checksum mismatch: passphrase wrong, expected "
												+ Integer.toHexString(cs)
												+ " found " + Integer.toHexString(calcCs));
						}
					}
				}
				catch (PGPException e)
				{
					throw e;
				}
				catch (Exception e)
				{
					throw new PGPException("Exception decrypting key", e);
				}
			}
			else
			{
				data = encData;
			}

			KeyFactory		fact;
			BCPGInputStream	in = new BCPGInputStream(new ByteArrayInputStream(data));
		
			switch (pub.getAlgorithm())
			{
			case PGPPublicKey.RSA_ENCRYPT:
			case PGPPublicKey.RSA_GENERAL:
				RSAPublicBCPGKey   	rsaPub = (RSAPublicBCPGKey)pub.getKey();
				RSASecretBCPGKey		rsaPriv = new RSASecretBCPGKey(in);
				RSAPrivateCrtKeySpec	rsaPrivSpec = new RSAPrivateCrtKeySpec(
													rsaPriv.getModulus(), 
													rsaPub.getPublicExponent(),
													rsaPriv.getPrivateExponent(),
													rsaPriv.getPrimeP(),
													rsaPriv.getPrimeQ(),
													rsaPriv.getPrimeExponentP(),
													rsaPriv.getPrimeExponentQ(),
													rsaPriv.getCrtCoefficient());
            						
				fact = KeyFactory.getInstance("RSA", provider);

				return new PGPPrivateKey(fact.generatePrivate(rsaPrivSpec), this.getKeyID());	
			case PGPPublicKey.DSA:
				DSAPublicBCPGKey  	dsaPub = (DSAPublicBCPGKey)pub.getKey();
				DSASecretBCPGKey   dsaPriv = new DSASecretBCPGKey(in);
				DSAPrivateKeySpec	dsaPrivSpec =
					new DSAPrivateKeySpec(dsaPriv.getX(), dsaPub.getP(), dsaPub.getQ(), dsaPub.getG());

				fact = KeyFactory.getInstance("DSA", provider);

				return new PGPPrivateKey(fact.generatePrivate(dsaPrivSpec), this.getKeyID());
			case PGPPublicKey.ELGAMAL_ENCRYPT:
			case PGPPublicKey.ELGAMAL_GENERAL:
				ElGamalPublicBCPGKey  	elPub = (ElGamalPublicBCPGKey)pub.getKey();
				ElGamalSecretBCPGKey	elPriv = new ElGamalSecretBCPGKey(in);
				ElGamalPrivateKeySpec	elSpec = new ElGamalPrivateKeySpec(elPriv.getX(), new ElGamalParameterSpec(elPub.getP(), elPub.getG()));
			
				fact = KeyFactory.getInstance("ElGamal", provider);
			
				return new PGPPrivateKey(fact.generatePrivate(elSpec), this.getKeyID());
			default:
				throw new PGPException("unknown public key algorithm encountered");
			}
		}
		catch (PGPException e)
		{
			throw e;
		}
		catch (Exception e)
		{
			throw new PGPException("Exception constructing key", e);
		}
	}
	
	public byte[] getEncoded() 
		throws IOException
	{
		ByteArrayOutputStream	bOut = new ByteArrayOutputStream();
		
		this.encode(bOut);
		
		return bOut.toByteArray();
	}
	
	public void encode(
		OutputStream	outStream) 
		throws IOException
	{
		BCPGOutputStream	out;
		
		if (outStream instanceof BCPGOutputStream)
		{
			out = (BCPGOutputStream)outStream;
		}
		else
		{
			out = new BCPGOutputStream(outStream);
		}

		out.writePacket(secret);
		if (trust != null)
		{
			out.writePacket(trust);
		}
        
        if (subSigs == null)		// is not a sub key
        {
			for (int i = 0; i != ids.size(); i++)
			{
				if (ids.get(i) instanceof String)
				{
					String	id = (String)ids.get(i);
		        	
					out.writePacket(new UserIDPacket(id));
				}
				else
				{
					PGPUserAttributeSubpacketVector	v = (PGPUserAttributeSubpacketVector)ids.get(i);

					out.writePacket(new UserAttributePacket(v.toSubpacketArray()));
				}
				
				if (idTrusts.get(i) != null)
				{
					out.writePacket((ContainedPacket)idTrusts.get(i));
				}
				
				ArrayList	sigs = (ArrayList)idSigs.get(i);
				
				for (int j = 0; j != sigs.size(); j++)
				{
					((PGPSignature)sigs.get(j)).encode(out);
				}
			}
        }
        else
        {		
			for (int j = 0; j != subSigs.size(); j++)
			{
				((PGPSignature)subSigs.get(j)).encode(out);
			}
        }
    }
}
