Introduction to Cryptography

    static String parse(String msg, String key, BiFunction<Byte, Byte, Byte> f){

        // ensure we use only our alphabet
        msg = msg.toUpperCase ();
        key = key.toUpperCase ();

        // remove the white spaces because these decrease the randomness of the message
        // and they are not part of our alphabet
        byte[] msgb = msg.replaceAll ("\\s+","") .getBytes (StandardCharsets.US_ASCII);
        byte[] keyb = key.replaceAll ("\\s+",""). getBytes (StandardCharsets.US_ASCII);

        byte[] ret = new byte[msgb.length];

        for(int i = 0; i < msgb.length; i++){
            ret[i] = f.apply (msgb[i], keyb[i%keyb.length]);
        }
        return new String (ret);

    }

    static String encode(String msg, String key){
        return parse (msg, key, (m, k) -> (byte)((m-'A'+k-'A') % COUNT + 'A'));
    }

    static String decode(String msg, String key){
        int n = COUNT;
        return parse (msg, key, (m, k) ->(byte)((m-k>=0)? (m-k+'A') : (m-k+n+'A')));
    }

    static double symbol_entropy(){
        // log2(x) = log(x) / log(2)
        return (Math.log (COUNT) / Math.log (2));
    }

    static double theoretical_entropy(byte[] sequence){
        return  symbol_entropy () * sequence.length;
    }

    static byte randomLetter(int A, int B, int seed){
        int length = COUNT;

        long lA = A;
        long lB = B;
        long lSeed = seed;

        long n = Math.abs (lA * lSeed + lB) % length;

        // there is no additional source of randomness except for the seed
        return (byte)(n + 'A');
    }

    static byte[] random_sequence(int A, int B, int seed, int cnt){
        if (cnt <= 0) throw new IllegalArgumentException ("Count must be larger than 0");

        byte[] ret = new byte[cnt];
        for (int i = 0; i < cnt; i++){
            seed = ret[i] = randomLetter (A, B, seed);
        }
        return ret;
    }

    static class Actor {

        private byte[] sharedPass = null;
        private int[] internalRandoms = null;

        private static int PASS_LEN = 10;
        private static int INITIAL_BASE = 5;

        private Actor (){
            Random r = new Random ();

            // STEP: generate a set of internal numbers
            internalRandoms =  IntStream.range (0, PASS_LEN).map (x -> r.nextInt ()).toArray ();
        }

        private Actor(int[] randoms){
            internalRandoms = randoms;
        }

        public static int powModulo(int base, int x, int modulo){
            // https://en.wikipedia.org/wiki/Modular_exponentiation
            if (modulo == 1)
                return 0;

            int result = 1;
            base = base % modulo;

            while (x > 0) {
                if (x % 2 == 1) {
                    result = (result * base) % modulo;
                }
                x = x >> 1;
                base =(base * base) % modulo;
            }
            return result;
        }

        public static int powModulo(int base, int x) {
            return powModulo(base, x, COUNT);
        }

        public static int initialBaseModulo(int x){
            return powModulo (INITIAL_BASE, x);
        }

        public int[] getPublicMessage() {
            // STEP2: share the result of exponentiation on the public channel
            return IntStream.of (internalRandoms). map (Actor::initialBaseModulo).toArray ();
        }

        public void createSharedPass(int[] othersRand){
            // STEP3: raise to the other's public message to the power of internal randoms
            // to get to the commonly shared password
            sharedPass = new byte[PASS_LEN];
            for(int i = 0; i < PASS_LEN; i++)
                sharedPass[i] = (byte)(powModulo (othersRand[i], internalRandoms[i]) + 'A');
        }

        public String encodeMessage(String s){
            // STEP4: communicate
            return encode (s, new String (sharedPass));
        }

        public String decodeMessage(String msg){
            // STEP4: communicate
            return decode (msg, new String (sharedPass));
        }

        public String getPrivatePassword(){
            return new String (sharedPass);
        }

    }
    /**
    * Implementation of the actual exchange
    */ 
    public static void test_diffieHellman() throws Exception {

        var actor1 = new Actor ();
        var actor2 = new Actor ();

        actor2.createSharedPass (actor1.getPublicMessage ());
        actor1.createSharedPass (actor2.getPublicMessage ());

        var msg = actor1.encodeMessage ("HELLOWORLD");
        var decodedMsg = actor2.decodeMessage (msg);

        var pass1 = actor1.getPrivatePassword ();
        var pass2 = actor2.getPrivatePassword ();

        if(!pass1.equals (pass2) || !decodedMsg.equals ("HELLOWORLD"))
            throw new Exception ("Wrong Algorithm");
    }

    public static void test_diffieHellman() throws Exception {  

        /*
        [....code from above, new lines added below... ]
        */

        var crackedMessage = crackEncodedMessage (actor1.getPublicMessage (), actor2.getPublicMessage (), msg);

        if (!decodedMsg.equals (crackedMessage))
            throw new Exception ("Wrong Algorithm");

    }

    public static String crackEncodedMessage(
        int[] actor1PublicMessage, 
        int[] actor2PublicMessage, 
        String encodedMessage){

        var crackedInternalRandoms = new int[actor1PublicMessage.length];

        // STEP 1: 
        // for each element in the public message from the first actor,
        // find a number for which the initialPowModulo is equal to the observed value
        // O(pass length * modulo = pass length * 2^no_of_bits_in_modulo_prime_number) algorithm:
        for(int i = 0; i < crackedInternalRandoms.length; i++) {
            for (int j = 0; j < COUNT; j++) {
                if (Actor.initialPowModulo (j) == actor1PublicMessage[i]) {
                    crackedInternalRandoms[i] = j;
                    break;
                }
            }
        }

        // STEP 2: create an actor with these randoms as the internal password
        var actor3 = new Actor (crackedInternalRandoms);

        // STEP 3: use this fake internal password and the observed responses from 
        // the second actor to re-create the shared password
        actor3.createSharedPass (actor2PublicMessage);

        // STEP 4: decode the message
        return actor3.decodeMessage (encodedMessage);
    }