package iotcloud;

import java.io.*;
import java.net.*;
import java.util.Arrays;
import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.SecureRandom;

/**
 * This class provides a communication API to the webserver.  It also
 * validates the HMACs on the slots and handles encryption.
 * @author Brian Demsky <bdemsky@uci.edu>
 * @version 1.0
 */


class CloudComm {
	private static final int SALT_SIZE = 8;
	private static final int TIMEOUT_MILLIS = 25; // 100

	/** Sets the size for the HMAC. */
	static final int HMAC_SIZE = 32;

	private String baseurl;
	private Cipher encryptCipher;
	private Cipher decryptCipher;
	private Mac mac;
	private String password;
	private SecureRandom random;
	private byte salt[];
	private Table table;
	private int listeningPort = -1;
	private Thread localServerThread = null;
	private boolean doEnd = false;

	/**
	 * Empty Constructor needed for child class.
	 */
	CloudComm() {
	}

	/**
	 * Constructor for actual use. Takes in the url and password.
	 */
	CloudComm(Table _table,  String _baseurl, String _password, int _listeningPort) {
		this.table = _table;
		this.baseurl = _baseurl;
		this.password = _password;
		this.random = new SecureRandom();
		this.listeningPort = _listeningPort;

		if (this.listeningPort > 0) {
			localServerThread = new Thread(new Runnable() {
				public void run() {
					localServerWorkerFunction();
				}
			});
			localServerThread.start();
		}
	}

	/**
	 * Generates Key from password.
	 */
	private SecretKeySpec initKey() {
		try {
			PBEKeySpec keyspec = new PBEKeySpec(password.toCharArray(),
			                                    salt,
			                                    65536,
			                                    128);
			SecretKey tmpkey = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256").generateSecret(keyspec);
			return new SecretKeySpec(tmpkey.getEncoded(), "AES");
		} catch (Exception e) {
			e.printStackTrace();
			throw new Error("Failed generating key.");
		}
	}

	/**
	 * Inits all the security stuff
	 */
	public void initSecurity() throws ServerException {
		// try to get the salt and if one does not exist set one
		if (!getSalt()) {
			//Set the salt
			setSalt();
		}

		initCrypt();
	}

	/**
	 * Inits the HMAC generator.
	 */
	private void initCrypt() {

		if (password == null) {
			return;
		}

		try {
			SecretKeySpec key = initKey();
			password = null; // drop password
			mac = Mac.getInstance("HmacSHA256");
			mac.init(key);
			encryptCipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
			encryptCipher.init(Cipher.ENCRYPT_MODE, key);
			decryptCipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
			decryptCipher.init(Cipher.DECRYPT_MODE, key);
		} catch (Exception e) {
			e.printStackTrace();
			throw new Error("Failed To Initialize Ciphers");
		}
	}

	/*
	 * Builds the URL for the given request.
	 */
	private URL buildRequest(boolean isput, long sequencenumber, long maxentries) throws IOException {
		String reqstring = isput ? "req=putslot" : "req=getslot";
		String urlstr = baseurl + "?" + reqstring + "&seq=" + sequencenumber;
		if (maxentries != 0)
			urlstr += "&max=" + maxentries;
		return new URL(urlstr);
	}

	private void setSalt() throws ServerException {

		if (salt != null) {
			// Salt already sent to server so dont set it again
			return;
		}

		try {
			byte[] saltTmp = new byte[SALT_SIZE];
			random.nextBytes(saltTmp);

			URL url = new URL(baseurl + "?req=setsalt");
			URLConnection con = url.openConnection();
			HttpURLConnection http = (HttpURLConnection) con;

			http.setRequestMethod("POST");
			http.setFixedLengthStreamingMode(saltTmp.length);
			http.setDoOutput(true);
			http.setConnectTimeout(TIMEOUT_MILLIS);
			http.connect();

			OutputStream os = http.getOutputStream();
			os.write(saltTmp);
			os.flush();

			int responsecode = http.getResponseCode();
			if (responsecode != HttpURLConnection.HTTP_OK) {
				// TODO: Remove this print
				System.out.println(responsecode);
				throw new Error("Invalid response");
			}

			salt = saltTmp;
		} catch (Exception e) {
			// e.printStackTrace();
			throw new ServerException("Failed setting salt", ServerException.TypeConnectTimeout);
		}
	}

