package org.openslx.taskmanager.tasks;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.UnknownHostException;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.X509Certificate;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;

import org.openslx.taskmanager.api.AbstractTask;

import com.google.gson.annotations.Expose;

public class SSLRelayTask extends AbstractTask {

	@Expose
	private String clientAIp;
	@Expose
	private int clientAPort;
	
	@Expose
	private String clientBIp;
	@Expose
	private int clientBPort;
	
	@Expose
	private boolean auth;

	private SSLSocket sockA;
	private SSLSocket sockB;
	
	private Relay aToB;
	private Relay bToA;
	
	private boolean enabled;
	
	private Output status;
	
	@Override
	protected boolean initTask() {
	
		this.setStatusObject(status);	
		return true;
	}
	
	@Override
	protected boolean execute() {
		SSLSocketFactory ssf = initSSLSocketFactory();
		if (ssf==null) {
			status.error = "Could not initialize SSLSocketFactory";
			return false;
		}
		try {
			initSockets(ssf);	
			initRelays();
		} catch (UnknownHostException uhx) {
			status.error = "One of the hosts to relay to/from is unknown: " + uhx.getMessage();
			return close();
		} catch (IOException iox) {
			status.error = iox.getMessage();
			return close();
		}
		
		Thread aToBThread = initThread('a');
		Thread bToAThread = initThread('b');
		
		if (aToBThread == null || bToAThread == null) {
			status.error = "Could not initialize Threads.";
			return close();
		}

		while(enabled) {
			aToBThread.start();
			bToAThread.start();
		}
		return close();
	}
	
	/**
	 * During testing phase there is the option to disable SSL/TLS authentication.
	 * THIS SHOULD BE REMOVED FOR ROLLOUT
	 * If auth == true, set system properties for key- and truststore and return default.
	 * Else, return SocketFactory from all-trusting SSLContext (returned by trustAll().
	 * @return
	 */
	private SSLSocketFactory initSSLSocketFactory () {
		if (auth) {
			System.setProperty("javax.net.ssl.keyStore", "keystore.jks");
			System.setProperty("javax.net.ssl.trustStore", "cacerts.jks");
			
			return (SSLSocketFactory) SSLSocketFactory.getDefault();
		}
		else {
			try {
				return trustAll().getSocketFactory();
			} catch (NoSuchAlgorithmException nax) {
				status.error = nax.getMessage();
				return null;
			} catch (KeyManagementException kmx) {
				status.error = kmx.getMessage();
				return null;
			}
		}
	}
	
	/**
	 * Returns one Thread to run the relay. Is there a nicer way to do this?
	 * Scoping posed some difficulties
	 * @param source
	 * @return
	 */
	private Thread initThread (char source) {
		if (source == 'a') {
			return new Thread() {
				public void run() {
					try {
						aToB.relay();
					} catch (IOException iox) {
						status.error = iox.getMessage();
						return;
					} catch (InterruptedException ix) {
						status.error = ix.getMessage();
						return;
					}
				};
			};
		}
		else if (source == 'b') {
			return new Thread() {
				public void run() {
					try {
						bToA.relay();
					} catch (IOException iox) {
						status.error = iox.getMessage();
						return;
					} catch (InterruptedException ix) {
						status.error = ix.getMessage();
						return;
					}
				};
			};
		}
		else
			return null;
	}
	/**
	 * Creates sockA and sockB
	 * @param ssf
	 * @throws IOException
	 * @throws UnknownHostException
	 */
	private void initSockets (SSLSocketFactory ssf) throws IOException, UnknownHostException {
		sockA = (SSLSocket) ssf.createSocket(clientAIp, clientAPort);
		System.out.println("connected to " + clientAIp + " on port " + Integer.toString(clientAPort));
		sockB = (SSLSocket) ssf.createSocket(clientBIp, clientBPort);
		System.out.println("connected to " + clientBIp + " on port " + Integer.toString(clientBPort));
	}
	/**
	 * Creates relays aToB and bToA
	 * @throws IOException
	 */
	private void initRelays () throws IOException {
		aToB = new Relay(sockA, sockB);
		System.out.println("relay created from " + clientAIp + " to " + clientBIp);
		bToA = new Relay(sockB, sockA);
		System.out.println("relay created from " + clientBIp + " to " + clientAIp);
	}
	
	/**
	 * Create all-trusting TrustManager for no-auth mode and return SSLContext.
	 */
	private SSLContext trustAll () throws NoSuchAlgorithmException, KeyManagementException {
		TrustManager[] trustAllMan = new TrustManager[] {new X509TrustManager() {
				public java.security.cert.X509Certificate[] getAcceptedIssuers() {
					return null;
				}
				public void checkClientTrusted(X509Certificate[] certs, String authType) {}
				public void checkServerTrusted(X509Certificate[] certs, String authType) {}
			}
		};
		SSLContext ctx = SSLContext.getInstance("SSL");
		ctx.init(null, trustAllMan, new java.security.SecureRandom());
		return ctx;
	}
	
	private boolean close() {
		try {
			if (aToB != null)
				aToB.close();
			if (bToA != null)
				bToA.close();
			sockA.close();
			sockB.close();
		} catch (IOException iox) {
			status.error = iox.getMessage();
			return false;
		}
		return true;
	}
	
	/**
	 * Do the actual relaying in one direction
	 */
	private class Relay {
		private InputStream in;
		private OutputStream out;
		
		private byte[] buffer = new byte[16384];
		
		public Relay (SSLSocket sIn, SSLSocket sOut) throws IOException {
			in = sIn.getInputStream();
			out = sOut.getOutputStream();
		}
		
		public void relay() throws IOException, InterruptedException {
			int readBytes = in.read(buffer);
			
			out.write(buffer, 0, readBytes);
		}
		
		public void close() throws IOException {
			in.close();
			out.close();
		}
	}
	
	/**
	 * Output - contains additional status data of this task
	 */
	@SuppressWarnings( "unused" )
	private static class Output
	{
		protected String error = null;
	}
}