import java.io.*;
import java.net.*;
import java.math.BigInteger;

public class Client {
	String nickname; // Brugerens navn på chat-serveren.
	String server; // Adressen på serveren.
	CListeningThread ServerListener;
	BigInteger AESkey;
	RSA rsa;
	AES aes;
	
	Socket clientSocket;

	BufferedReader inFromUser;
	BufferedReader inFromServer;
	DataOutputStream outToServer;	
	
	public static void main(String args[]) {
		new Client();
	}
	
	public Client() {
		inFromUser = new BufferedReader( new InputStreamReader(System.in));				
				
		Setup(); // Bed brugeren om navn og chat-server.
		// opret forbindelser til serveren og lav RSA-nøgler.
		try {
			clientSocket = new Socket(server, 6789);
			ServerListener = new CListeningThread(clientSocket);
			outToServer = new DataOutputStream(clientSocket.getOutputStream());
			inFromServer = new BufferedReader ( new InputStreamReader(clientSocket.getInputStream()));
			
			rsa = new RSA();
			aes = new AES();
			
			if (!rsa.Load("keys.dat")) {
				System.out.println("Generating RSA-system. Please wait...");
				rsa.GenerateKeys(true);		
				rsa.Save("keys.dat");
			}	
			
			ExchangeAESkey();
			
			// Start den tråd der venter på input fra serveren.
			ServerListener.start();
		} catch (UnknownHostException e) {
			System.out.println("Unknown host: " + e.getMessage());
			return;
		} catch (IOException e) {
			System.out.println("Network error: " + e.getMessage());
			return;
		}
		
		System.out.println("Kontakt til server etableret.");
		
		// Etabler forbindelse til serveren.		
		try {								
			WriteToServer(nickname);
			String userline = "";			
			System.out.print(">");			
			// Krypter hvad brugeren skriver og send det til serveren.
			do {			
				userline = ReadFromUser();			
				outToServer.writeBytes(Encrypt(userline) + '\n');
			} while (!userline.equals(""));
		} catch (IOException e) {
			System.out.println("Error writing to server: " + e.getMessage());
		}
		
		Disconnect();		
	}
	
	// Krypter str med AES.
	public String Encrypt(String str) {
		int i = 0;						
		StringBuffer res = new StringBuffer();
		
		// Krypter strengen 16 karakterer ad gangen. 
		while (i + 16 < str.length() ) {
			aes.state = str.substring(i, i + 16).getBytes();
			aes.encrypt();
			
			for (int j = 0 ; j < 16 ; j++) {
				res.append( AES.Byte2Hex(aes.state[j]));
			}
			i += 16;
		}
		
		// Krypter resten af strengen. Indsæt 0 hvor der ikke er en karakter.		
		byte[] byte_array = new byte[16];
		byte[] tmp;
		tmp = str.substring(i, str.length() ).getBytes();
		for (int j = 0 ; j < tmp.length ; j++) {
			byte_array[j] = tmp[j];
		}
		
		aes.state = byte_array;
		aes.encrypt();
		
		for (int j = 0 ; j < 16 ; j++) {
			res.append( AES.Byte2Hex(aes.state[j]));
		}		
		
		return res.toString();
	}		
	
	// Indledende handshaking med serveren. Her sendes den offentlige nøgle,
	// og AES-nøglen modtages.
	private void ExchangeAESkey() {
		StringQueue incoming = new StringQueue(5);
		StringBuffer res = new StringBuffer();
		WriteToServer(rsa.GetPublicKey().toString());
		WriteToServer(rsa.GetModulos().toString());
		String fromServer;
		while(true) {
			fromServer = ReadFromServer();
			if (fromServer.equals("&END"))
				break;
			incoming.Insert(fromServer);			
		}
		while (!incoming.isEmpty())
			res.append(rsa.DecryptPrivate(new BigInteger(incoming.Retrieve())));
		
		ServerListener.LoadKey(res.toString().trim());
		LoadKey(res.toString().trim());
	}
	
	// Konverter fra AES-nøglen som streng til et byte-array.
	public void LoadKey(String key) {
		AESkey = new BigInteger(key);
		// Big-endian order. MSB af tallet ligger i byte_key[0]. 
		byte[] byte_key = AESkey.toByteArray();				
		int[] int_key = new int[16];					
		
		// Math.min(..) skyldes at antallet af bytes i AESkey kan være fra
		// 17 og nedefter. (pga. sign bit).
		for (int i = 0 ; i < Math.min(byte_key.length,16) ; i++) {
			int_key[i] = byte_key[i];
		}

		aes.setCipherkey(int_key);
	}
	
	private String ReadFromServer() {
		try {
			return inFromServer.readLine();
		} catch (IOException e) {
			return null;
		}
	}
	
	private void WriteToServer(String str) {
		try {
			outToServer.writeBytes(str + '\n');
		} catch(IOException e) {
			return;
		}
	}
	