	private boolean getSalt() throws ServerException {
		URL url = null;
		URLConnection con = null;
		HttpURLConnection http = null;

		try {
			url = new URL(baseurl + "?req=getsalt");
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("getSlot failed");
		}
		try {

			con = url.openConnection();
			http = (HttpURLConnection) con;
			http.setRequestMethod("POST");
			http.setConnectTimeout(TIMEOUT_MILLIS);
			http.setReadTimeout(TIMEOUT_MILLIS);
			http.connect();
		} catch (SocketTimeoutException e) {
			throw new ServerException("getSalt failed", ServerException.TypeConnectTimeout);
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("getSlot failed");
		}

		try {

			int responsecode = http.getResponseCode();
			if (responsecode != HttpURLConnection.HTTP_OK) {
				// TODO: Remove this print
				// System.out.println(responsecode);
				throw new Error("Invalid response");
			}

			InputStream is = http.getInputStream();
			if (is.available() > 0) {
				DataInputStream dis = new DataInputStream(is);
				int salt_length = dis.readInt();
				byte [] tmp = new byte[salt_length];
				dis.readFully(tmp);
				salt = tmp;
				return true;
			} else {
				return false;
			}
		} catch (SocketTimeoutException e) {
			throw new ServerException("getSalt failed", ServerException.TypeInputTimeout);
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("getSlot failed");
		}
	}

	/*
	 * API for putting a slot into the queue.  Returns null on success.
	 * On failure, the server will send slots with newer sequence
	 * numbers.
	 */
	public Slot[] putSlot(Slot slot, int max) throws ServerException {
		URL url = null;
		URLConnection con = null;
		HttpURLConnection http = null;

		try {
			if (salt == null) {
				if (!getSalt()) {
					throw new ServerException("putSlot failed", ServerException.TypeSalt);
				}
				initCrypt();
			}

			long sequencenumber = slot.getSequenceNumber();
			byte[] bytes = slot.encode(mac);
			bytes = encryptCipher.doFinal(bytes);

			url = buildRequest(true, sequencenumber, max);
			con = url.openConnection();
			http = (HttpURLConnection) con;

			http.setRequestMethod("POST");
			http.setFixedLengthStreamingMode(bytes.length);
			http.setDoOutput(true);
			http.setConnectTimeout(TIMEOUT_MILLIS);
			http.setReadTimeout(TIMEOUT_MILLIS);
			http.connect();

			OutputStream os = http.getOutputStream();
			os.write(bytes);
			os.flush();

			// System.out.println("Bytes Sent: " + bytes.length);
		} catch (ServerException e) {
			throw e;
		} catch (SocketTimeoutException e) {
			throw new ServerException("putSlot failed", ServerException.TypeConnectTimeout);
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("putSlot failed");
		}



		try {
			InputStream is = http.getInputStream();
			DataInputStream dis = new DataInputStream(is);
			byte[] resptype = new byte[7];
			dis.readFully(resptype);

			if (Arrays.equals(resptype, "getslot".getBytes()))
				return processSlots(dis);
			else if (Arrays.equals(resptype, "putslot".getBytes()))
				return null;
			else
				throw new Error("Bad response to putslot");

		} catch (SocketTimeoutException e) {
			throw new ServerException("putSlot failed", ServerException.TypeInputTimeout);
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("putSlot failed");
		}
	}

	/**
	 * Request the server to send all slots with the given
	 * sequencenumber or newer.
	 */
	public Slot[] getSlots(long sequencenumber) throws ServerException {
		URL url = null;
		URLConnection con = null;
		HttpURLConnection http = null;

		try {
			if (salt == null) {
				if (!getSalt()) {
					throw new ServerException("getSlots failed", ServerException.TypeSalt);
				}
				initCrypt();
			}

			url = buildRequest(false, sequencenumber, 0);
			con = url.openConnection();
			http = (HttpURLConnection) con;
			http.setRequestMethod("POST");
			http.setConnectTimeout(TIMEOUT_MILLIS);
			http.setReadTimeout(TIMEOUT_MILLIS);
			http.connect();
		} catch (SocketTimeoutException e) {
			throw new ServerException("getSlots failed", ServerException.TypeConnectTimeout);
		} catch (ServerException e) {
			throw e;
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("getSlots failed");
		}

		try {
			InputStream is = http.getInputStream();
			DataInputStream dis = new DataInputStream(is);
			byte[] resptype = new byte[7];
			dis.readFully(resptype);
			if (!Arrays.equals(resptype, "getslot".getBytes()))
				throw new Error("Bad Response: " + new String(resptype));
			else
				return processSlots(dis);
		} catch (SocketTimeoutException e) {
			throw new ServerException("getSlots failed", ServerException.TypeInputTimeout);
		} catch (Exception e) {
			// e.printStackTrace();
			throw new Error("getSlots failed");
		}
	}

	/**
	 * Method that actually handles building Slot objects from the
	 * server response.  Shared by both putSlot and getSlots.
	 */
	private Slot[] processSlots(DataInputStream dis) throws Exception {
		int numberofslots = dis.readInt();
		int[] sizesofslots = new int[numberofslots];

		Slot[] slots = new Slot[numberofslots];
		for (int i = 0; i < numberofslots; i++)
			sizesofslots[i] = dis.readInt();

		for (int i = 0; i < numberofslots; i++) {

			byte[] data = new byte[sizesofslots[i]];
			dis.readFully(data);

			data = decryptCipher.doFinal(data);

			slots[i] = Slot.decode(table, data, mac);
		}
		dis.close();
		return slots;
	}

