rsa

/*
Copyright (C) 2018-2024 Geoffrey Daniels. https://gpdaniels.com/

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License only.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#pragma once
#ifndef GTL_CRYPTO_RSA_HPP
#define GTL_CRYPTO_RSA_HPP

// Summary: An implementation of the RSA (Rivest-Shamir-Adleman) asymmetric encryption algorithm.

#ifndef NDEBUG
#   if defined(_MSC_VER)
#       define __builtin_trap() __debugbreak()
#   endif
/// @brief A simple assert macro to break the program if the rsa is misused.
#   define GTL_RSA_ASSERT(ASSERTION, MESSAGE) static_cast<void>((ASSERTION) || (__builtin_trap(), 0))
#else
/// @brief At release time the assert macro is implemented as a nop.
#   define GTL_RSA_ASSERT(ASSERTION, MESSAGE) static_cast<void>(0)
#endif

#include <math/big_unsigned>

#if defined(_MSC_VER)
#pragma warning(push, 0)
#endif

#include <functional>

#if defined(_MSC_VER)
#pragma warning(pop)
#endif

namespace gtl {
    /// @brief  The rsa class implements all the routines needed to generate primes and used them to perform asymmetric encryption.
    class rsa final {
    public:
        /// @brief  The public key is used to encrypt data or verify signed data.
        class public_key_type final {
        public:
            gtl::big_unsigned public_modulus;
            gtl::big_unsigned public_exponent;
        };

        /// @brief  The private key is used to decrypt data or sign data.
        /// @note   Only the public_modulus and private_exponent are needed for decryption.
        class private_key_type final {
        public:
            gtl::big_unsigned public_modulus;
            gtl::big_unsigned public_exponent;
            gtl::big_unsigned private_exponent;
            gtl::big_unsigned primes[2];
            gtl::big_unsigned exponents[2];
            gtl::big_unsigned coefficient;
        };

        /// @brief  Storage class for a public and private key pair.
        class key_type final {
        public:
            public_key_type public_key;
            private_key_type private_key;
        };

    private:
        /// @brief  Generate a random big number.
        /// @param  minimum The lower inclusive bound for the returned random number.
        /// @param  maximum The upper exclusive bound for the returned random number.
        /// @return A random number between the (inclusive) minimum and (exclusive) maximum bounds.
        static gtl::big_unsigned generate_random(const std::function<unsigned int()>& random_generator, const gtl::big_unsigned& minimum, const gtl::big_unsigned& maximum) {
            // Generate random bits
            gtl::big_unsigned random = minimum;
            while (random < maximum) {
                random = (random << gtl::big_unsigned::chunk_bits) | random_generator();
            }
            return (minimum + random) % maximum;
        }

        /// @brief  Generate a random prime number.
        /// @param  size_bytes The number of bytes in the prime.
        /// @param  miller_rabin_iterations The number of miller rabin test iterations to validate primality.
        /// @return A random prime number.
        static gtl::big_unsigned generate_prime(const std::function<unsigned int()>& random_generator, unsigned int size_bytes, unsigned int miller_rabin_iterations) {
            gtl::big_unsigned prime;
            do {
                // Generate a number.
                prime = rsa::generate_random(random_generator, 1u, gtl::big_unsigned(1u) << (size_bytes * 8));
                // Ensure it's got the lsb and msb set, aka odd and full size.
                prime |= (gtl::big_unsigned(1u) << ((size_bytes * 8u) - 1u)) | 1u;
                // Check if the number is prime, if not try again.
            } while (!rsa::is_prime(random_generator, prime, miller_rabin_iterations));
            return prime;
        }

        /// @brief  Check if a given number is a prime number.
        /// @param  prime_candidate The prime number candidate to test.
        /// @param  miller_rabin_iterations The number of miller rabin test iterations to validate primality.
        /// @return true if the prime number candidate is probably a prime number.
        static bool is_prime(const std::function<unsigned int()>& random_generator, const gtl::big_unsigned& prime_candidate, unsigned int miller_rabin_iterations = 64) {
            // Corner cases.
            if (prime_candidate < 2) {
                return false;
            }

            // List of the first 256 primes.
            constexpr static const unsigned int first_256_primes[256] = {
                   2,    3,    5,    7,   11,   13,   17,   19,   23,   29,   31,   37,   41,   43,   47,   53,
                  59,   61,   67,   71,  179,  181,  191,  193,  197,  199,  211,  223,  227,  229,  233,  239,
                 241,  251,  257,  263,  269,  271,  277,  281,  283,  293,  307,  311,  313,  317,  331,  337,
                 347,  349,  353,  359,  367,  373,  379,  383,  389,  397,  401,  409,  419,  421,  431,  433,
                 439,  443,  449,  457,  461,  463,  467,  479,  487,  491,  499,  503,  509,  521,  523,  541,
                 547,  557,  563,  569,  571,  577,  587,  593,  599,  601,  607,  613,  617,  619,  631,  641,
                 643,  647,  653,  659,  661,  673,  677,  683,  691,  701,  709,  719,  727,  733,  739,  743,
                 751,  757,  761,  769,  773,  787,  797,  809,  811,  821,  823,  827,  829,  839,  853,  857,
                 859,  863,  877,  881,  883,  887,  907,  911,  919,  929,  937,  941,  947,  953,  967,  971,
                 977,  983,  991,  997, 1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069,
                1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163, 1171, 1181, 1187, 1193,
                1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249, 1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301,
                1303, 1307, 1319, 1321, 1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439,
                1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511, 1523, 1531, 1543, 1549,
                1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601, 1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657,
                1663, 1667, 1669, 1693, 1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783
            };
            // Check if the candidate is divisible by one of the first 256 primes.
            for (unsigned int i = 0; i < 256; ++i) {
                if ((prime_candidate % first_256_primes[i]) == 0) {
                    // If it is divisible it is only prime if it is exactly that value.
                    return prime_candidate == first_256_primes[i];
                }
            }

            // Calculate a helper variable for the miller rabin to avoid recalculating it each iteration.
            // power_of_two_multiplier is an odd number such that "power_of_two_multiplier = (prime_candidate - 1) / (2 ^ power_of_two)" for "power_of_two >= 1".
            gtl::big_unsigned power_of_two_multiplier = prime_candidate - 1;
            while (power_of_two_multiplier % 2 == 0) {
                power_of_two_multiplier /= 2;
            }

            // Iterate given nber of 'k' times
            for (unsigned int i = 0; i < miller_rabin_iterations; ++i) {
                if (!rsa::miller_rabin(random_generator, prime_candidate, power_of_two_multiplier)) {
                    return false;
                }
            }

            return true;
        }

        /// @brief  The miller rabin test to evaluate the primality of a number.
        /// @param  prime_candidate The prime number candidate to test.
        /// @param  power_of_two_multiplier Helper variable equal to "(prime_candidate - 1) / (2 ^ power_of_two)" for "power_of_two >= 1"
        /// @return true if the test past and the number could be a prime number, false if the number is definitely not prime.
        static bool miller_rabin(const std::function<unsigned int()>& random_generator, const gtl::big_unsigned& prime_candidate, gtl::big_unsigned power_of_two_multiplier) {
            // Pick a random number in "[2 ... prime_candidate - 2]".
            gtl::big_unsigned base = rsa::generate_random(random_generator, 2, prime_candidate - 2);

            // Compute "(base ^ power_of_two_multiplier) % prime_candidate".
            gtl::big_unsigned test = rsa::modular_exponentiation(base, power_of_two_multiplier, prime_candidate);

            // Corner case checks.
            if ((test == 1) || (test == (prime_candidate - 1))) {
                return true;
            }

            // Keep squaring "test" until:
            // - "power_of_two_multiplier" does not reach "prime_candidate - 1".
            // - "(test ^ 2) % prime_candidate" is not one.
            // - "(test ^ 2) % prime_candidate" is not "prime_candidate - 1".
            while (power_of_two_multiplier != (prime_candidate - 1)) {
                test = (test * test) % prime_candidate;
                power_of_two_multiplier <<= 1;
                if (test == 1) {
                    return false;
                }
                if (test == (prime_candidate - 1)) {
                    return true;
                }
            }
            return false;
        }

        /// @brief  Calculates the greatest common denominator of two big numbers.
        /// @param  value_a One of the numbers to calculate the greatest common denominator of.
        /// @param  value_b One of the numbers to calculate the greatest common denominator of.
        /// @return The greatest common denominator of the two numbers.
        /// @note   Also known as the greatest common devisor, or the highest common factor, or highest common divisor.
        static gtl::big_unsigned greatest_common_denominator(gtl::big_unsigned value_a, gtl::big_unsigned value_b) {
            // Ensure value_a is smaller than value_b.
            if (value_a > value_b) {
                gtl::big_unsigned swap;
                swap = value_a;
                value_a = value_b;
                value_b = swap;
            }
            // Reduce values until remainder is zero.
            while (true) {
                gtl::big_unsigned remainder = value_b % value_a;
                if (remainder == 0) {
                    return value_a;
                }
                value_b = value_a;
                value_a = remainder;
            }
        }

        /// @brief  Calculate the modular multiplicative inverse of a number.
        /// @param  value The value to calculate the modular multiplicative inverse of.
        /// @param  modulus The modulus used in the calculation.
        /// @return The modular multiplicative inverse of the input value.
        /// @note   The modular inverse of "value % modulus" is "value^-1" such that "(value * value^-1) % modulus" is one.
        static gtl::big_unsigned modular_inverse(gtl::big_unsigned value, const gtl::big_unsigned& modulus) {
            GTL_RSA_ASSERT(rsa::greatest_common_denominator(value, modulus) == 1, "Value and modulus must be co-primes.");
            gtl::big_unsigned result = 1;
            gtl::big_unsigned previous = 0;
            gtl::big_unsigned quotient = modulus;
            bool inverse = false;
            while (true) {
                gtl::big_unsigned remainder;
                quotient = gtl::big_unsigned::divide(quotient, value, remainder);
                if (remainder == 0) {
                    if (value != 1) {
                        return 0;
                    }
                    else if (inverse) {
                        return modulus - result;
                    }
                    else {
                        return result;
                    }
                }
                gtl::big_unsigned temp = result;
                result = result * quotient + previous;
                previous = temp;
                quotient = value;
                value = remainder;
                inverse = !inverse;
            }
        }

        /// @brief  Calculate the modular exponentiation of a number.
        /// @param  value The value to calculate the modular exponentiation of.
        /// @param  exponent The exponent to raise the value to.
        /// @param  modulus The modulus used in the calculation.
        /// @return The modular exponentiation of the input value.
        /// @note   The modular exponentiation of "value" is "(value ^ exponent) % modulus".
        static gtl::big_unsigned modular_exponentiation(const gtl::big_unsigned& value, const gtl::big_unsigned& exponent, const gtl::big_unsigned& modulus){
            gtl::big_unsigned result = 1;
            gtl::big_unsigned square = value;
            long long int bits = static_cast<long long int>(exponent.get_length_bits());
            for (long long int i = bits - 1; i >= 0; --i) {
                if (exponent.get_bit(static_cast<unsigned long long int>(i))) {
                    result = (result * square) % modulus;
                    square = (square * square) % modulus;
                }
                else {
                    square = (result * square) % modulus;
                    result = (result * result) % modulus;
                }
            }
            return result;
        }

    public:
        /// @brief  Generate a public and private key set.
        /// @param  size_bytes The number of bytes in the keys..
        /// @param  miller_rabin_iterations The number of miller rabin test iterations to validate primality.
        /// @return A public and private key set.
        static key_type generate_key_pair(const std::function<unsigned int()>& random_generator, unsigned int size_bytes, unsigned int exponent = 65537, unsigned int miller_rabin_iterations = 64) {
            key_type keys;
            keys.private_key.public_exponent = exponent;
            keys.public_key.public_exponent = exponent;
            gtl::big_unsigned phi;
            do {
                keys.private_key.primes[0] = rsa::generate_prime(random_generator, size_bytes / 2, miller_rabin_iterations);
                keys.private_key.primes[1] = rsa::generate_prime(random_generator, size_bytes / 2, miller_rabin_iterations);
                keys.private_key.public_modulus = keys.private_key.primes[0] * keys.private_key.primes[1];
                keys.public_key.public_modulus = keys.private_key.public_modulus;
                phi = (keys.private_key.primes[0] - 1) * (keys.private_key.primes[1] - 1);
            } while (
                (keys.private_key.primes[0] == keys.private_key.primes[1]) ||
                (rsa::greatest_common_denominator(phi, exponent) != 1)
            );
            keys.private_key.private_exponent = rsa::modular_inverse(keys.public_key.public_exponent, phi);
            keys.private_key.exponents[0] = keys.private_key.private_exponent % (keys.private_key.primes[0] - 1);
            keys.private_key.exponents[1] = keys.private_key.private_exponent % (keys.private_key.primes[1] - 1);
            keys.private_key.coefficient = rsa::modular_inverse(keys.private_key.primes[1], keys.private_key.primes[0]);
            return keys;
        }

    public:
        /// @brief  Transform a block of data writing the transformed data into the output.
        /// @param  data The block of data to encrypt, must be at least length bytes long.
        /// @param  length The length of the data, must be equal to the key length.
        /// @param  exponent Exponent used to transform the data, must be length bytes long.
        /// @param  modulus Modulus used to transform the data, must be length bytes long.
        /// @param  output Pointer to an output buffer that will receive the transformed data, must be at least length bytes long
        /// @note   For encryption the exponent is the public exponent, and modulus is the public modulus.
        /// @note   For decryption the exponent is the private exponent, and modulus is the public modulus.
        static void transform_block(const unsigned char* data, const unsigned int length, const unsigned char* exponent, const unsigned char* modulus, unsigned char* output) {
            rsa::modular_exponentiation(
                gtl::big_unsigned(data, length),
                gtl::big_unsigned(exponent, length),
                gtl::big_unsigned(modulus, length)
            ).to_bytes(output, length);
        }

    public:
        /// @brief  Encrypt a block of data, writing the edcrypted data into the output.
        /// @param  data The block of data to encrypt, must be at least length bytes long.
        /// @param  length The length of the data, must be equal to the key length.
        /// @param  public_key The public key data used to encrypt the data.
        /// @param  output Pointer to an output buffer that will receive the encrypted data, must be at least length bytes long.
        static void encrypt(const unsigned char* data, const unsigned int length, const public_key_type& public_key, unsigned char* output) {
            rsa::modular_exponentiation(
                gtl::big_unsigned(data, length),
                public_key.public_exponent,
                public_key.public_modulus
            ).to_bytes(output, length);
        }

        /// @brief  Decrypt a block of data, writing the decrypted data into the output.
        /// @param  data The block of data to decrypt, must be at least length bytes long.
        /// @param  length The length of the data, must be equal to the key length.
        /// @param  private_key The private key data used to decrypt the data.
        /// @param  output Pointer to an output buffer that will receive the decrypted data, must be at least length bytes long.
        static void decrypt(const unsigned char* data, const unsigned int length, const private_key_type& private_key, unsigned char* output) {
            rsa::modular_exponentiation(
                gtl::big_unsigned(data, length),
                private_key.private_exponent,
                private_key.public_modulus
            ).to_bytes(output, length);
        }
    };
}

#undef GTL_RSA_ASSERT

#endif // GTL_CRYPTO_RSA_HPP