package org.ourfilesystem.security;

/*
OurFileSystem is a peer2peer file sharing program.
Copyright (C) 2012  Robert Gass

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

*/

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.security.SecureRandom;
import java.util.Arrays;

import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SHA512Digest;
import org.bouncycastle.crypto.engines.RSAEngine;
import org.bouncycastle.crypto.generators.RSAKeyPairGenerator;
import org.bouncycastle.crypto.params.RSAKeyGenerationParameters;
import org.bouncycastle.crypto.params.RSAKeyParameters;
import org.bouncycastle.crypto.params.RSAPrivateCrtKeyParameters;
import org.bouncycastle.crypto.encodings.PKCS1Encoding;
import org.ourfilesystem.db.Peer;
import org.ourfilesystem.db.Post;
import org.ourfilesystem.utilities.BBytes;
import org.ourfilesystem.utilities.FileUtils;

public class SecurityTools {

	/*
	 * SecureRandom should be thread safe. 
	 */
	public static SecureRandom Random = new SecureRandom();
	
	public static AsymmetricCipherKeyPair  generateKeyPair() {
		/*
		 * The probability that the new BigInteger represents a prime number will exceed (1 - 1/(2^certainty))
		 * 2^40 > 1T 
		 */
		RSAKeyGenerationParameters parms = new RSAKeyGenerationParameters(BigInteger.valueOf(65537),   // exponent
																		  Random,                      // Random generator
																		  2048,                        // Key size 
																		  40);                         // Certainty
		RSAKeyPairGenerator gen = new RSAKeyPairGenerator();
		gen.init(parms);	
		return gen.generateKeyPair();
	}

	public static void writeRSAPublicKey(RSAKeyParameters p, OutputStream fos) throws IOException {
		byte exp[] = p.getExponent().toByteArray();
		byte mod[] = p.getModulus().toByteArray();
		FileUtils.writeBytes(fos, exp);
		FileUtils.writeBytes(fos, mod);
	}
	
	public static RSAKeyParameters readRSAPublicKey(InputStream fis) throws IOException {
		byte [] exp = FileUtils.readBytes(fis);
		byte [] mod = FileUtils.readBytes(fis);
		RSAKeyParameters pub = new RSAKeyParameters(false,                 //boolean isPrivate,            
													new BigInteger(mod),   //java.math.BigInteger modulus, 
													new BigInteger(exp));  //java.math.BigInteger exponent 
		return pub;
	}
	
	public static void writeRSAPrivateKey(RSAPrivateCrtKeyParameters p, OutputStream fos) throws IOException {
		FileUtils.writeBytes(fos, p.getDP().toByteArray());
		FileUtils.writeBytes(fos, p.getDQ().toByteArray());
		FileUtils.writeBytes(fos, p.getExponent().toByteArray());
		FileUtils.writeBytes(fos, p.getModulus().toByteArray());
		FileUtils.writeBytes(fos, p.getP().toByteArray());
		FileUtils.writeBytes(fos, p.getPublicExponent().toByteArray());
		FileUtils.writeBytes(fos, p.getQ().toByteArray());
		FileUtils.writeBytes(fos, p.getQInv().toByteArray());
	}
	
	public static RSAPrivateCrtKeyParameters readRSAPrivateKey(InputStream fis) throws IOException {
		byte [] dp = FileUtils.readBytes(fis);
		byte [] dq = FileUtils.readBytes(fis);
		byte [] exp = FileUtils.readBytes(fis);
		byte [] mod = FileUtils.readBytes(fis);
		byte [] p = FileUtils.readBytes(fis);
		byte [] pe = FileUtils.readBytes(fis);
		byte [] q = FileUtils.readBytes(fis);
		byte [] qi = FileUtils.readBytes(fis);
		RSAPrivateCrtKeyParameters priv = new RSAPrivateCrtKeyParameters(new BigInteger(mod),   //java.math.BigInteger modulus,          
																		 new BigInteger(pe),    //java.math.BigInteger publicExponent,   
																		 new BigInteger(exp),   //java.math.BigInteger privateExponent,  
																		 new BigInteger(p),     //java.math.BigInteger p,                
																		 new BigInteger(q),     //java.math.BigInteger q,                
																		 new BigInteger(dp),    //java.math.BigInteger dP,               
																		 new BigInteger(dq),    //java.math.BigInteger dQ,               
																		 new BigInteger(qi));   //java.math.BigInteger qInv              
		return priv;
	}
	
