package de.tudarmstadt.informatik.hostage.protocol;

import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.List;

import de.tudarmstadt.informatik.hostage.commons.HelperUtils;
import de.tudarmstadt.informatik.hostage.ssh.crypto.KeyMaterial;
import de.tudarmstadt.informatik.hostage.ssh.crypto.PEMDecoder;
import de.tudarmstadt.informatik.hostage.ssh.crypto.cipher.CBCMode;
import de.tudarmstadt.informatik.hostage.ssh.crypto.cipher.DESede;
import de.tudarmstadt.informatik.hostage.ssh.crypto.dh.DhExchange;
import de.tudarmstadt.informatik.hostage.ssh.crypto.digest.MAC;
import de.tudarmstadt.informatik.hostage.ssh.signature.DSAPrivateKey;
import de.tudarmstadt.informatik.hostage.ssh.signature.DSASHA1Verify;
import de.tudarmstadt.informatik.hostage.ssh.signature.DSASignature;
import de.tudarmstadt.informatik.hostage.ssh.util.TypesWriter;
import de.tudarmstadt.informatik.hostage.wrapper.Packet;

/**
 * SSH protocol.
 * @author Wulf Pfeiffer
 */
public class SSH implements Protocol {
	/**
	 * Represents the states of the protocol.
	 */
	private enum STATE {
		NONE,
		SERVER_VERSION,
		CLIENT_VERSION,
		KEX_INIT,
		USERAUTH,
		CONNECTION,
		CHANNEL,
		LOOP, // TODO for debugging
		CLOSED
	}
	
	/**
	 * Denotes in which state the protocol is right now.
	 */
	private STATE state = STATE.NONE;
		
	@Override
	public int getPort() {
		return 22;
	}

	@Override
	public TALK_FIRST whoTalksFirst() {
		return TALK_FIRST.SERVER;
	}

	@Override
	public List<Packet> processMessage(Packet packet) {
		List<Packet> response = new ArrayList<Packet>();
		byte[] request = null;
		if(packet != null) {
			request = packet.getMessage();
		}
		switch(state) {
		case NONE:			
			response.add(new Packet(serverVersion + serverType + "\r\n"));
			state = STATE.SERVER_VERSION;
			break;
		case SERVER_VERSION:
			extractType(request);
			extractPayload(request);
			response.add(kexInit());
			state = STATE.CLIENT_VERSION;
			break;
		case CLIENT_VERSION:
			extractPubKey(request);
			response.add(dhKexReply());
			response.add(newKeys());
			state = STATE.KEX_INIT;
			break;
		case KEX_INIT:
			response.add(serviceReply(request));
			state = STATE.USERAUTH;
			break;
		case USERAUTH:
			response.add(userauthReply(request));
			state = STATE.CONNECTION;
			break;
		case CONNECTION:
			response.add(channelOpenReply(request));
			state = STATE.CHANNEL;
			break;
		case CHANNEL:
			response.add(channelSuccessReply(request));
			response.add(data(request));
			state = STATE.LOOP;
			break;
		case LOOP:
			decryptPacket(request);
			response.add(data(request));
			state = STATE.LOOP;
			break;
		case CLOSED:
			break;
		default:
			state = STATE.CLOSED;
			break;
		}
		
		return response;
	}

	@Override
	public boolean isClosed() {
		return (state == STATE.CLOSED);
	}

	@Override
	public boolean isSecure() {
		return false;
	}

	@Override
	public Class<byte[]> getType() {
		return byte[].class;
	}

	@Override
	public String toString() {
		return "SSH";
	}

	/**
	 * Wraps the packets with packet length and padding.
	 * @param packet content that is wrapped.
	 * @return wrapped packet.
	 */
	private Packet wrapPacket(byte[] packet, boolean encrypt) {
		int packetLength = 5 + packet.length; 	//4 byte packet length, 1 byte padding length, payload length
		int paddingLengthCBS = cipherBlockSize - (packetLength % cipherBlockSize);
		int paddingLength8 = 8 - (packetLength % 8);
		int paddingLength = paddingLengthCBS > paddingLength8 ? paddingLengthCBS : paddingLength8;
		if(paddingLength < 4) paddingLength += cipherBlockSize;
		packetLength = packetLength + paddingLength - 4;					//add padding string length to packet length
		
		byte[] packetLen = ByteBuffer.allocate(4).putInt(packetLength).array();
		byte[] paddingLen = {(byte) paddingLength};
		byte[] paddingString = new byte[paddingLength];
		for(int i = 0; i < paddingLength; i++) {
			SecureRandom rndm = new SecureRandom();
			paddingString[i] = (byte) rndm.nextInt(255);
		}
		byte[] response = HelperUtils.concat(packetLen, paddingLen, packet, paddingString);
		if(encrypt) {
			byte[] mac = new byte[20];
			macEnc.initMac(send_packet_number);
			macEnc.update(response, 0, response.length);
			macEnc.getMac(mac, 0);		
			
			byte[] responseEnc = new byte[response.length];
			for(int i = 0; i < response.length; i+=8) { 
				cbcEnc.transformBlock(response, i, responseEnc, i);
			}
			response = HelperUtils.concat(responseEnc,mac);
		}
		send_packet_number++;
		
		return new Packet(response);
	}
	