	// Lukker forbindelsen til serveren. 
	private void Disconnect() {
		try {
			ServerListener.running = false;
			clientSocket.close();
			outToServer.close();
		} catch (IOException e) {}
	}
	
	private void Setup() {
		System.out.print("Server-address: ");
		server = ReadFromUser();
		System.out.print("Username: ");
		nickname = ReadFromUser();
	}
	
	// Wrapper funktion. Bruges for ikke hele tiden at skulle catch'e IOException.
	private String ReadFromUser() {
		try {
			return inFromUser.readLine();
		} catch (IOException e) {
			System.out.println("Communication error! User not found! :)");
			return null;
		}		
	}
	
}

class CListeningThread extends Thread {	
	BufferedReader inFromServer;	
	public boolean running;
	InputScanner scanner;	
	BigInteger AESkey;
	AES aes;
	RSA rsa;

	public CListeningThread(Socket clientSocket) throws IOException {
		inFromServer = new BufferedReader ( new InputStreamReader(clientSocket.getInputStream()));
		running = true;				
		
		scanner = new InputScanner();
		aes = new AES();
	}
	
	// Dekrypterer str.	
	private String Decrypt(String str) {
		int i = 0;
		StringBuffer res = new StringBuffer();
		byte[] cipher_text = new byte[16];
		// Vi ved at strengens længde altid er et multiplum af 32.
		while (i != str.length()) {
			for (int j = 0 ; j < 16 ; j++) {
				cipher_text[j] = AES.ToByte(Integer.parseInt(str.substring(i,i+2), 16));
				i += 2;
			}
			aes.state = cipher_text;
			aes.decrypt();
			for (int j = 0 ; j < 16 ; j++)
				res.append( (char) aes.state[j]);
		}
		return res.toString();
	}
	
	public void LoadKey(String key) {
		AESkey = new BigInteger(key);
		// Big-endian order. MSB af tallet ligger i byte_key[0]. 
		byte[] byte_key = AESkey.toByteArray();				
		int[] int_key = new int[16];					

		for (int i = 0 ; i < Math.min(byte_key.length,16) ; i++) {
			int_key[i] = byte_key[i];
		}

		aes.setCipherkey(int_key);
	}
	
	// Håndter tekst sendt fra serveren.
	public void run() {
		String fromServer = null;
		try {			
			while(running) {
				fromServer = inFromServer.readLine();
				
				scanner.Scan(fromServer);
					
				switch (scanner.type) {
					case InputScanner.TEXT : 						
						System.out.println(Decrypt(fromServer));
					break;
				} // End of switch.								
			} 
			
			inFromServer.close();
		} catch (IOException e) {
			System.out.println("Network error: " + e.getMessage());
			return;		
		}		
	}	
}

// Benyttes til at genkende inputtet fra serveren. Benyttes langt mindre
// end oprindeligt tilsigtet. Gode muligheder for udvidelser her. :)
class InputScanner {
	int index; 
	public int type;
	String input; 
	String word;
	
	public static final int TEXT = 2;
	
	
	public void Scan(String str) {
		index = 1;		
		input = str;
		Next();
	}
	
	public int Next() {
		if (input.charAt(0) != '#') {
			type = TEXT;
			return TEXT;
		}
		
		word = "";
		
		SkipWhiteSpace();
		
		// Få fat i næste ord.
		while (Character.isLetterOrDigit(input.charAt(index))) {
			word = word + input.charAt(index);
			index++;			
			if (index == input.length()) {
				break;
			}
		}	
					
		return TEXT;
	}
	
	private void SkipWhiteSpace() {
		while (Character.isWhitespace(input.charAt(index))) {
			index++;
			if (index == input.length()) {
				break;
			}			
		}
	}
}

// En kø af tekst-strenge.
class StringQueue {
	String[] array;
	int top; // Peger på det første ubrugte element.
	int bottom; // Peger på den forreste streng i køen.
	int size; 
	
	public StringQueue(int s) {
		array = new String[s];
		top = bottom = 0;
	}		
	
	public synchronized void Insert(String str) {
		array[top] = str;
		top++;
		size++;
		
		// Check for overflow.
		if (top == array.length) 
			top = 0;

		// Check for størrelse.
		if ( top % array.length == bottom)
			resize();
	}	

	public boolean isEmpty() {
		if (size == 0) 
			return true;
		else 
			return false;
	}
	
	private void resize() {
		String[] newarray = new String[2 * array.length];
		for (int i = bottom ; i < bottom + size ; i++) {
			newarray[i] = array[i % array.length];
		}
		top = bottom + size;
		array = newarray;
	}
	
	public void PrintRel() {
		for (int i = bottom ; i < bottom + size ; i++) {
			System.out.println( (i - bottom) + " - " + array[i % array.length]);
		}
	}
	
	public synchronized String Retrieve() {
		if (bottom == top) 
			return null;
			
		String tmp = array[bottom];
		array[bottom] = null;
		bottom++;
		size--;
		
		if (bottom == array.length)
			bottom = 0;
		
		return tmp;
	}
}