package de.tudarmstadt.informatik.hostage.protocol;

import java.io.IOException;
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.TypesReader;
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,
		NEW_KEYS,
		USERAUTH,
		CONNECTION,
		CHANNEL,
		TERMINAL_CMD,
		TERMINAL_ENTER,
		CLOSED
	}
	
	/**
	 * Denotes in which state the protocol is right now.
	 */
	private STATE state = STATE.NONE;
	private boolean useEncryption = false;
		
	@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();
			if(useEncryption) {
				request = decryptBytes(request);
			}
		}
		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());
			state = STATE.KEX_INIT;
			break;
		case KEX_INIT:
			response.add(newKeys());
			useEncryption = true;
			state = STATE.NEW_KEYS;
			break;
		case NEW_KEYS:
			response.add(serviceReply(request));
			state = STATE.USERAUTH;
			break;
		case USERAUTH:
			response.add(connectionReply(request));
			state = STATE.CONNECTION;
			break;
		case CONNECTION:
			response.add(channelOpenReply(request));
			state = STATE.CHANNEL;
			break;
		case CHANNEL:
			response.add(channelSuccessReply(request));
			response.add(terminalPrefix());
			state = STATE.TERMINAL_CMD;
			break;
		case TERMINAL_CMD:
			response.add(terminalReply(request));
			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) {
		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 = HelperUtils.randomBytes(paddingLength);
		byte[] response = HelperUtils.concat(packetLen, paddingLen, packet, paddingString);
		if(useEncryption) {
			byte[] mac = getMac(response);
			byte[] responseEnc = encryptBytes(response);
			response = HelperUtils.concat(responseEnc,mac);
		}
		packetNumber++;
		
		return new Packet(response);
	}
	
	/**
	 * Encrypts a request with triple DES.
	 * @param request that is encrypted.
	 * @return encrypted request.
	 */
	private byte[] encryptBytes(byte[] bytes) {
		byte[] responseEnc = new byte[bytes.length];
		for(int i = 0; i < bytes.length; i+=8) { 
			cbcEnc.transformBlock(bytes, i, responseEnc, i);
		}
		return responseEnc;
	}
	
	/**
	 * Decrypts a request with triple DES.
	 * @param request that is decrypted.
	 * @return decrypted request.
	 */
	private byte[] decryptBytes(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);
		}
		return message;
	}
	
	/**
	 * Creates the SHA1 Mac with the given bytes.
	 * @param bytes that are used for the Mac.
	 * @return Mac.
	 */
	private byte[] getMac(byte[] bytes) {
		byte[] mac = new byte[20];
		macEnc.initMac(packetNumber);
		macEnc.update(bytes, 0, bytes.length);
		macEnc.getMac(mac, 0);
		return mac;
	}
	
	/**
	 * Builds the Kex Init packet that contains all the allowed algorithms by the server.
	 * @return Kex Init packet.
	 */
	private Packet kexInit() {
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x14);
		tw.writeBytes(HelperUtils.randomBytes(16)); //cookie
		tw.writeString(kex_alg);
		tw.writeString(server_alg);
		tw.writeString(encrypt_alg_c);
		tw.writeString(encrypt_alg_s);
		tw.writeString(mac_alg_c);
		tw.writeString(mac_alg_s);
		tw.writeString(comp_alg_c);
		tw.writeString(comp_alg_s);
		tw.writeBytes(new byte[]{0x00, 0x00, 0x00, 0x00}); //language client to server
		tw.writeBytes(new byte[]{0x00, 0x00, 0x00, 0x00}); //language server to client
		tw.writeByte(0x00); //no guess from server
		tw.writeBytes(new byte[]{0x00, 0x00, 0x00, 0x00}); //reserved
		byte[] response = tw.getBytes();
		I_S = response;
		
		return wrapPacket(response);
	}
	
	/**
	 * Builds the Diffie-Hellman Kex Reply, containing the host key,f and the signature.
	 * @return Diffie-Hellman Kex Reply packet.
	 */
	private Packet dhKexReply() {
		byte[] response = null;
		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);
			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);
		} catch (Exception e) {
			e.printStackTrace();
		}
		
		return wrapPacket(response);
	}

	/**
	 * New Keys response.
	 * @return New Keys response.
	 */
	private Packet newKeys() {
		byte[] msgCode = {0x15};
		return wrapPacket(msgCode);
	}
	
	/**
	 * Service ssh-userauth reply.
	 * @param request from the client.
	 * @return Service reply.
	 */
	private Packet serviceReply(byte[] request) {
		byte[] message;
		if(request[5] == 0x15) {	//if newkeys request is included in the same packet
			message = new byte[request.length - 16]; //remove it
			System.arraycopy(request, 16, message, 0, request.length-16);
		} else {
			message = request;
		}
		if (message[5] != 0x05 && !(HelperUtils.byteToStr(message).contains("ssh-userauth"))) {
			return disconnectReply(7); // disconnect because its not servicerequest ssh-userauth
		}
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x06);
		tw.writeString("ssh-userauth");
		return wrapPacket(tw.getBytes());		
	}
	
	/**
	 * Userauth ssh-connection reply.
	 * @param request from the client.
	 * @return ssh-connection reply.
	 */
	private Packet connectionReply(byte[] request) {
		if (request[5] != 0x32 && !(HelperUtils.byteToStr(request).contains("ssh-connection"))) {
			return disconnectReply(14);// disconnect because its not servicerequest ssh-connect
		}
		try {
			TypesReader tr = new TypesReader(request, 6);
			userName = tr.readString();
			terminalPrefix = "[" + userName + "@" + serverName + "]$";
		} catch (IOException e) {
			e.printStackTrace();
		}
		byte[] msgcode = {0x34};		
		return wrapPacket(msgcode);		
	}
	
	/**
	 * Channel Open Reply.
	 * @param request from client.
	 * @return Channel Open Reply.
	 */
	private Packet channelOpenReply(byte[] request) {
		if (!(HelperUtils.byteToStr(request).contains("session"))) {
			return disconnectReply(2); // if contains "session" ok else disc
		}
		TypesReader tr = new TypesReader(request, 6);
		TypesWriter tw = new TypesWriter();
		try {
			tr.readString();
			recipientChannel = tr.readUINT32();
			int senderChannel = recipientChannel;
			int initialWindowSize = tr.readUINT32();
			int maximumPacketSize = tr.readUINT32();
			
			tw.writeByte(0x5b); //msgcode
			tw.writeUINT32(recipientChannel);
			tw.writeUINT32(senderChannel);
			tw.writeUINT32(initialWindowSize);
			tw.writeUINT32(maximumPacketSize);
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		return wrapPacket(tw.getBytes());
	}
	
	/**
	 * Channel Success Reply.
	 * @param request from client.
	 * @return Channel Success Reply.
	 */
	private Packet channelSuccessReply(byte[] request) {
		if(!(HelperUtils.byteToStr(request)).contains("pty-req")) {
			return disconnectReply(2);
		}
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x63); //msgcode
		tw.writeUINT32(recipientChannel);
		
		return wrapPacket(tw.getBytes());
	}
	
	/**
	 * Returns the terminal prefix for the client.
	 * @return terminal prefix.
	 */
	private Packet terminalPrefix() {
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x5e);
		tw.writeUINT32(recipientChannel);
		tw.writeString(terminalPrefix);
		return wrapPacket(tw.getBytes());
	}
	
	/**
	 * Computes the reply for the client input.
	 * @param request client input.
	 * @return input reply.
	 */
	private Packet terminalReply(byte[] request) {
		//TODO
		TypesReader tr = new TypesReader(request, 6);
		String msg = "";
		System.out.println(HelperUtils.bytesToHexString(request));
		try {
			tr.readUINT32();
			msg = tr.readString();
			System.out.println(msg);
			if(msg.contains("\r")) {
				System.out.println("hi");
				if(cmd.toString().contains("exit")) {
					state = STATE.CLOSED; // ugly style
					return disconnectReply(2);
				}
				msg = "\r\nbash: " + cmd + " :command not found\r\n"	+ terminalPrefix;
				cmd = new StringBuffer();
			} else if(msg.contains(new String(new char[]{'\u007F'})) && cmd.length() > 0) {
				cmd = cmd.delete(cmd.length()-1, cmd.length());
			} else {
				cmd.append(msg);
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x5e); //msgcode
		tw.writeUINT32(recipientChannel);
		tw.writeString(msg);
		return wrapPacket(tw.getBytes());
	}
	
	/**
	 * Disconnect Reply using the given number as reason code.
	 * @param reasonCode for disconnect reply. Must be between 1 and 15, default is 2.
	 * @return Disconnect Reply.
	 */
	private Packet disconnectReply(int reasonCode) {
		TypesWriter tw = new TypesWriter();
		tw.writeByte(0x01);
		switch (reasonCode) {
		case 1: 
			tw.writeUINT32(1);
			tw.writeString("SSH_DISCONNECT_HOST_NOT_ALLOWED_TO_CONNECT");
			break;
		case 7: 
			tw.writeUINT32(7);
			tw.writeString("SSH_DISCONNECT_SERVICE_NOT_AVAILABLE");
			break;
		case 14: 
			tw.writeUINT32(14);
			tw.writeString("SSH_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE");
			break;
		default: 
			tw.writeUINT32(2);
			tw.writeString("SSH_DISCONNECT_PROTOCOL_ERROR");
			break;
		}
		return wrapPacket(tw.getBytes());
	}
	
	/**
	 * 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;
	}
	
	//server infos
	private String serverVersion = ProtocolSettings.getSshVersion();
	private String serverType = ProtocolSettings.getSshType();
	private String serverName = ProtocolSettings.getSshName();
	private int packetNumber = 0;
	int recipientChannel;
	String userName;
	String terminalPrefix;
	StringBuffer cmd = new StringBuffer();
	
	//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;
	
	//for en- and decryption
	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();

}