package sun.security.ssl;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.crypto.SecretKey;
import javax.net.ssl.SSLException;

import sun.security.ssl.HandshakeMessage.CertificateMsg;
import sun.security.ssl.HandshakeMessage.CertificateRequest;
import sun.security.ssl.HandshakeMessage.ClientHello;
import sun.security.ssl.HandshakeMessage.Finished;
import sun.security.ssl.HandshakeMessage.RSA_ServerKeyExchange;
import sun.security.ssl.HandshakeMessage.ServerHello;
import sun.security.ssl.HandshakeMessage.ServerHelloDone;
import sun.security.ssl.HandshakeMessage.ServerKeyExchange;

/**
 * <p>Java class used to parse TLS/SSL messages.
 * 
 * <p>Contains TLS/SSL handshake messages structures that are filled when using the {@link parseHandshakeMessage } 
 */

/**
 * @author Thibault Lalevée
 *
 */
public class MessageParser {

	public static final int HEADER_SIZE = 5;
	public static final byte TYPE_CLIENT_HELLO = HandshakeMessage.ht_client_hello;
	public static final byte TYPE_SERVER_HELLO = HandshakeMessage.ht_server_hello;
	public static final byte TYPE_HELLO_REQUEST = HandshakeMessage.ht_hello_request;
	public static final byte TYPE_CERTIFICATE = HandshakeMessage.ht_certificate;
	public static final byte TYPE_SERVER_KEYEXCH = HandshakeMessage.ht_server_key_exchange;
	public static final byte TYPE_CERT_RQST = HandshakeMessage.ht_certificate_request;
	public static final byte TYPE_SVR_HELLO_DONE = HandshakeMessage.ht_server_hello_done;
	public static final byte TYPE_CLIENT_KEYEXCH = HandshakeMessage.ht_client_key_exchange;
	public static final byte TYPE_FINISHED = HandshakeMessage.ht_finished;
	public static final byte TYPE_SERVER_HELLO_DONE = HandshakeMessage.ht_server_hello_done;
	public static final byte TYPE_CERTIFICATE_VERIFY = HandshakeMessage.ht_certificate_verify;
	public static final byte TYPE_CHG_CIPHER_SPEC = InputRecord.ct_change_cipher_spec;
	public static final byte TYPE_ALERT = InputRecord.ct_alert;
	public static final byte TYPE_HANDSHAKE = InputRecord.ct_handshake;
	public static final byte TYPE_APP_DATA = InputRecord.ct_application_data;

	private PrivateKey serverPrivateKey;
	private PrivateKey clientPrivateKey;
	private ClientHello clientHello;
	private ServerHello serverHello;
	private CertificateMsg certificateServer;
	private CertificateMsg certificateClient;
	private CertificateRequest certificateRequest;
	private ServerHelloDone serverHelloDone;
	private RSAClientKeyExchange rsaClientKeyExchange;
	private ServerKeyExchange serverKeyExchange;
	private Finished clientFinished;
	private Finished serverFinished;
	private Handshaker clientHandshaker;
	private Handshaker serverHandshaker;
	private MAC clientSigner;
	private CipherBox clientCipher;
	private MAC serverSigner;
	private CipherBox serverCipher;
	private boolean serverFinishedDone;
	private boolean clientFinishedDone;
	private boolean serverChgCipherDone;
	private boolean clientChgCipherDone;
	private SecretKey premaster;
	private SecretKey master;
	private ByteArrayOutputStream xmlOutputStream;
	public int connectionId;

	/**
	 * <p>
	 * Empty constructor of the MessageParser object
	 * 
	 */
	public MessageParser() {
		serverPrivateKey = null;
		clientPrivateKey = null;
		clientHello = null;
		serverHello = null;
		certificateServer = null;
		certificateClient = null;
		certificateRequest = null;
		serverHelloDone = null;
		rsaClientKeyExchange = null;
		serverKeyExchange = null;
		clientFinished = null;
		serverFinished = null;
		clientHandshaker = null;
		serverHandshaker = null;
		clientSigner = MAC.NULL;
		clientCipher = CipherBox.NULL;
		serverSigner = MAC.NULL;
		serverCipher = CipherBox.NULL;
		serverFinishedDone = false;
		clientFinishedDone = false;
		serverChgCipherDone = false;
		clientChgCipherDone = false;
		premaster = null;
		master = null;
		xmlOutputStream = null;
		connectionId = 0;
	}

