Crown-OTP keystream gen time, attack speed and KPA

package fields.tests;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

import org.junit.jupiter.api.Test;

import fields.floatingpoint.FiniteRealVectorSpace;
import fields.floatingpoint.Reals;
import fields.floatingpoint.Reals.Real;
import fields.integers.Integers;
import fields.integers.Integers.IntE;
import fields.vectors.RealLattice;
import fields.vectors.SubRealInnerProductSpace;
import fields.vectors.Vector;


class CrownOTPTest {

    @Test
    void Test1WhatSizePlaintextTakes10SecondsToEncrypt() {
        System.out.println("\n-------------------------------------------------------------------------");
        System.out.println("TEST 1: Test what size plaintext takes unacceptable 10 seconds to encrypt");
        System.out.println("-------------------------------------------------------------------------");

        int key_bits = 256;
        int key_stream_len_bytes = 128;
        long KeystreamGenTimeMS;
        while (true) {
            System.out.printf("Testing %s-byte keystream...", key_stream_len_bytes);

            // Generate key
            int key_stream_len_bits = 8 * key_stream_len_bytes;
            Reals r = Reals.r(key_stream_len_bits);
            Integers z = Integers.z();
            IntE key = z.getRandomElement(z.power(z.getInteger(2), key_bits));

            // Generate keystream and measure time required to generate it
            long KeystreamGenStart = System.nanoTime();
            Real sqrt = r.positiveSqrt(r.getInteger(key));
            Real keystream = sqrt.fractionalPart();
            long KeystreamGenEnd = System.nanoTime();
            KeystreamGenTimeMS = (KeystreamGenEnd - KeystreamGenStart) / 1000000;
            System.out.printf("%d ms\n", KeystreamGenTimeMS);

            if (KeystreamGenTimeMS > 10000) {
                break;
            } else {
                key_stream_len_bytes *= 2;
            }

        }
        System.out.printf("\nEncrypting just %s kilobytes takes %d seconds.\n",
                key_stream_len_bytes/1024,
                KeystreamGenTimeMS/1000);
    }


    @Test
    void Test2WhenKeystreamGenBecomesSlowerThanBreakingIt() {
        System.out.println("\n------------------------------------------------------------------------------");
        System.out.println("TEST 2: Test when keystream generation becomes slower than breaking the cipher");
        System.out.println("------------------------------------------------------------------------------");

        int key_bits = 256;
        System.out.printf("\nKey size: %d bytes\n", key_bits/8);

        for (int key_stream_len_bytes = 128; key_stream_len_bytes < 1024 * 1024; key_stream_len_bytes *= 2) {

            // Generate key
            int key_stream_len_bits = 8 * key_stream_len_bytes;
            Reals r = Reals.r(key_stream_len_bits);
            Integers z = Integers.z();
            IntE key = z.getRandomElement(z.power(z.getInteger(2), key_bits));

            // Generate keystream and measure time required to generate it
            long KeystreamGenStart = System.nanoTime();
            Real sqrt = r.positiveSqrt(r.getInteger(key));
            Real keystream = sqrt.fractionalPart();
            long KeystreamGenEnd = System.nanoTime();
            long KeystreamGenTimeMS = (KeystreamGenEnd - KeystreamGenStart) / 1000000;

            // Break key from keystream and measure time required to break it
            long BreakRNGStart = System.nanoTime();
            IntE reconstructed = breakCrownOTP(keystream, 3 * key_bits, key_bits);
            long BreakRNGEnd = System.nanoTime();
            long BreakRNGTimeMS = (BreakRNGEnd - BreakRNGStart) / 1000000;

            if (key.equals(reconstructed)) {
                if (KeystreamGenTimeMS > BreakRNGTimeMS) {
                    System.out.printf("\nBreaking the cipher is faster than encrypting messages,\n" +
                            "when the length of sent message is longer than %d bytes.\n", key_stream_len_bytes);
                    break;
                }
            }

        }
    }


