package aefs.protocols.authorization;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.security.spec.InvalidKeySpecException;
import java.util.Arrays;
import java.util.Set;

import javax.crypto.NoSuchPaddingException;

import misc.io.ArrayUtils;
import misc.io.StreamUtils;
import misc.logging.SimpleLogger;
import misc.numbers.BigIntegerUtilities;
import abe.Ciphertext;
import abe.MasterPublicParameters;
import abe.PrivateKey;
import abe.schemes.waters08.Waters08ABEScheme;
import abe.schemes.waters08.Waters08Ciphertext;
import aefs.encryption.RandomlyPaddedDecryptionStream;
import aefs.encryption.RandomlyPaddedEncryptionStream;
import aefs.encryption.rsa.GenericRSAKey;

/**
 * Token generated by a token authority which contains an encrypted
 * session communication keys, expiration information, and a list of 
 * authorized attributes.
 *
 * @author Aubrey Alston (ada2145@columbia.edu)
 */
public class MasterSessionToken {
	
	public Ciphertext encryptedServerKey;
	public long expiration;
	
	public byte[] signature;
	public GenericRSAKey verificationKey;
	
	public byte[] challenge;
	public byte[] token;
	
	public MasterSessionToken(){ }
	
	/**
	 * Creates a MasterSessionToken.
	 * @param sessionKey session key to be used between server and client
	 * @param serverKey server key to be used by the server
	 * @param encryptedServerKey attribute-encrypted server key
	 * @param signingKey public key to use to sign the token
	 * @param verificationKey key to use to verify the signature on the token
	 * @param expiration expiration time of this token
	 * @param attributes attributes authorized by this token
	 * @param rng random number generator to use
	 */
	public MasterSessionToken(byte[] sessionKey, byte[] serverKey,
			Ciphertext encryptedServerKey, GenericRSAKey signingKey, 
			GenericRSAKey verificationKey, long expiration, int[] attributes, SecureRandom rng)
				throws NoSuchAlgorithmException, IOException, SignatureException, KeyException,
					InvalidKeySpecException, NoSuchPaddingException, NoSuchAlgorithmException,
					InvalidAlgorithmParameterException {
		
		challenge = new byte[32];
		rng.nextBytes(challenge);
		
		this.encryptedServerKey = encryptedServerKey;
		this.expiration = expiration;
		// E_s(K_s), expiration, challenge
		signature = signingKey.sign(StreamUtils.combineByteArrays(new byte[][]{ 
				encryptedServerKey.serialize(),
				BigIntegerUtilities.longToBytes(expiration),
				challenge
		}));
		
		this.verificationKey = verificationKey;
		
		// todo: encryption stream
		
		byte[] key = ArrayUtils.copyOfRange(serverKey, 0, 32);
		byte[] iv = ArrayUtils.copyOfRange(serverKey, 32, 48);
		
		// generate token
		byte[] pad = new byte[16];
		rng.nextBytes(pad);
		
		ByteArrayOutputStream tokenOut = new ByteArrayOutputStream();
		RandomlyPaddedEncryptionStream eStream = new RandomlyPaddedEncryptionStream(tokenOut,
				key, iv, pad);
		DataOutputStream dout = new DataOutputStream(eStream);
		
		try{
			// (K1 | expiry | A | challenge)
			StreamUtils.writeAdvertisedBytes(dout, sessionKey);
			dout.writeLong(expiration);
			dout.writeInt(attributes.length);
			for(int i = 0; i < attributes.length; ++i){
				dout.writeInt(attributes[i]);
			}
			StreamUtils.writeAdvertisedBytes(dout, challenge);
		} finally{
			eStream.close();
		}
		
		this.token = tokenOut.toByteArray();
	}
	