	private byte[] decryptPacket(byte[] request) {
		byte[] message = new byte[request.length - ((request.length % 8 == 0) ? 0 : 20)];
		for(int i = 0; i < message.length; i += 8) { // -12 wegen MAC
			cbcDec.transformBlock(request, i, message, i);
		}
		System.out.println("Encrypted packet:");
		System.out.println(HelperUtils.bytesToHexString(request));
		System.out.println("\nDecrypted packet:");
		System.out.println(HelperUtils.bytesToHexString(message));
		System.out.println("\nDecrypted string packet:");
		System.out.println(HelperUtils.byteToStr(message));
		System.out.println();
		return message;
	}
	
	/**
	 * Builds the Kex Init packet that contains all the allowed algorithms by the server.
	 * @return Kex Init packet.
	 */
	private Packet kexInit() {
		byte[] msgCode = {0x14};
		byte[] cookie = HelperUtils.randomBytes(16);
		byte[] kexLength = ByteBuffer.allocate(4).putInt(kex_alg.getBytes().length).array();
		byte[] serverLength = ByteBuffer.allocate(4).putInt(server_alg.getBytes().length).array();
		byte[] encrypt_c_Length = ByteBuffer.allocate(4).putInt(encrypt_alg_c.getBytes().length).array();
		byte[] encrypt_s_Length = ByteBuffer.allocate(4).putInt(encrypt_alg_s.getBytes().length).array();
		byte[] mac_c_Length = ByteBuffer.allocate(4).putInt(mac_alg_c.getBytes().length).array();
		byte[] mac_s_Length = ByteBuffer.allocate(4).putInt(mac_alg_s.getBytes().length).array();
		byte[] comp_c_Length = ByteBuffer.allocate(4).putInt(comp_alg_c.getBytes().length).array();
		byte[] comp_s_Length = ByteBuffer.allocate(4).putInt(comp_alg_s.getBytes().length).array();
		byte[] language_c_s = {0x00, 0x00, 0x00, 0x00};
		byte[] language_s_c = {0x00, 0x00, 0x00, 0x00};
		byte[] kexFirsPckt = {0x00};
		byte[] reserved = {0x00, 0x00, 0x00, 0x00};	
		
		byte[] response = HelperUtils.concat(msgCode, cookie, kexLength, kex_alg.getBytes(), serverLength, server_alg.getBytes(),
							encrypt_c_Length, encrypt_alg_c.getBytes(), encrypt_s_Length, encrypt_alg_s.getBytes(), mac_c_Length, mac_alg_c.getBytes(),
							mac_s_Length, mac_alg_s.getBytes(), comp_c_Length, comp_alg_c.getBytes(), comp_s_Length, comp_alg_s.getBytes(),
							language_c_s, language_s_c, kexFirsPckt, reserved);
		I_S = response;
		return wrapPacket(response, false);
	}
	