    @Test
    void Test3KPAAndExistentialForgeryWithoutOffset() {
        System.out.println("\n-----------------------------------");
        System.out.println("TEST 3: KPA and existential forgery");
        System.out.println("-----------------------------------");

        int key_bits = 256;

        System.out.println("\nSender:\n");

        // Define secret plaintext message
        String plaintext = "Hello Linda, this is Robert Grant, the CEO of the company." +
                " Please transfer 1,000 ETH to account 0x0123456789abcdef0123456789abcdef0123456";
        byte[] plaintextBytes = plaintext.getBytes(StandardCharsets.UTF_8);

        System.out.printf("  Original Plaintext: %s\n", plaintext);

        // Generate key
        int key_stream_bit_len = 8 * plaintext.length();
        Reals r = Reals.r(key_stream_bit_len);
        Integers z = Integers.z();
        IntE key = z.getRandomElement(z.power(z.getInteger(2), key_bits));

        // Generate the keystream
        Real sqrt = r.positiveSqrt(r.getInteger(key));
        String discardedOffset = String.valueOf(sqrt).split("\\.")[0];
        Real fractionalPart = sqrt.fractionalPart();
        String keystream = String.valueOf(fractionalPart).split("\\.")[1];
        byte[] keystreamBytes = new BigInteger(keystream).toByteArray();

        // Encrypt the plaintext by XORing it with the keystream
        byte[] ciphertextBytes = xorByteArrays(plaintextBytes, keystreamBytes);

        // Output
        String dashes = new String(new char[1096]).replace('\0', '-');
        System.out.printf("  Original key:       %s\n", key);
        System.out.printf("  Discarded KS:       %s\n", discardedOffset);
        System.out.printf("  Original Base10 KS: %s\n\n", keystream);

        printByteArray(plaintextBytes, "  PT");
        printByteArray(keystreamBytes, "  KS");
        System.out.printf("  XOR %s\n", dashes);
        printByteArray(ciphertextBytes, "  CT");

        // -------------------------------------------------------------------------------------------------------------

        System.out.println("\nMITM Attacker:\n");

        String knownPlaintext = "Hello Linda, this is Robert Grant, the CEO of the company.";
        System.out.printf("  Known plaintext:                       '%s'.\n", knownPlaintext);

        byte[] knownPlaintextBytes = knownPlaintext.getBytes(StandardCharsets.UTF_8);
        byte[] recoveredKeystreamBytes = xorByteArrays(ciphertextBytes, knownPlaintextBytes);

        String CTBits = byteArrayToBinaryString(ciphertextBytes);
        String recoveredOnesAndZeroes = byteArrayToBinaryString(recoveredKeystreamBytes);

        // Concatenate recovered keystream with trailing zeroes to match ciphertext length:
        int numberOfTrailingZeroes = CTBits.length() - recoveredOnesAndZeroes.length();
        recoveredOnesAndZeroes += "0".repeat(Math.max(0, numberOfTrailingZeroes));

        BigInteger partiallyCorrectRecoveredKSInt = new BigInteger(recoveredOnesAndZeroes, 2);
        String partialKPAKeystream = String.valueOf(partiallyCorrectRecoveredKSInt).substring(0, 90);
        String partialKPAFractional = String.format("0.%s", partialKPAKeystream);

        System.out.printf("  KPA Partial Base10 keystream recovery: %s\n", partialKPAKeystream);

        // Break the cipher by recovering key from partial keystream
        IntE reconstructedKey = breakCrownOTP(r.fromString(partialKPAFractional), 4 * key_bits, key_bits);

        // Reconstruct the full keystream
        Reals r_attacker = Reals.r(key_stream_bit_len);
        Real sqrt_attacker = r_attacker.positiveSqrt(r.getInteger(reconstructedKey));
        Real fractionalPartReconstruct = sqrt_attacker.fractionalPart();
        String keyStreamReconstruct = String.valueOf(fractionalPartReconstruct).split("\\.")[1];
        byte[] keyStreamReconstructBytes = new BigInteger(keyStreamReconstruct).toByteArray();
        System.out.printf("  Full keystream recovery with lattices: %s\n", keyStreamReconstruct);

        // Decrypt the message
        byte[] recoveredPlaintextBytes = xorByteArrays(keyStreamReconstructBytes, ciphertextBytes);
        String recoveredPlaintext = new String(recoveredPlaintextBytes, StandardCharsets.UTF_8);
        System.out.printf("  Recovered Plaintext:                   '%s'\n", recoveredPlaintext);

        String newPlaintext = "Hello Linda, this is Robert Grant, the CEO of the company." +
                " Please transfer 1,000,000,000 ETH to 0xfffffffffffffffffffffffffffffffffffffff";
        System.out.printf("  New, forged plaintext:                 '%s'\n", newPlaintext);
        byte[] newPlaintextBytes = newPlaintext.getBytes(StandardCharsets.UTF_8);
        byte[] newCiphertextBytes = xorByteArrays(newPlaintextBytes, keyStreamReconstructBytes);

        printByteArray(newPlaintextBytes, "\n  NEW PT");
        printByteArray(keyStreamReconstructBytes, "  REC KS");
        System.out.printf("  XOR %s\n", dashes);
        printByteArray(newCiphertextBytes, "  NEW CT");

        // -------------------------------------------------------------------------------------------------------------

        System.out.println("\nRecipient:\n");

        byte[] finalPlaintextBytes = xorByteArrays(newCiphertextBytes, keystreamBytes);
        String finalPlaintext = new String(finalPlaintextBytes, StandardCharsets.UTF_8);

        printByteArray(newCiphertextBytes, "  CT");
        printByteArray(keystreamBytes, "  KS");
        System.out.printf("  XOR %s\n", dashes);
        printByteArray(finalPlaintextBytes, "  PT");

        System.out.printf("\n  Recipient's plaintext: '%s'\n", finalPlaintext);
    }


