package de.bwlehrpool.bwlp_guac;

import java.io.Closeable;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.SecretKeySpec;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class VncConnection implements Closeable {

	private static final Logger LOGGER = LoggerFactory.getLogger(VncConnection.class);

	private final Socket sock;
	private final DataOutputStream out;
	private final DataInputStream in;

	public VncConnection(String host, int port) throws IOException {
		sock = new Socket();
		sock.connect(new InetSocketAddress(host, port), 1200);
		sock.setSoTimeout(1000);
		out = new DataOutputStream(sock.getOutputStream());
		in = new DataInputStream(sock.getInputStream());
	}

	public String handshake() throws IOException {
		byte[] buffer = new byte[12];
		in.readFully(buffer);
		out.write("RFB 003.008\n".getBytes());
		out.flush();
		return new String(buffer).substring(4, 11);
	}

	public boolean tryLogin(String passwd) throws IOException {
		if (passwd == null)
			return false; // Paswordless not supported yet (although simpler..)
		int numTypes = in.read();
		if (numTypes == 0) {
			LOGGER.info("VNC Server @ " + sock.getRemoteSocketAddress() + " does not support any auth methods");
			printError();
			return false;
		}
		boolean ok = false;
		for (int i = 0; i < numTypes; ++i) {
			if (in.read() == 2) {
				ok = true;
				break; // Found "VNC Authentication"
			}
		}
		if (!ok) {
			LOGGER.info("VNC Server @ " + sock.getRemoteSocketAddress() + " does not support VNC Authentication");
			return false;
		}
		out.write(2); // Pick passwd auth
		// Get challenge data
		byte[] challenge = new byte[16];
		int ret = in.read(challenge);
		if (ret != 16) {
			LOGGER.info("Didn't receive challenge from VNC server @ " + sock.getRemoteSocketAddress());
			return false;
		}
		// pad pw to 8 bytes
		byte[] pw_bytes = passwd.getBytes();
		pw_bytes = Arrays.copyOf(pw_bytes, 8);
		// Encrypt
		Cipher des;
		try {
			des = Cipher.getInstance("DES/ECB/NoPadding");
			des.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(reverseBits(pw_bytes), 0, pw_bytes.length, "DES"));
			out.write(des.doFinal(challenge));
		} catch (InvalidKeyException | NoSuchAlgorithmException | NoSuchPaddingException | IllegalBlockSizeException
				| BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		// check reply
		int securityReply = in.readInt();
		if (securityReply != 0) {
			LOGGER.info("Security reply = " + securityReply + " for VNC server @ " + sock.getRemoteSocketAddress());
			return false;
		}
		return true;
	}

	private void printError() throws IOException {
		int len = in.readInt();
		byte[] msg = new byte[len];
		in.readFully(msg);
		LOGGER.info(new String(msg, StandardCharsets.ISO_8859_1));
	}

	@Override
	public void close() throws IOException {
		try {
			in.close();
		} catch (Exception e) {
		}
		try {
			out.close();
		} catch (Exception e) {
		}
		try {
			sock.close();
		} catch (Exception e) {
		}
	}

	/*
	 * 
	 */

	private byte[] reverseBits(byte[] b) {
		byte[] result = new byte[b.length];
		for (int i = 0; i < b.length; i++) {
			result[i] = reverseBits(b[i]);
		}
		return result;
	}

	private byte reverseBits(byte input) {
		byte result = 0x00;
		for (int i = 0; i < 8; i++) {
			result |= ((byte) ((input & (0x01 << i)) >>> i) << 7 - i);
		}
		return result;
	}
	
}