	/**
	 * Builds the Diffie-Hellman Kex Reply, containing the host key,f and the signature.
	 * @return Diffie-Hellman Kex Reply packet.
	 */
	private Packet dhKexReply() {
		//TODO make it small		
		try {
			SecureRandom rnd = new SecureRandom();
			DhExchange dhx = new DhExchange();
			dhx.serverInit(1, rnd);
			dhx.setE(new BigInteger(e));
			f = dhx.getF();
			DSAPrivateKey dsa = (DSAPrivateKey) PEMDecoder.decode(dsa_pem, null);
			K_S = DSASHA1Verify.encodeSSHDSAPublicKey(dsa.getPublicKey());
			h = dhx.calculateH(V_C, V_S, I_C, I_S, K_S);
			k = dhx.getK();
			DSASignature ds = DSASHA1Verify.generateSignature(h, dsa, rnd);
			sig = DSASHA1Verify.encodeSSHDSASignature(ds);
			TypesWriter tw = new TypesWriter();
			tw.writeByte(31);
			tw.writeString(K_S, 0, K_S.length);
			tw.writeMPInt(f);
			tw.writeString(sig, 0, sig.length);
			byte[] response = tw.getBytes();
			
			//init for decryption and encryption
			KeyMaterial km = KeyMaterial.create("SHA1", h, k, h, 24, 8, 20, 24, 8, 20); // alg, h, k, keylength, blocklength, maclength, keylength, blocklength, maclength
			enc = new DESede();
			dec = new DESede();
			enc.init(true, km.enc_key_server_to_client);
			dec.init(false, km.enc_key_client_to_server);
			cbcEnc = new CBCMode(enc, km.initial_iv_server_to_client, true);
			cbcDec = new CBCMode(dec, km.initial_iv_client_to_server, false);
			macEnc = new MAC("hmac-sha1", km.integrity_key_server_to_client);
			macDec = new MAC("hmac-sha1", km.integrity_key_client_to_server);
			
			return wrapPacket(response, false);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return wrapPacket(null, false);
	}

	/**
	 * New Keys response.
	 * @return New Keys response.
	 */
	private Packet newKeys() {
		byte[] msgCode = {0x15};
		return wrapPacket(msgCode, false);
	}
	
	// TODO
	private Packet serviceReply(byte[] request) {
		byte[] messageBuf;
		if(request[5] == 0x15) {	//if newkeys request is included in the same packet
			messageBuf = new byte[request.length - 16];
			System.arraycopy(request, 16, messageBuf, 0, request.length-16);
		} else {
			messageBuf = request;
		}
		byte[] message = decryptPacket(messageBuf);
		
		if (message[5] != 0x05 && message[9] != 0x0c) {
			// TODO disconnect because its not servicerequest ssh-userauth
		}
		byte[] msgcode = {0x06};
		byte[] serviceLength;
		byte[] service = "ssh-userauth".getBytes();
		serviceLength = ByteBuffer.allocate(4).putInt(service.length).array();
		
		byte[] response = HelperUtils.concat(msgcode, serviceLength, service);
		return wrapPacket(response, true);		
	}
	
	// TODO
	private Packet userauthReply(byte[] request) {
		byte[] message = decryptPacket(request);
		if (message[5] != 0x32) {
			// TODO disconnect because its not servicerequest ssh-connect
		}
		byte[] msgcode = {0x34};
		byte[] response = msgcode;
		
		return wrapPacket(response, true);		
	}
	
	// TODO
	private Packet channelOpenReply(byte[] request) {
		byte[] message = decryptPacket(request);
		// TODO if contains "session" ok else disc
		byte[] msgcode = {0x5b};
		recipientChannel = new byte[4];
		byte[] senderChannel = new byte[4];
		byte[] initialWindowSize = new byte[4];
		byte[] maximumPacketSize = new byte[4];
		System.arraycopy(message, 17, recipientChannel, 0, 4); //TODO dynamic position
		System.arraycopy(message, 17, senderChannel, 0, 4);
		System.arraycopy(message, 21, initialWindowSize, 0, 4);
		System.arraycopy(message, 25, maximumPacketSize, 0, 4);
		byte[] response = HelperUtils.concat(msgcode, recipientChannel, senderChannel, initialWindowSize, maximumPacketSize);
		System.out.println("CHANNEL REPLY:");
		System.out.println(HelperUtils.bytesToHexString(response));
		return wrapPacket(response, true);
	}
	
	byte[] recipientChannel; //TODO
	private Packet channelSuccessReply(byte[] request) {
		decryptPacket(request);
		byte[] msgcode = {0x63}; //TODO check msg type
		byte[] response = HelperUtils.concat(msgcode, recipientChannel);
		System.out.println(HelperUtils.bytesToHexString(response));
		return wrapPacket(response, true);
	}
	
	private Packet data(byte[] request) {
		byte[] msgcode = {0x5E};
		byte[] msglength = new byte[4];
		byte[] msg = "lolol".getBytes();
		msglength = ByteBuffer.allocate(4).putInt(msg.length).array();
		byte[] response = HelperUtils.concat(msgcode, recipientChannel, msglength, msg);
		return wrapPacket(response, true);
	}
	
	/**
	 * Extracts the type of the client
	 * @param request containing the clients type
	 */
	private void extractType(byte[] request) {
		int length = 0;
		for(int i = 0; i < request.length; i++, length++) { 	
			if(request[i] == 0x0d) break;			//find the end of the type: '\r'
		}
		V_C = new byte[length];
		System.arraycopy(request, 0, V_C, 0, length);
	}
	
	/**
	 * Extracts the payload of a packet and writes it in I_C.
	 * @param request packet of which the payload is extracted.
	 */
	private void extractPayload(byte[] request) {
		int pos = 0;
		if(request[5] != 0x14) {
			pos = 1;
			for(int i = 0; i < request.length; i++, pos++) {
				if(request[i] == 0xa) break;
			}
		}
		int packetLength = byteToInt(new byte[]{request[pos], request[1+pos], request[2+pos], request[3+pos]});
		int paddingLength = byteToInt(new byte[]{request[4+pos]});
		byte[] payload = new byte[packetLength - paddingLength - 1];
		for(int i = 5; i < packetLength - paddingLength - 1; i++) {
			payload[i-5] = request[i+pos];
		}
		I_C = payload;
	}
	
	/**
	 * Extracts the public key from the DH Kex Request
	 * @param request containing the clients public key
	 */
	private void extractPubKey(byte[] request) {
		e = new byte[byteToInt(new byte[] {request[6], request[7], request[8], request[9]})];
		for(int i = 0; i < e.length; i++) {
			e[i] = request[i+10];
		}
	}
	
	/**
	 * Converts a byte[] to int
	 * @param bytes that are converted
	 * @return converted byte[] as int
	 */
	private static int byteToInt(byte[] bytes) {
		  int ret = 0;
		  for (int i=0; i < bytes.length; i++) {
		    ret <<= 8;
		    ret |= bytes[i] & 0xFF;
		  }
		  return ret;
	}
	
	private String serverVersion = ProtocolSettings.getSshVersion();
	private String serverType = ProtocolSettings.getSshType();
	private int send_packet_number = 0;
	
	//SSH Parameters for Kex etc.
    private byte[] V_S = (serverVersion + serverType).getBytes();
    private byte[] V_C;
    private byte[] I_S;
    private byte[] I_C;
    private byte[] e;
    private BigInteger f;
    private byte[] h;
    private BigInteger k;
    private byte[] K_S;
    private byte[] sig;
			
    //allowed algorithms for kexinit
	private String kex_alg = "diffie-hellman-group1-sha1";
	private String server_alg = "ssh-dss";
	private String encrypt_alg_c = "3des-cbc";
	private String encrypt_alg_s = "3des-cbc";
	private String mac_alg_c = "hmac-sha1";
	private String mac_alg_s = "hmac-sha1";
	private String comp_alg_c = "none";
	private String comp_alg_s = "none";
	
	private int cipherBlockSize = 16;
	
	DESede enc;
	DESede dec;
	CBCMode cbcEnc;
	CBCMode cbcDec;
	MAC macEnc;
	MAC macDec;
	
	//dsa private key
	private final char[] dsa_pem = ("-----BEGIN DSA PRIVATE KEY-----\n" +
			"MIIBugIBAAKBgQCDZ9R2vfCPwjv5vKF1igIv9drrZ7G0dhMkGT9AZTjgI34Qm4w0\n" +
			"0iWeCqO7SmqiaMIjbRIm91MeDed4ObAq4sAkqRE/2P4mTbzFx5KhEczRRiDoqQBX\n" +
			"xYa0yWKJpeZ94SGM6DEPuBTxKo0T4uMjbq2FzHL2FXT1/WoNCmRU6gFSiwIVAMK4\n" +
			"Epz3JiwDUbkSpLOjIqtEhJmVAoGAL6zlXRI4Q8iwvSDh0vDf1j9a5Aaaq+93LTjK\n" +
			"SwL4nvUWBl2Aa0vqu05ZS5rOD1I+/naLMg0fNgFJRhA03sl+12MI3a2HXJWXRSdj\n" +
			"m1Vq9cUXqiYrX6+iGfEaA/y9UO4ZPF6if6eLypXB8VuqjtjDCiMMsM6+qQki7L71\n" +
			"yN4M75ICgYAcFXUhN2zRug3JvwmGxW8gMgHquSiBnbx1582KGh2B/ukE/kOrbKYD\n" +
			"HUkBzolcm4x1Odq5apowlriFxY6zMQP615plIK4x9NaU6dvc/HoTkjzT5EYSMN39\n" +
			"eAGufJ0jrtIpKL4lP8o8yrAHfmbR7bjecWc0viTH0+OWlyVsex/bZAIUEKn310Li\n" +
			"v62Zs4hlDvhwvx8MQ+A=\n" +
			"-----END DSA PRIVATE KEY-----").toCharArray();

}