	public static void writePublicKeySet(PublicKeySet k, OutputStream fos) throws IOException {
		writeRSAPublicKey((RSAKeyParameters)k.getPublicEncryptionKey(), fos);
		writeRSAPublicKey((RSAKeyParameters)k.getPublicSigningKey(), fos);
	}
	
	public static PublicKeySet readPublicKeySet(InputStream fis) throws IOException {
		PublicKeySet p = new PublicKeySet();
		p.setPublicEncryptionKey(readRSAPublicKey(fis));
		p.setPublicSigningKey(readRSAPublicKey(fis));
		return p;
	}
	
	public static void writeSignedDigest(SignedDigest s, OutputStream fos) throws IOException {
		FileUtils.writeBytes(fos, ((BBytes)s.getDigest()).getBytes());
		if (s.getPeerIdentifier() == null) {
			fos.write(2);
		}
		else {
			fos.write(1);
			FileUtils.writeBytes(fos, ((BBytes)s.getPeerIdentifier()).getBytes());
		}
		FileUtils.writeBytes(fos, ((BBytes)s.getSignature()).getBytes());
	}
	
	public static SignedDigest readSignedDigest(InputStream fis) throws IOException {
		SignedDigest s = new SignedDigest();
		s.setDigest(new BBytes(FileUtils.readBytes(fis)));
		int v = fis.read();
		if (v == 2) { s.setPeerIdentifier(null); }
		else { s.setPeerIdentifier(new BBytes(FileUtils.readBytes(fis))); }
		s.setSignature(new BBytes(FileUtils.readBytes(fis)));
		return s;
	}
	
	public static void writePublicKeySetSigned(PublicKeySetSigned k, OutputStream fos) throws IOException {
		writePublicKeySet(k, fos);
		writeSignedDigest(k.getSignature(), fos);
	}
	
	public static PublicKeySetSigned readPublicKeySetSigned(InputStream fis) throws IOException {
		PublicKeySet p = readPublicKeySet(fis);
		PublicKeySetSigned ps = new PublicKeySetSigned();
		ps.setPublicEncryptionKey(p.getPublicEncryptionKey());
		ps.setPublicSigningKey(p.getPublicSigningKey());
		ps.setSignature(readSignedDigest(fis));
		return ps;
	}
	
	public static void writeKeySet(KeySet k, OutputStream fos) throws IOException {
		writeRSAPrivateKey((RSAPrivateCrtKeyParameters) k.getPrivateEncryptionKey(), fos);
		writeRSAPrivateKey((RSAPrivateCrtKeyParameters) k.getPrivateSigningKey(), fos);
		writePublicKeySet(k.getPublicKeySet(), fos);
	}
	
	public static KeySet readKeySet(InputStream fis) throws IOException {
		KeySet k = new KeySet();
		k.setPrivateEncryptionKey(readRSAPrivateKey(fis));
		k.setPrivateSigningKey(readRSAPrivateKey(fis));
		k.setPublicKeySet(readPublicKeySet(fis));
		return k;
	}
	
	public static void DigestRSAPublicKey(Digest dig, RSAKeyParameters pub) {
		byte [] a = null;
		a = pub.getExponent().toByteArray();
		dig.update(a, 0, a.length);
		a = pub.getModulus().toByteArray();
		dig.update(a, 0, a.length);
	}
	
	public static byte[] digestPublicKey(PublicKeySet k) {
		SHA512Digest dig = new SHA512Digest();
		DigestRSAPublicKey(dig, (RSAKeyParameters)k.getPublicEncryptionKey());
		DigestRSAPublicKey(dig, (RSAKeyParameters)k.getPublicSigningKey());
		byte sig[] = new byte[dig.getDigestSize()];
		dig.doFinal(sig, 0);
		return sig;
	}
	
	public static SignedDigest signDigest(byte dig[], BBytes peerid, RSAPrivateCrtKeyParameters key) {
		SignedDigest s = new SignedDigest();
		RSAEngine eng = new RSAEngine();
		PKCS1Encoding enc = new PKCS1Encoding(eng);
		enc.init(true, key);
		try {
			byte sig[] = enc.processBlock(dig, 0, dig.length);
			s.setDigest(new BBytes(dig));
			s.setPeerIdentifier(peerid);
			s.setSignature(new BBytes(sig));
			return s;
		} catch (InvalidCipherTextException e) {
			e.printStackTrace();
			throw new RuntimeException(e.getMessage());
		}
	}
	
