import javax.crypto.spec.SecretKeySpec;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.math.BigInteger;
import java.net.InetAddress;
import java.net.Socket;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Base64;
import java.util.Random;
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class Client {

    private int port;

    private BigInteger p;
    private BigInteger g;
    private BigInteger secret;

    private static BufferedReader in = null;
    private static PrintWriter out = null;

    public Client(int port, BigInteger p, BigInteger g, BigInteger secret){

        this.port = port;
        this.p = p;
        this.g = g;
        this.secret = secret;
    }

    public void run(){

        for(;;){

            try{

                InetAddress addr = InetAddress.getByName("127.0.0.1") ;

                Socket socket = new Socket(addr, port);

                in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
                out = new PrintWriter(new java.io.BufferedWriter(new OutputStreamWriter(socket.getOutputStream())), true);

                BigInteger diffHellA = g.modPow(secret, p);
                System.out.println("Computed A Value: " + diffHellA);

                out.println("**DHA**" + diffHellA + "****");

                String diffHellBString = in.readLine().trim();

                BigInteger diffHellB = new BigInteger(
                        diffHellBString.substring(7, diffHellBString.length() - 4));

                BigInteger dHKey = diffHellB.modPow(secret, p);
                System.out.println("Shared Diffie Hellman key computed: " + dHKey);

                String nonce = generateRandomString(4);

                String dhPadded = String.format("%012d", dHKey);

                String base = nonce + dhPadded;

                byte[] aesBytes = sha256(base);
                byte[] newArray = Arrays.copyOfRange(aesBytes, 0, 16);
                System.out.println("128 bit key: " + Base64.getEncoder().encodeToString(newArray));
                SecretKeySpec aesKey = new SecretKeySpec(newArray, "AES");

                out.println("**NONCE**" + nonce + "****");
                out.println("**REQ****");

                String response = in.readLine().trim();
                String encryptedHash = response.substring(7, response.length() - 4);

                String file = decrypt(encryptedHash, base.getBytes(), aesKey);
                String fileTrim = file.substring(5, file.length());

                String fileHash = Base64.getEncoder().encodeToString(sha256(fileTrim));

                byte[] encryptedFileHash = encrypt(fileHash, base.getBytes(), aesKey);

                out.println("**VERIFY**" + fileHash + "****");

                System.out.println(new String(encryptedFileHash).length());

                String verified = in.readLine().trim();
                System.out.println(verified);

                socket.close();
                break;

            }catch(Exception e){

                e.printStackTrace(System.out);
            }
        }
    }

    public byte[] sha256(String key) throws NoSuchAlgorithmException{

        MessageDigest md = MessageDigest.getInstance("MD5");

        md.update(key.getBytes());
        return md.digest();
    }

    public static String generateRandomString(int len){

        String alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
        Random random = new Random();

        StringBuilder sb = new StringBuilder();

        alphabet.charAt(random.nextInt(alphabet.length()));

        for (int i = 0; i < len; i++) {
            sb.append(alphabet.charAt(random.nextInt(alphabet.length())));
        }

        return sb.toString();
    }

    public static String decrypt(String encrypted, byte[] bytesIv, SecretKeySpec key){

        try {

            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5PADDING");
            IvParameterSpec iv = new IvParameterSpec(bytesIv);
            cipher.init(2, key, iv);
            byte[] toDec = Base64.getDecoder().decode(encrypted);

            return new String(cipher.doFinal(toDec));
        }
        catch (Exception e) {

            System.out.println("Error while decrypting: " + e.toString());
        }

        return null;
    }

    public static byte[] encrypt(String plain, byte[] bytesIv, SecretKeySpec key){

        try {

            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5PADDING");

            IvParameterSpec iv = new IvParameterSpec(bytesIv);
            cipher.init(1, key, iv);

            return cipher.doFinal(plain.getBytes("UTF-8"));
        }

        catch (Exception e) {

            System.out.println("Error while encrypting: " + e.toString());
        }

        return null;
    }

    public static void main(String[] args){

        int port = 8080;
        BigInteger p = new BigInteger("53");
        BigInteger g = new BigInteger("27");
        BigInteger secret = new BigInteger("15");


        Client c = new Client(port, p, g, secret);
        c.run();
    }

}