	public TokenDetails verifyFully(String source, Set<GenericRSAKey> trustedVerifiers,
			MasterPublicParameters context, PrivateKey serverABEKey){
		if(! trustedVerifiers.contains(verificationKey)){
			SimpleLogger.error("Token received from "+source+
					" signed by untrusted token authority.");
			return null;
		} if(expiration < System.currentTimeMillis()){
			SimpleLogger.error("Expired token received from "+source+".");
			return null;
		}
		
		// verify validity of token
		try{
			byte[] toVerify = getBytesToVerify();
			if(! verificationKey.verify(toVerify, signature)){
				SimpleLogger.error("Signature verification failed for token received from "+source);
				return null;
			}
		} catch(Exception e){ 
			SimpleLogger.error("Unable to verify public properties in token received from "+source); 
			return null;
		}
		
		// decrypt encrypted section; test challenge
		Waters08ABEScheme s = new Waters08ABEScheme();
		try{
			s.setPublicParameters(context);
			byte[] o = s.decrypt(encryptedServerKey, serverABEKey);
			
			byte[] selfKey = ArrayUtils.copyOfRange(o, 0, 32);
			byte[] selfIV = ArrayUtils.copyOfRange(o, 32, 48);
			
			TokenDetails details = accessEncryptedDetails(selfKey, selfIV);
			if(details == null){
				SimpleLogger.error("Unable to retrieve token details from token sent by "+source);
			}if(! Arrays.equals(challenge, details.challenge) ){
				SimpleLogger.error("Token received from "+source+" contains mismatched "+
						"challenge bytes.");
				return null;
			} if( details.expiration != expiration ){
				SimpleLogger.error("Token received from "+source+" contains mismatched "+
						"expiration fields.");
				return null;
			}
			
			return details;
		} catch(Exception e){ 
			SimpleLogger.error("Unable to verify public properties in token received from "+source); 
			return null;
		}
	
	}
	
	public TokenDetails accessEncryptedDetails(byte[] key, byte[] iv){
		RandomlyPaddedDecryptionStream eStream = null;
		try{
			ByteArrayInputStream tokenIn = new ByteArrayInputStream(this.token);
			eStream = new RandomlyPaddedDecryptionStream(tokenIn,
					key, iv, true);
			DataInputStream din = new DataInputStream(eStream);
			
			byte[] sessionKey = StreamUtils.readAdvertisedBytes(din);
			long expiration = din.readLong();
			int[] attributes = new int[din.readInt()];
			for(int i = 0; i < attributes.length; ++i){
				attributes[i] = din.readInt();
			}
			byte[] challenge = StreamUtils.readAdvertisedBytes(din);
			
			return new TokenDetails(sessionKey, expiration, attributes,
					challenge);
		} catch(Exception e){
			return null;
		}finally{
			try{
				eStream.close();
			} catch(Exception e){ }
		}
	}
	
	public byte[] getBytesToVerify() throws IOException{
		return StreamUtils.combineByteArrays(new byte[][]{ 
				encryptedServerKey.serialize(),
				BigIntegerUtilities.longToBytes(expiration),
				challenge
		});
	}
	
	/**
	 * Serialize this session token to a byte array.
	 */
	public byte[] serialize() throws IOException{
		ByteArrayOutputStream b = new ByteArrayOutputStream();
		DataOutputStream out = new DataOutputStream(b);
		
		b.write(encryptedServerKey.serialize());
		out.writeLong(expiration);
		StreamUtils.writeAdvertisedBytes(out, signature);
		out.write(verificationKey.serialize());
		StreamUtils.writeAdvertisedBytes(out, challenge);
		StreamUtils.writeAdvertisedBytes(out, token);
		
		return b.toByteArray();
	}
	
	/**
	 * Deserializes this session token from an input source.
	 * @param source InputStream source
	 * @param context public parameters context
	 */
	public void deserialize(InputStream source, MasterPublicParameters context)
			throws IOException{
		DataInputStream in = new DataInputStream(source);
		
		encryptedServerKey = new Waters08Ciphertext();
		encryptedServerKey.deserialize(source, context);
		expiration = in.readLong();
		
		signature = StreamUtils.readAdvertisedBytes(in);
		
		verificationKey = new GenericRSAKey();
		verificationKey.deserialize(source);
		
		challenge = StreamUtils.readAdvertisedBytes(in);
		token = StreamUtils.readAdvertisedBytes(in);
	}
	
	public void printSignature(){
		for(int i = 0; i < signature.length; i++){
			System.out.print(signature[i]+" ");
		}
		System.out.println();
	}
	
	/**
	 * Decrypted token contents (communication key|expiration|attribute list|challenge)
	 */
	public static class TokenDetails{
		
		public byte[] keyMaterial;
		public long expiration;
		public int[] attributes;
		public byte[] challenge;
		
		public TokenDetails(byte[] keyMaterial, long expiration, int[] attributes,
				byte[] challenge){
			this.keyMaterial = keyMaterial;
			this.expiration = expiration;
			this.attributes = attributes;
			this.challenge = challenge;
		}
	}
}
