package org.bouncycastle.openpgp;

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.PublicKey;
import java.security.interfaces.DSAParams;
import java.security.interfaces.DSAPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.DSAPublicKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.ArrayList;
import java.util.Date;
import java.util.Iterator;

import org.bouncycastle.bcpg.*;

import org.bouncycastle.jce.interfaces.ElGamalPublicKey;
import org.bouncycastle.jce.spec.ElGamalParameterSpec;
import org.bouncycastle.jce.spec.ElGamalPublicKeySpec;

/**
 * general class to handle a PGP public key object.
 */
public class PGPPublicKey
	implements PublicKeyAlgorithmTags
{
    private long 	keyID;
	private byte[]	fingerprint;
	private int		keyStrength;
	
	PublicKeyPacket		publicPk;
	TrustPacket				trustPk;
	ArrayList					ids = new ArrayList();
	ArrayList					idTrusts = new ArrayList();
	ArrayList					idSigs = new ArrayList();
	
	private ArrayList		subSigs = null;

	private void init(
		MessageDigest	sha)
		throws IOException
	{
		BCPGKey				key = publicPk.getKey();
		
		if (publicPk.getVersion() <= 3)
		{
			RSAPublicBCPGKey	rK = (RSAPublicBCPGKey)key;
			
			this.keyID = rK.getModulus().longValue();
				
			sha.update(new MPInteger(rK.getModulus()).getEncoded());
			sha.update(new MPInteger(rK.getPublicExponent()).getEncoded());
			
			this.fingerprint = sha.digest();
			this.keyStrength = rK.getModulus().bitLength();
		}
		else
		{
				byte[]							kBytes = publicPk.getEncodedContents();

				sha.update((byte)0x99);
				sha.update((byte)(kBytes.length >> 8));
				sha.update((byte)kBytes.length);
				sha.update(kBytes);
				
				this.fingerprint = sha.digest();
				this.keyID = ((long)(fingerprint[fingerprint.length - 8] & 0xff) << 56)
								| ((long)(fingerprint[fingerprint.length - 7] & 0xff) << 48)
								| ((long)(fingerprint[fingerprint.length - 6] & 0xff) << 40)
								| ((long)(fingerprint[fingerprint.length - 5] & 0xff) << 32)
								| ((long)(fingerprint[fingerprint.length - 4] & 0xff) << 24)
								| ((long)(fingerprint[fingerprint.length - 3] & 0xff) << 16)
								| ((long)(fingerprint[fingerprint.length - 2] & 0xff) << 8)
								| ((fingerprint[fingerprint.length - 1] & 0xff));
				
				if (key instanceof RSAPublicBCPGKey)
				{
					this.keyStrength = ((RSAPublicBCPGKey)key).getModulus().bitLength();
				}
				else if (key instanceof DSAPublicBCPGKey)
				{
					this.keyStrength = ((DSAPublicBCPGKey)key).getG().bitLength();
				}
				else if (key instanceof ElGamalPublicBCPGKey)
				{
					this.keyStrength = ((ElGamalPublicBCPGKey)key).getG().bitLength();
				}
		}
	}
	
	/**
	 * Create a PGPPublicKey from the passed in JCA one.
	 * <p>
	 * Note: the time passed in affects the value of the key's keyID, so you probably only want
	 * to do this once for a JCA key, or make sure you keep track of the time you used.
	 * 
	 * @param algorithm
	 * @param pubKey
	 * @param time
	 * @param provider
	 * @throws PGPException
	 * @throws NoSuchProviderException
	 */
	public PGPPublicKey(		
		int					algorithm,
		PublicKey		pubKey,
		Date				time,
		String				provider) 
		throws PGPException, NoSuchProviderException
	{
		PublicKeyPacket			pubPk;
		MessageDigest 			sha;
		
		try
		{
			sha = MessageDigest.getInstance("SHA1", provider);
		}
		catch (NoSuchAlgorithmException e)
		{
			throw new PGPException("can't find SHA1 digest");
		}
		
		if (pubKey instanceof RSAPublicKey)
		{
			RSAPublicKey		rK = (RSAPublicKey)pubKey;
			
			pubPk = new PublicKeyPacket(algorithm, time, new RSAPublicBCPGKey(rK.getModulus(), rK.getPublicExponent()));
		}
		else if (pubKey instanceof DSAPublicKey)
		{
			DSAPublicKey	dK = (DSAPublicKey)pubKey;
			DSAParams		dP = dK.getParams();
			
			pubPk = new PublicKeyPacket(algorithm, time, new DSAPublicBCPGKey(dP.getP(), dP.getQ(), dP.getG(),  dK.getY()));
		}
		else if (pubKey instanceof ElGamalPublicKey)
		{
			ElGamalPublicKey			eK = (ElGamalPublicKey)pubKey;
			ElGamalParameterSpec	eS = eK.getParams();
			
			pubPk = new PublicKeyPacket(algorithm, time, new ElGamalPublicBCPGKey(eS.getP(), eS.getG(), eK.getY()));
		}
		else
		{
			throw new PGPException("unknown key class");
		}
		
		this.publicPk = pubPk;
		this.ids = new ArrayList();
		this.idSigs = new ArrayList();
		
		try
        {
            init(sha);
        }
        catch (IOException e)
        {
			throw new PGPException("exception calculating keyID", e);
        }
	}
	
	/**
	  * Constructor for a sub-key.
	  * 
	  * @param pk
	  * @param sha
	  * @param sig
	  */
	PGPPublicKey(
		PublicKeyPacket 	publicPk, 
		TrustPacket			trustPk,
		MessageDigest 	sha, 
		ArrayList 				sigs)
		throws IOException
	 {
		this.publicPk = publicPk;
		this.trustPk = trustPk;
		this.subSigs = sigs;
		
		init(sha);
	 }
	 
	/**
	 * Copy constructor.
	 * @param pubKey
	 */
	PGPPublicKey(
		PGPPublicKey	pubKey)
	 {
        this.publicPk = pubKey.publicPk;
        
		this.ids = new ArrayList(pubKey.ids);
		this.idTrusts = new ArrayList(pubKey.idTrusts);
		this.idSigs = new ArrayList(pubKey.idSigs.size());
		for (int i = 0; i != pubKey.idSigs.size(); i++)
		{
			this.idSigs.add(new ArrayList((ArrayList)pubKey.idSigs.get(i)));
		}
        
        if (this.subSigs != null)
        {
			this.subSigs = new ArrayList(pubKey.subSigs.size());
			for (int i = 0; i != pubKey.subSigs.size(); i++)
			{
				this.subSigs.add(new ArrayList((ArrayList)pubKey.idSigs.get(i)));
			}
        }
        
		this.fingerprint = pubKey.fingerprint;
		this.keyID = pubKey.keyID;
		this.keyStrength = pubKey.keyStrength;
	 }

	PGPPublicKey(
		PublicKeyPacket 	publicPk,
		TrustPacket			trustPk,
		MessageDigest	sha,
		ArrayList			ids,
		ArrayList			idTrusts,
		ArrayList			idSigs)
		throws IOException
	{
		this.publicPk = publicPk;
		this.trustPk = trustPk;
		this.ids = ids;
		this.idTrusts = idTrusts;
		this.idSigs = idSigs;
		
		init(sha);
	}
	
	PGPPublicKey(
		PublicKeyPacket 	publicPk,
		ArrayList			ids,
		ArrayList			idSigs)
		throws IOException
	{
		this.publicPk = publicPk;
		this.ids = ids;
		this.idSigs = idSigs;
	}
	
	/**
	 * @return the version of this key.
	 */
	public int getVersion()
	{
		return publicPk.getVersion();
	}
	
	/**
	 * @return creation time of key.
	 */
	public Date getCreationTime()
	{
		return publicPk.getTime();
	}
	
	/**
	 * @return number of valid days from creation time - zero means no
	 * expiry.
	 */
	public int getValidDays()
	{
		return publicPk.getValidDays();
	}
	
	/**
	 * Return the keyID associated with the public key.
	 * 
	 * @return long
	 */
	public long getKeyID()
	{
		return keyID;
	}
	
	/**
	 * Return the 16 byte fingerprint of the key.
	 * 
	 * @return key fingerprint.
	 */
	public byte[] getFingerprint()
	{
		byte[]	tmp = new byte[fingerprint.length];
		
		System.arraycopy(fingerprint, 0, tmp, 0, tmp.length);
		
		return tmp;
	}
	
    /**
     * return true if this key is marked as suitable for encryption.
     */
    public boolean isEncryptionKey()
    {
        int algorithm = publicPk.getAlgorithm();

        return ((algorithm == RSA_GENERAL) || (algorithm == RSA_ENCRYPT)
                || (algorithm == ELGAMAL_ENCRYPT) || (algorithm == ELGAMAL_GENERAL));
    }

	/**
	 * Return true if this is a master key.
	 * @return true if a master key.
	 */
	public boolean isMasterKey()
	{
		return (subSigs == null);
	}
	
	/**
	 * Return the algorithm code associated with the public key.
	 * 
	 * @return int
	 */
	public int getAlgorithm()
	{
		return publicPk.getAlgorithm();
	}
	
	/**
	 * Return the strength of the key in bits.
	 * 
	 * @return bit strenght of key.
	 */
	public int getBitStrength()
	{
		return keyStrength;
	}
	
	/**
	 * Return the public key contained in the object.
	 * 
	 * @param provider provider to construct the key for.
	 * 
	 * @return PublicKey
	 * @throws PGPException
	 * @throws NoSuchProviderException
	 */
	public PublicKey getKey(
		String				provider)
		throws PGPException, NoSuchProviderException
	{
		KeyFactory						fact;
		
		try
		{
			switch (publicPk.getAlgorithm())
			{
			case RSA_ENCRYPT:
			case RSA_GENERAL:
				RSAPublicBCPGKey     rsaK = (RSAPublicBCPGKey)publicPk.getKey();
				RSAPublicKeySpec    rsaSpec = new RSAPublicKeySpec(rsaK.getModulus(), rsaK.getPublicExponent());
	
				fact = KeyFactory.getInstance("RSA", provider);
				
				return fact.generatePublic(rsaSpec);
			case DSA:
				DSAPublicBCPGKey	dsaK = (DSAPublicBCPGKey)publicPk.getKey();
				DSAPublicKeySpec 	dsaSpec = new DSAPublicKeySpec(dsaK.getY(), dsaK.getP(), dsaK.getQ(), dsaK.getG());
			
				fact = KeyFactory.getInstance("DSA", provider);
				
				return fact.generatePublic(dsaSpec);
			case ELGAMAL_ENCRYPT:
			case ELGAMAL_GENERAL:
				ElGamalPublicBCPGKey	elK = (ElGamalPublicBCPGKey)publicPk.getKey();
				ElGamalPublicKeySpec	elSpec = new ElGamalPublicKeySpec(elK.getY(), new ElGamalParameterSpec(elK.getP(), elK.getG()));
				
				fact = KeyFactory.getInstance("ElGamal", provider);
				
				return fact.generatePublic(elSpec);
			default:
				throw new PGPException("unknown public key algorithm encountered");
			}
		}
		catch (PGPException e)
		{
			throw e;
		}
		catch (Exception e)
		{
			throw new PGPException("exception constructing public key", e);
		}
	}
	
	/**
	 * Return any userIDs associated with the key.
	 * 
	 * @return an iterator of Strings.
	 */
	public Iterator getUserIDs()
	{
		ArrayList	temp = new ArrayList();
		
		for (int i = 0; i != ids.size(); i++)
		{
			if (ids.get(i) instanceof String)
			{
				temp.add(ids.get(i));
			}
		}
		
		return temp.iterator();
	}
	
	/**
	 * Return any user attribute vectors associated with the key.
	 * 
	 * @return an iterator of Strings.
	 */
	public Iterator getUserAttributes()
	{
		ArrayList	temp = new ArrayList();
		
		for (int i = 0; i != ids.size(); i++)
		{
			if (ids.get(i) instanceof PGPUserAttributeSubpacketVector)
			{
				temp.add(ids.get(i));
			}
		}
		
		return temp.iterator();
	}
	
	public Iterator getSignaturesForID(
		String	id)
	{
		for (int i = 0; i != ids.size(); i++)
		{
			if (id.equals(ids.get(i)))
			{
				return ((ArrayList)idSigs.get(i)).iterator();
			}
		}
		
		return null;
	}
	
	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(publicPk);
		if (trustPk != null)
		{
			out.writePacket(trustPk);
		}
        
        if (subSigs == null)	// 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);
			}
        }
	}
	
	/**
	 * Add a certfication to the given public key.
	 * 
	 * @param key the key the certification is to be added to.
	 * @param id the id the certification is associated with.
	 * @param certification the new certification.
	 * @return the re-certified key.
	 */
	public static PGPPublicKey addCertification(
		PGPPublicKey	key,
		String					id,
		PGPSignature	certification)
	{
		PGPPublicKey	returnKey = new PGPPublicKey(key);
		
		Iterator	idIt = returnKey.getUserIDs();
		ArrayList	sigList = null;
		
		for (int i = 0; i != returnKey.ids.size(); i++)
		{
			if (id.equals(returnKey.ids.get(i)))
			{
				sigList = (ArrayList)returnKey.idSigs.get(i);
			}
		}
		
		if (sigList != null)
		{
			sigList.add(certification);
		}
		else
		{
			sigList = new ArrayList();
			
			sigList.add(certification);
			returnKey.ids.add(id);
			returnKey.idTrusts.add(null);
			returnKey.idSigs.add(sigList);
		}
		
		return returnKey;
	}
	
	/**
	 * Remove any certifications associated with a given id on a key.
	 * 
	 * @param key the key the certifications are to be removed from.
	 * @param id the id that is to be removed.
	 * @return the re-certified key, null if the id was not found on the key.
	 */
	public static PGPPublicKey removeCertification(
		PGPPublicKey	key,
		String					id)
	{
		PGPPublicKey	returnKey = new PGPPublicKey(key);
		Iterator				idIt = returnKey.getUserIDs();
		ArrayList				sigList = null;
		boolean				found = false;
		
		for (int i = 0; i < returnKey.ids.size(); i++)
		{
			if (id.equals(returnKey.ids.get(i)))
			{
				found = true;
				returnKey.ids.remove(i);
				returnKey.idTrusts.remove(i);
				returnKey.idSigs.remove(i);
			}
		}
		
		if (!found)
		{
			return null;
		}
		
		return returnKey;
	}
	
	/**
	 * Remove any certifications associated with a given id on a key.
	 * 
	 * @param key the key the certifications are to be removed from.
	 * @param id the id that the certfication is to be removed from.
	 * @param certification the certfication to be removed.
	 * @return the re-certified key, null if the certification was not found.
	 */
	public static PGPPublicKey removeCertification(
		PGPPublicKey	key,
		String					id,
		PGPSignature	certification)
	{
		PGPPublicKey	returnKey = new PGPPublicKey(key);
		
		Iterator	idIt = returnKey.getUserIDs();
		ArrayList	sigList = null;
		boolean	found = false;
		
		for (int i = 0; i < returnKey.ids.size(); i++)
		{
			if (id.equals(returnKey.ids.get(i)))
			{
				found = ((ArrayList)returnKey.idSigs.get(i)).remove(certification);
			}
		}
		
		if (!found)
		{
			return null;
		}
		
		return returnKey;
	}
}