	/**
	 * <p>
	 * Parse a single record layer message (could be that there are several
	 * handshake messages after a single record header)
	 * 
	 * @param data
	 *            byte array of a single record layer frame
	 * @param dataFromServer
	 *            true if the data comes from server, false instead
	 */
	@SuppressWarnings("restriction")
	public void parseHandshakeMulti(byte[] data, boolean dataFromServer) throws Exception {

		// Create inputstream from data byte array
		InputStream dataStream = new ByteArrayInputStream(data);

		// This outputstream is required to use the HandshakeInStream.read
		// method, not used then
		OutputStream output = new ByteArrayOutputStream();

		// HandshakeInStream with a null handshakeHash (not necessary to
		// decrypt)
		HandshakeInStream input = new HandshakeInStream(null);

		// Put byte values in the InputRecord
		input.r.read(dataStream, output);
		while (input.available() != 0) {
			HandshakeInStream inputBis = parseHandshakeMessage(data, dataFromServer, input);
			input.close();
			input = inputBis;
		}
		if (input != null) {
			try {
				input.close();

			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	/**
	 * <p>
	 * Initialize the handshakers that are used to calculate keys and Mac and
	 * Cipher units
	 * 
	 * @param dataFromServer
	 *            True if data comes from the server, false otherwise
	 * @throws Exception
	 */
	@SuppressWarnings("restriction")
	private void initializeHandshakers(boolean dataFromServer) throws Exception {
		// Need to instantiate an empty context to initialize handshakers
		SecureRandom secure = JsseJce.getSecureRandom();
		SSLContextImpl context = new SSLContextImpl.TLS12Context();
		context.engineInit(null, null, secure);

		// when and only when (not after finished msg) chgcipherspec is
		// done, need to elaborate MAC and Cipher units to decode message
		// MAC and Cipher depends on the handshaker CipherSuite parameter
		// MAC and Cipher are set to their null values (= payload remain
		// unchanged)
		if (serverChgCipherDone && !dataFromServer && !serverFinishedDone) {

			clientHandshaker.sslContext = context;
			serverCipher = clientHandshaker.newReadCipher();
			serverSigner = clientHandshaker.newReadMAC();
		} else if (clientChgCipherDone && dataFromServer && !clientFinishedDone) {

			serverHandshaker.sslContext = context;
			clientCipher = serverHandshaker.newReadCipher();
			clientSigner = serverHandshaker.newReadMAC();
		}
	}

	/**
	 * <p>
	 * Set the handshakers values needed to compute the Mac and Cipher units for
	 * example
	 * 
	 * @param protocolVer
	 *            Protocol version of the SSL exchange
	 */
	@SuppressWarnings("restriction")
	private void setHandshakers(ProtocolVersion protocolVer) {
		// set handshakers with negotiated parameters (serverHello
		// params)
		clientHandshaker = new MyClientHandshaker(null, protocolVer, null);
		clientHandshaker.protocolVersion = protocolVer;
		clientHandshaker.clnt_random = clientHello.clnt_random;
		clientHandshaker.svr_random = serverHello.svr_random;
		clientHandshaker.setCipherSuite(serverHello.cipherSuite);
		clientHandshaker.session = new SSLSessionImpl(serverHello.protocolVersion, serverHello.cipherSuite, null,
				serverHello.sessionId, "127.0.0.1", 10002);

		serverHandshaker = new MyServerHandshaker(null, protocolVer, null);
		serverHandshaker.protocolVersion = protocolVer;
		serverHandshaker.clnt_random = clientHello.clnt_random;
		serverHandshaker.svr_random = serverHello.svr_random;
		serverHandshaker.setCipherSuite(serverHello.cipherSuite);
		serverHandshaker.session = new SSLSessionImpl(serverHello.protocolVersion, serverHello.cipherSuite, null,
				serverHello.sessionId, "127.0.0.1", 10002);
	}

	/**
	 * <p>
	 * Generation of the masterSecret and then the session keys from the
	 * encrypted premasterSecret exctracted from client key exchange message
	 * 
	 * @param input
	 *            Data stream
	 * @param messageLen
	 *            Length of the record layer data
	 * @param protocolVer
	 *            Protocol version of the SSL exchange
	 * @throws IOException
	 * 
	 */
	@SuppressWarnings("restriction")
	private void generateKeys(HandshakeInStream input, int messageLen, ProtocolVersion protocolVer) throws IOException {
		// Get client random as a SecureRandom
		SecureRandom generator = new SecureRandom();
		byte[] gen = clientHello.clnt_random.random_bytes.clone();

		generator.nextBytes(gen);

		// case RSAKeyExchange
		if (clientHandshaker.keyExchange == CipherSuite.KeyExchange.K_RSA) {

			// Parse ClientKeyExchange, create premastersecret from
			// server privateKey
			rsaClientKeyExchange = new RSAClientKeyExchange(serverHello.protocolVersion, clientHello.protocolVersion,
					generator, input, messageLen, serverPrivateKey);

			premaster = rsaClientKeyExchange.preMaster;

			// Calculate the masterSecret from premaster and
			// protocol version
			clientHandshaker.calculateKeys(premaster, protocolVer);
			serverHandshaker.calculateKeys(premaster, protocolVer);
		} else {
			// TODO fill this part to return something like "impossible
			// to decrypt the message"

			throw new SSLException("This key exchange algorithm is not supported");
		}
	}

	/**
	 * <p>
	 * Handles a single message from the {@link data} byte array.
	 * 
	 * <p>
	 * Several cases can occur, if the byte array is:
	 * <p>
	 * 1. handshake or alert message : the informations are parsed and put in
	 * the various {@link HandshakeMessage} of {@link MessageParser}
	 * <p>
	 * 2. appData message : is deciphered
	 * <p>
	 * 3. chgCipherSpec : nothing happens
	 * 
	 * @param data
	 *            handshake message in a byte array form
	 * @param dataFromServer
	 *            True if the data comes from the server, false in the other way
	 * 
	 */
	@SuppressWarnings("restriction")
	public HandshakeInStream parseHandshakeMessage(byte[] data, boolean dataFromServer, HandshakeInStream input)
			throws Exception {

		TlsMsgToXml serializer = new TlsMsgToXml(this);
		// RecordLength value needed for serialization
		byte[] recordLengthByte = new byte[2];
		recordLengthByte[0] = data[3];
		recordLengthByte[1] = data[4];
		int recordLength = java.nio.ByteBuffer.wrap(recordLengthByte).getShort();

		byte messageType, contentType = 0;
		int messageLen = 0;

		ProtocolVersion protocolVer = null;

		// Get the content type (handshake, appdata, chgcipherspec or alert)
		contentType = input.r.contentType();

		// Initialize the handshakers
		initializeHandshakers(dataFromServer);

		// We use the MAC and Cipher units created with the secret keys to
		// decrypt the message
		if (dataFromServer) {
			input.r.decrypt(clientSigner, clientCipher);
		} else {
			input.r.decrypt(serverSigner, serverCipher);
		}

		// Skip the 5 first bytes (position the beginning at the array cell
		// number 4)
		input.mark(4);

		// Record layer content type
		switch (contentType) {

		case TYPE_HANDSHAKE:

			// Read 4 bytes : handshake type + length
			messageType = (byte) input.getInt8();
			messageLen = input.getInt24();

			byte[] version = new byte[2];
			version[0] = data[1];
			version[1] = data[2];

			// Process handshake
			switch (messageType) {
			case TYPE_CLIENT_HELLO:

				// Parse ClientHello
				clientHello = new ClientHello(input, messageLen);

				xmlOutputStream = (ByteArrayOutputStream) serializer.messageToXml(recordLength, contentType, version,
						0x01, dataFromServer);
				xmlOutputStream.close();
				break;

			case TYPE_SERVER_HELLO:

				// parse ServerHello
				serverHello = new ServerHello(input, messageLen);

				// get negotiated protocol version
				protocolVer = serverHello.protocolVersion;

				setHandshakers(protocolVer);

				xmlOutputStream = (ByteArrayOutputStream) serializer.messageToXml(recordLength, contentType, version,
						0x02, dataFromServer);
				xmlOutputStream.close();
				break;

			case TYPE_CERTIFICATE:

				// Parse certificate (server or client)
				if (dataFromServer) {
					certificateClient = new CertificateMsg(input);
				} else {
					certificateServer = new CertificateMsg(input);
				}

				xmlOutputStream = (ByteArrayOutputStream) serializer.messageToXml(recordLength, contentType, version,
						0x0b, dataFromServer);
				xmlOutputStream.close();

				break;

			case TYPE_SERVER_KEYEXCH:

				// Parse serverKeyExchange
				serverKeyExchange = new RSA_ServerKeyExchange(input);

				break;

			case TYPE_CERT_RQST:

				// Parse CertificateRequest
				protocolVer = serverHello.protocolVersion;
				certificateRequest = new CertificateRequest(input, protocolVer);

				break;

			case TYPE_SVR_HELLO_DONE:
				// Parse ServerHelloDone
				serverHelloDone = new ServerHelloDone(input);

				break;

			case TYPE_CLIENT_KEYEXCH:

				generateKeys(input, messageLen, protocolVer);
				xmlOutputStream = (ByteArrayOutputStream) serializer.messageToXml(recordLength, contentType, version,
						0x10, dataFromServer);
				xmlOutputStream.close();
				break;
			case TYPE_FINISHED:

				protocolVer = serverHello.protocolVersion;
				// Parse decrypted finishedMessage (decryption happened
				// earlier with decrypt method)
				if (dataFromServer) {
					clientFinished = new Finished(protocolVer, input, serverHello.cipherSuite);
					// Flag
					clientFinishedDone = true;
				} else {
					serverFinished = new Finished(protocolVer, input, serverHello.cipherSuite);
					// Flag
					serverFinishedDone = true;
				}

				master = clientHandshaker.session.getMasterSecret();

				break;
			default:
				throw new SSLException("Unknown handshake message");

			}

			break;

		case TYPE_APP_DATA:

			// AppData : get the decrypted message (buffer is set so that
			// available method is the length of the decrypted msg)
			byte[] message = new byte[input.r.available()];
			input.read(message);
			String messageString = new String(message, Charset.defaultCharset());
			System.out.println("Decoded message : " + messageString);

			break;
		case TYPE_CHG_CIPHER_SPEC:
			if (dataFromServer) {
				// Flag
				input.read();
				clientChgCipherDone = true;
			} else {
				// Flag
				input.read();
				serverChgCipherDone = true;
			}
			break;
		case TYPE_ALERT:
			// read the 2 bytes of alert
			byte level = (byte) input.read();
			byte code = (byte) input.read();
			// get the string value associated with the byte value
			String alertString = Alerts.alertDescription(code);
			switch (level) {
			case Alerts.alert_warning:
				alertString = "Warning :" + alertString;
				break;
			case Alerts.alert_fatal:

				alertString = "Fatal :" + alertString;
				break;
			default:
				alertString = "Unknown error";
				break;
			}
			break;
		default:
			throw new SSLException("Unknown content type message");
		}

		return input;

	}

	/**
	 * <p>
	 * Handles several messages from the data byte array
	 * 
	 * <p>
	 * Has to be used following the handshake order steps by steps on the same
	 * MessageParser object (First phase 1, then phase 2, ...)
	 * 
	 * @param data
	 *            one or several TLS messages to be parsed/decoded
	 * @param dataFromServer
	 *            True if the data comes from the server, false in the other way
	 * @return Array of the lengths of the different messages that have been
	 *         parsed
	 */
	public List<Integer> parseMsgMultiple(byte[] data, boolean dataFromServer) throws Exception {
		short length = 0;
		List<Integer> handshakesLengths = new ArrayList<Integer>();
		ByteBuffer messageBuffer = ByteBuffer.wrap(data);

		// Check if there is another msg in the frame
		while (messageBuffer.remaining() > 0) {

			// Mark position 1 as the byte to reset to
			messageBuffer.mark();

			// Skip the first byte (content type)
			messageBuffer.get();

			// Skip the next 2 bytes (protocol version)
			messageBuffer.getShort();

			// Get 2 bytes of length
			length = messageBuffer.getShort();

			// Back to byte 1
			messageBuffer.reset();

			// Cut and store the msg to a byte array
			byte[] messageFull = new byte[HEADER_SIZE + length];
			messageBuffer.get(messageFull);

			// Parse the message
			parseHandshakeMulti(messageFull, dataFromServer);

			// Add the msg length to a resizable array
			handshakesLengths.add((int) length + HEADER_SIZE);
		}
		return handshakesLengths;
	}

	/**
	 * <p>
	 * Get the length of a single TLS message
	 * 
	 * @param data
	 *            message data header (5 bytes) as a byte array
	 * @return length of the whole message
	 */
	public static int getMsgLength(byte[] data) {
		short length = 0;
		ByteBuffer messageBuffer = ByteBuffer.wrap(data);
		messageBuffer.get();
		messageBuffer.getShort();
		length = messageBuffer.getShort();

		return length + HEADER_SIZE;
	}

	public ByteArrayOutputStream getXmlStream() {
		return this.xmlOutputStream;
	}

	/**
	 * <p>
	 * Sets the {@link MessageParser} {@link serverPrivateKey} field from a .jks
	 * file
	 * 
	 * @param privateKeyJks
	 *            Jks file containing a PrivateKey
	 * @param pwd
	 *            Password of the Jks file
	 * @throws Exception
	 */
	public void setRsaPrivateKeyFromJks(File privateKeyJks, String pwd) throws UnrecoverableKeyException, IOException {
		FileInputStream fis = new FileInputStream(privateKeyJks);
		try {
			KeyStore ks = KeyStore.getInstance("JKS");
			String alias = null;

			ks.load(fis, pwd.toCharArray());

			Enumeration aliases = ks.aliases();
			for (; aliases.hasMoreElements();) {
				alias = (String) aliases.nextElement();

				// Does alias refer to a private key?
				boolean b = ks.isKeyEntry(alias);

				// Does alias refer to a trusted certificate?
				b = ks.isCertificateEntry(alias);
			}

			serverPrivateKey = (PrivateKey) ks.getKey(alias, pwd.toCharArray());
		} catch (KeyStoreException e) {

		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (CertificateException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} finally {
			try {
				fis.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	public ClientHello getClientHello() {
		return clientHello;
	}

	public ServerHello getServerHello() {
		return serverHello;
	}

	public CertificateMsg getCertificateServer() {
		return certificateServer;
	}

	public CertificateMsg getCertificateClient() {
		return certificateClient;
	}

	public CertificateRequest getCertificateRequest() {
		return certificateRequest;
	}

	public ServerHelloDone getServerHelloDone() {
		return serverHelloDone;
	}

	public RSAClientKeyExchange getRsaClientKeyExchange() {
		return rsaClientKeyExchange;
	}

	public ServerKeyExchange getServerKeyExchange() {
		return serverKeyExchange;
	}

	public Finished getClientFinished() {
		return clientFinished;
	}

	public Finished getServerFinished() {
		return serverFinished;
	}

	public PrivateKey getServerPrivateKey() {
		return serverPrivateKey;
	}

	public PrivateKey getClientPrivateKey() {
		return clientPrivateKey;
	}

	public SecretKey getPremaster() {
		return premaster;
	}

	public SecretKey getMaster() {
		return master;
	}

}