Main.java

package org.frost.superkey.test;

import java.io.IOException;
import java.security.*;
import java.util.concurrent.ThreadLocalRandom;

import org.apache.commons.lang3.time.*;

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

public class Main {
    public static boolean simulateFailure;

    public static void main(String[] args) throws IOException, NoSuchAlgorithmException, NoSuchPaddingException,
            InvalidKeyException, BadPaddingException, IllegalBlockSizeException, InterruptedException {
        // write your code here
        StopWatch watch = new StopWatch();
        System.out.print("Simulate wrong key attempt?[y/n]:");
        char response = (char) System.in.read();
        System.out.println();
        if (Character.toUpperCase(response) == 'Y') {
            simulateFailure = true;
            System.out.println("Simulating key failure...");
        }
        watch.start();
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(2048, new SecureRandom());
        System.out.println("@ " + watch.getTime() + "ms : Initialized key pair.");

        KeyPair keyPair = keyGen.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        if (simulateFailure) {
            KeyPair falseKeyPair = keyGen.generateKeyPair();
            privateKey = falseKeyPair.getPrivate();
        }
        System.out.println("@ " + watch.getTime() + "ms : Created public and private key instances accordingly.");

        byte[] challenge = new byte[245];
        ThreadLocalRandom.current().nextBytes(challenge);
        System.out.println("@ " + watch.getTime() + "ms : Created a random challenge.");

        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedChallenge = cipher.doFinal(challenge);
        System.out.println("@ " + watch.getTime() + "ms : Encrypted the challenge. Simulating network latency...");
        Thread.sleep(ThreadLocalRandom.current().nextLong(40, 300));
        System.out.println("@ "+ watch.getTime() + "ms : Client recieved, now calculating challenge.");

        cipher = Cipher.getInstance("RSA");
        byte[] newChallenge = new byte[245];
        try {
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            newChallenge = cipher.doFinal(encryptedChallenge);
        }
        catch (BadPaddingException e){
            System.out.println("@ "+ watch.getTime() + "ms : Detected incorrect key! Sending random stream of bytes to compensate.");
            ThreadLocalRandom.current().nextBytes(newChallenge);
        }

        System.out.println("@ "+ watch.getTime() + "ms : Client finished. Simulating network latency...");
        Thread.sleep(ThreadLocalRandom.current().nextLong(40, 300));
        System.out.println("@ " + watch.getTime() + "ms : Server recieved. To validate the key now...");

        if (java.util.Arrays.equals(challenge, newChallenge)) {
            System.out.println("@ " + watch.getTime() + "ms : Yes! Client has been validated.");
        }
        else {
            System.out.println("@ " + watch.getTime() + "ms : No! Client is a hacker and will be denied.");
        }
        System.in.read();
    }
}