    private static String byteArrayToBinaryString(byte[] byteArray) {
        StringBuilder BinaryString = new StringBuilder();
        for (byte b : byteArray) {
            String binary = String.format("%8s", Integer.toBinaryString(b & 0xFF)).replace(' ', '0');
            BinaryString.append(binary);
        }
        return String.valueOf(BinaryString);
    }


    private static void printByteArray(byte[] keystreamBytes, String name) {
        StringBuilder keystreamOnesZeroes = new StringBuilder();
        keystreamOnesZeroes.append(String.format("%s: ", name));
        for (byte b : keystreamBytes) {
            String binary = String.format("%8s", Integer.toBinaryString(b & 0xFF)).replace(' ', '0');
            keystreamOnesZeroes.append(binary);
        }
        System.out.println(keystreamOnesZeroes);
    }


    private static byte[] xorByteArrays(byte[] array1, byte[] array2) {
        int length = Math.min(array1.length, array2.length);
        byte[] result = new byte[length];

        for (int i = 0; i < length; i++) {
            result[i] = (byte) (array1[i] ^ array2[i]);
        }

        return result;
    }


    private IntE breakCrownOTP(Real sqrt, int precision, int bitRate) {
        Reals r = Reals.r(precision);
        Integers z = Integers.z();
        FiniteRealVectorSpace space = new FiniteRealVectorSpace(r, 3);
        Real power = r.getPowerOfTwo(bitRate + 4);
        List<Vector<Real>> generators = new ArrayList<>();
        generators.add(new Vector<>(r.one(),  r.zero(), r.multiply(2, sqrt, power)));
        generators.add(new Vector<>(r.zero(), r.one(),  r.negative(power)));
        SubRealInnerProductSpace<Real, Vector<Real>> subSpace = new SubRealInnerProductSpace<>(space, generators);
        RealLattice lattice  = new RealLattice(space, generators);
        Vector<Real> target = new Vector<>(r.zero(), r.zero(), r.multiply(-1, power, sqrt, sqrt));
        target = subSpace.project(target);
        Vector<IntE> result  = lattice.asVector(space.closestLatticePoint(target, lattice));
        Vector<Real> closest = lattice.fromVector(result);
        IntE intPart = closest.get(1).round();
        IntE diff = closest.get(2).round();
        IntE seed = z.add(z.multiply(intPart, intPart), diff);

        return seed;
    }
}