	public byte[] sendLocalData(byte[] sendData, String host, int port) {

		if (salt == null) {
			return null;
		}
		try {

			System.out.println("Passing Locally");

			mac.update(sendData);
			byte[] genmac = mac.doFinal();
			byte[] totalData = new byte[sendData.length + genmac.length];
			System.arraycopy(sendData, 0, totalData, 0, sendData.length);
			System.arraycopy(genmac, 0, totalData, sendData.length, genmac.length);

			// Encrypt the data for sending
			// byte[] encryptedData = encryptCipher.doFinal(totalData);
			byte[] encryptedData = encryptCipher.doFinal(totalData);

			// Open a TCP socket connection to a local device
			Socket socket = new Socket(host, port);
			socket.setReuseAddress(true);
			DataOutputStream output = new DataOutputStream(socket.getOutputStream());
			DataInputStream input = new DataInputStream(socket.getInputStream());

			// Send data to output (length of data, the data)
			output.writeInt(encryptedData.length);
			output.write(encryptedData, 0, encryptedData.length);
			output.flush();

			int lengthOfReturnData = input.readInt();
			byte[] returnData = new byte[lengthOfReturnData];
			input.readFully(returnData);
			returnData = decryptCipher.doFinal(returnData);

			// We are done with this socket
			socket.close();

			mac.update(returnData, 0, returnData.length - HMAC_SIZE);
			byte[] realmac = mac.doFinal();
			byte[] recmac = new byte[HMAC_SIZE];
			System.arraycopy(returnData, returnData.length - realmac.length, recmac, 0, realmac.length);

			if (!Arrays.equals(recmac, realmac))
				throw new Error("Local Error: Invalid HMAC!  Potential Attack!");

			byte[] returnData2 = new byte[lengthOfReturnData - recmac.length];
			System.arraycopy(returnData, 0, returnData2, 0, returnData2.length);

			return returnData2;
		} catch (SocketTimeoutException e) {

		} catch (BadPaddingException e) {

		} catch (IllegalBlockSizeException e) {

		} catch (UnknownHostException e) {

		} catch (IOException e) {

		}

		return null;
	}

	private void localServerWorkerFunction() {

		ServerSocket inputSocket = null;

		try {
			// Local server socket
			inputSocket = new ServerSocket(listeningPort);
			inputSocket.setReuseAddress(true);
			inputSocket.setSoTimeout(TIMEOUT_MILLIS);
		} catch (Exception e) {
			e.printStackTrace();
			throw new Error("Local server setup failure...");
		}

		while (!doEnd) {

			try {
				// Accept incoming socket
				Socket socket = inputSocket.accept();

				DataInputStream input = new DataInputStream(socket.getInputStream());
				DataOutputStream output = new DataOutputStream(socket.getOutputStream());

				// Get the encrypted data from the server
				int dataSize = input.readInt();
				byte[] readData = new byte[dataSize];
				input.readFully(readData);

				// Decrypt the data
				readData = decryptCipher.doFinal(readData);

				mac.update(readData, 0, readData.length - HMAC_SIZE);
				byte[] genmac = mac.doFinal();
				byte[] recmac = new byte[HMAC_SIZE];
				System.arraycopy(readData, readData.length - recmac.length, recmac, 0, recmac.length);

				if (!Arrays.equals(recmac, genmac))
					throw new Error("Local Error: Invalid HMAC!  Potential Attack!");

				byte[] returnData = new byte[readData.length - recmac.length];
				System.arraycopy(readData, 0, returnData, 0, returnData.length);

				// Process the data
				// byte[] sendData = table.acceptDataFromLocal(readData);
				byte[] sendData = table.acceptDataFromLocal(returnData);

				mac.update(sendData);
				byte[] realmac = mac.doFinal();
				byte[] totalData = new byte[sendData.length + realmac.length];
				System.arraycopy(sendData, 0, totalData, 0, sendData.length);
				System.arraycopy(realmac, 0, totalData, sendData.length, realmac.length);

				// Encrypt the data for sending
				byte[] encryptedData = encryptCipher.doFinal(totalData);

				// Send data to output (length of data, the data)
				output.writeInt(encryptedData.length);
				output.write(encryptedData, 0, encryptedData.length);
				output.flush();

				// close the socket
				socket.close();
			} catch (SocketTimeoutException e) {

			} catch (BadPaddingException e) {

			} catch (IllegalBlockSizeException e) {

			} catch (UnknownHostException e) {

			} catch (IOException e) {

			}
		}

		if (inputSocket != null) {
			try {
				inputSocket.close();
			} catch (Exception e) {
				e.printStackTrace();
				throw new Error("Local server close failure...");
			}
		}
	}

	public void close() {
		doEnd = true;

		if (localServerThread != null) {
			try {
				localServerThread.join();
			} catch (Exception e) {
				e.printStackTrace();
				throw new Error("Local Server thread join issue...");
			}
		}

		// System.out.println("Done Closing Cloud Comm");
	}

	protected void finalize() throws Throwable {
		try {
			close();        // close open files
		} finally {
			super.finalize();
		}
	}

}