	public static boolean verifySignedDigest(SignedDigest s, byte checkdig[], RSAKeyParameters pub) {
		if (!Arrays.equals(((BBytes)s.getDigest()).getBytes(), checkdig)) return false;   
		RSAEngine eng = new RSAEngine();
		PKCS1Encoding enc = new PKCS1Encoding(eng);
		enc.init(false, pub);
		byte sig[] = ((BBytes)s.getSignature()).getBytes();
		try {
			byte decsig[] = enc.processBlock(sig, 0, sig.length);
			return Arrays.equals(decsig, checkdig);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return false;
	}
	
	public static BBytes digestFile(File f) throws IOException {
		SHA512Digest d = new SHA512Digest();
		digestFile(d, f);
		byte sig[] = new byte[d.getDigestSize()];
		d.doFinal(sig, 0);
		return new BBytes(sig);
	}
	
	public static void digestFile(Digest d, File f) throws IOException {
		FileInputStream fis = new FileInputStream(f);
		byte buf[] = new byte[1024];
		int len = fis.read(buf);
		while (len != -1){
			if (len > 0) {
				d.update(buf, 0, len);
			}
			len = fis.read(buf);
		}
		fis.close();
	}
	
	public static void digestLong(Digest d, long val) {
		byte b[] = new byte[Long.SIZE/Byte.SIZE];
		ByteBuffer buf = ByteBuffer.wrap(b);
		buf.putLong(val);
		d.update(b, 0, b.length);
	}
	
	public static void digestString(Digest d, String str) {
		if (str != null) {
			byte strb[] = str.getBytes(Charset.forName("UTF-16BE"));
			d.update(strb, 0, strb.length);
		}
	}
	
	public static BBytes digestPeerLocation(Peer p) {
		SHA512Digest d = new SHA512Digest();
		digestString(d, p.getIntroduction());
		digestString(d, (String)p.getLocation());
		digestString(d, (String)p.getNickname());
		digestLong(d, p.getUpdateCount());
		byte sig[] = new byte[d.getDigestSize()];
		d.doFinal(sig, 0);
		return new BBytes(sig);
	}
	
	public static BBytes digestPost(Post p, BBytes peerid) throws IOException {
		BBytes fr = (BBytes)p.getFileReferenceDigest();
		File f = (File)p.getMessage();
		long pnum = p.getPostNumber();
		SHA512Digest d = new SHA512Digest();
		if (fr != null) {
			if (fr.getBytes() != null) {
				d.update(fr.getBytes(), 0, fr.getBytes().length);
			}
		}
		if (f != null) {
			if (f.exists() && f.isFile()) {
				digestFile(d, f);
			}
		}
		digestLong(d, pnum);
		if (p.isPosterHasFile()) { d.update((byte)1); }
		else { d.update((byte)0); }
		if (peerid != null) {
			d.update(peerid.getBytes(), 0, peerid.getBytes().length);
		}
		byte sig[] = new byte[d.getDigestSize()];
		d.doFinal(sig, 0);
		return new BBytes(sig);
	}
	
	public static void signPost(Post p, BBytes peerid, RSAPrivateCrtKeyParameters key) throws IOException {
		BBytes dig = digestPost(p, peerid);
		p.setSignedDigest(signDigest(dig.getBytes(), peerid, key));
	}
	
	public static boolean verifyPost(Post p, RSAKeyParameters pub) throws IOException {
		BBytes dig = digestPost(p, (BBytes)p.getSignedDigest().getPeerIdentifier());
		return verifySignedDigest(p.getSignedDigest(), dig.getBytes(), pub);
	}
	
	public static PublicKeySetSigned signPublicKeySet(PublicKeySet k, BBytes peerid, RSAPrivateCrtKeyParameters key) {
		PublicKeySetSigned s = new PublicKeySetSigned();
		s.setPublicEncryptionKey(k.getPublicEncryptionKey());
		s.setPublicSigningKey(k.getPublicSigningKey());
		byte dig[] = digestPublicKey(k);
		s.setSignature(signDigest(dig, peerid, key));
		return s;
	}
	
	public static boolean verifyPublicKeySetSigned(PublicKeySetSigned sk, RSAKeyParameters pubkey) {
		byte cdig[] = digestPublicKey(sk);
		return verifySignedDigest(sk.getSignature(), cdig, pubkey);
	}
	
}
