/*
 *  Copyright (C) 2017 - This file is part of libecc project
 *
 *  Authors:
 *      Ryad BENADJILA <ryadbenadjila@gmail.com>
 *      Arnaud EBALARD <arnaud.ebalard@ssi.gouv.fr>
 *      Jean-Pierre FLORI <jean-pierre.flori@ssi.gouv.fr>
 *
 *  Contributors:
 *      Nicolas VIVET <nicolas.vivet@ssi.gouv.fr>
 *      Karim KHALFALLAH <karim.khalfallah@ssi.gouv.fr>
 *
 *  This software is licensed under a dual BSD and GPL v2 license.
 *  See LICENSE file at the root folder of the project.
 */
/*
 * The purpose of this example is to implement Pollard's rho
 * algorithm to find non-trivial factors of a composite natural
 * number.
 * The prime numbers decomposition of the natural number is
 * recovered through repeated Pollard's rho. Primality checking
 * is performed using a Miller-Rabin test.
 *
 * WARNING: the code in this example is only here to illustrate
 * how to use the NN layer API. This code has not been designed
 * for production purposes (e.g. no effort has been made to make
 * it constant time).
 *
 *
 */

/* We include the NN layer API header */
#include <libecc/libarith.h>

/* Declare our Miller-Rabin test implemented
 * in another module.
 */
ATTRIBUTE_WARN_UNUSED_RET int miller_rabin(nn_src_t n, const unsigned int t, int *check);

ATTRIBUTE_WARN_UNUSED_RET int pollard_rho(nn_t d, nn_src_t n, const word_t c);
/* Pollard's rho main function, as described in
 * "Handbook of Applied Cryptography".
 *
 * Pollard's rho:
 * ==============
 * See "Handbook of Applied Cryptography", alorithm 3.9:
 *
 *   Algorithm Pollard’s rho algorithm for factoring integers
 *   INPUT: a composite integer n that is not a prime power.
 *   OUTPUT: a non-trivial factor d of n.
 *      1. Set a←2, b←2.
 *      2. For i = 1, 2, ... do the following:
 *        2.1 Compute a←a^2 + 1 mod n, b←b^2 + 1 mod n, b←b^2 + 1 mod n.
 *        2.2 Compute d = gcd(a − b, n).
 *        2.3 If 1 < d < n then return(d) and terminate with success.
 *        2.4 If d = n then terminate the algorithm with failure (see Note 3.12).
 */
int pollard_rho(nn_t d, nn_src_t n, const word_t c)
{
        int ret, cmp, cmp1, cmp2;
        /* Temporary a and b variables */
        nn a, b, tmp, one, c_bignum;
        a.magic = b.magic = tmp.magic = one.magic = c_bignum.magic = WORD(0);

        /* Initialize variables */
        ret = nn_init(&a, 0); EG(ret, err);
        ret = nn_init(&b, 0); EG(ret, err);
        ret = nn_init(&tmp, 0); EG(ret, err);
        ret = nn_init(&one, 0); EG(ret, err);
        ret = nn_init(&c_bignum, 0); EG(ret, err);
        ret = nn_init(d, 0); EG(ret, err);

        MUST_HAVE((c > 0), ret, err);

        /* Zeroize the output */
        ret = nn_zero(d); EG(ret, err);
        ret = nn_one(&one); EG(ret, err);
        /* 1. Set a←2, b←2. */
        ret = nn_set_word_value(&a, WORD(2)); EG(ret, err);
        ret = nn_set_word_value(&b, WORD(2)); EG(ret, err);
        ret = nn_set_word_value(&c_bignum, c); EG(ret, err);

        /* For i = 1, 2, . . . do the following: */
        while (1) {
                int sign;
                /* 2.1 Compute a←a^2 + c mod n */
                ret = nn_sqr(&a, &a); EG(ret, err);
                ret = nn_add(&a, &a, &c_bignum); EG(ret, err);
                ret = nn_mod(&a, &a, n); EG(ret, err);
                /* 2.1 Compute b←b^2 + c mod n twice in a row */
                ret = nn_sqr(&b, &b); EG(ret, err);
                ret = nn_add(&b, &b, &c_bignum); EG(ret, err);
                ret = nn_mod(&b, &b, n); EG(ret, err);
                ret = nn_sqr(&b, &b); EG(ret, err);
                ret = nn_add(&b, &b, &c_bignum); EG(ret, err);
                ret = nn_mod(&b, &b, n); EG(ret, err);
                /* 2.2 Compute d = gcd(a − b, n) */
                ret = nn_cmp(&a, &b, &cmp); EG(ret, err);
                if (cmp >= 0) {
                        ret = nn_sub(&tmp, &a, &b); EG(ret, err);
                } else {
                        ret = nn_sub(&tmp, &b, &a); EG(ret, err);
                }
                ret = nn_gcd(d, &tmp, n, &sign); EG(ret, err);
                ret = nn_cmp(d, n, &cmp1); EG(ret, err);
                ret = nn_cmp(d, &one, &cmp2); EG(ret, err);
                if ((cmp1 < 0) && (cmp2 > 0)) {
                        ret = 0;
                        goto err;
                }
                ret = nn_cmp(d, n, &cmp); EG(ret, err);
                if (cmp == 0) {
                        ret = -1;
                        goto err;
                }
        }
 err:
        /* Uninitialize local variables */
        nn_uninit(&a);
        nn_uninit(&b);
        nn_uninit(&tmp);
        nn_uninit(&one);
        nn_uninit(&c_bignum);

        return ret;
}

ATTRIBUTE_WARN_UNUSED_RET int find_divisors(nn_src_t in);
/* Maximum number of divisors we support */
#define MAX_DIVISORS 10
/* Function to find prime divisors of the NN input */
int find_divisors(nn_src_t in)
{
        int n_divisors_found, i, found, ret, check, cmp;
        nn n;
        nn divisors[MAX_DIVISORS];
        word_t c;

        n.magic = WORD(0);
        for(i = 0; i < MAX_DIVISORS; i++){
                divisors[i].magic = WORD(0);
        }

        ret = nn_check_initialized(in); EG(ret, err);

        ext_printf("=================\n");
        nn_print("Finding factors of:", in);

        /* First, check primality of the input */
        ret = miller_rabin(in, 10, &check); EG(ret, err);
        if (check) {
                ext_printf("The number is probably prime, leaving ...\n");
                ret = -1;
                goto err;
        }
        ext_printf("The number is composite, performing Pollard's rho\n");

        ret = nn_init(&n, 0); EG(ret, err);
        ret = nn_copy(&n, in); EG(ret, err);
        for (i = 0; i < MAX_DIVISORS; i++) {
                ret = nn_init(&(divisors[i]), 0); EG(ret, err);
        }

        n_divisors_found = 0;
        c = 0;
        while (1) {
                c++;
                ret = pollard_rho(&(divisors[n_divisors_found]), &n, c);
                if (ret) {
                        continue;
                }
                found = 0;
                for (i = 0; i < n_divisors_found; i++) {
                        ret = nn_cmp(&(divisors[n_divisors_found]), &(divisors[i]), &cmp); EG(ret, err);
                        if (cmp == 0) {
                                found = 1;
                        }
                }
                if (found == 0) {
                        nn q, r;
                        ret = nn_init(&q, 0); EG(ret, err);
                        ret = nn_init(&r, 0); EG(ret, err);
                        ext_printf("Pollard's rho succeded\n");
                        nn_print("d:", &(divisors[n_divisors_found]));
                        /*
                         * Now we can launch the algorithm again on n / d
                         * to find new divisors. If n / d is prime, we are done!
                         */
                        ret = nn_divrem(&q, &r, &n, &(divisors[n_divisors_found])); EG(ret, err);
                        /*
                         * Check n / d primality with Miller-Rabin (security
                         * parameter of 10)
                         */
                        ret = miller_rabin(&q, 10, &check); EG(ret, err);
                        if (check == 1) {
                                nn_print("Last divisor is prime:", &q);
                                nn_uninit(&q);
                                nn_uninit(&r);
                                break;
                        }
                        nn_print("Now performing Pollard's rho on:", &q);
                        ret = nn_copy(&n, &q); EG(ret, err);
                        nn_uninit(&q);
                        nn_uninit(&r);
                        c = 0;
                        n_divisors_found++;
                        if (n_divisors_found == MAX_DIVISORS) {
                                ext_printf
                                        ("Max divisors reached, leaving ...\n");
                                break;
                        }
                }
        }
        ret = 0;

err:
        ext_printf("=================\n");
        nn_uninit(&n);
        for (i = 0; i < MAX_DIVISORS; i++) {
                nn_uninit(&(divisors[i]));
        }
        return ret;
}

#ifdef NN_EXAMPLE
int main(int argc, char *argv[])
{
        int ret;

        /* Fermat F5 = 2^32 + 1 = 641 x 6700417 */
        const unsigned char fermat_F5[] = { 0x01, 0x00, 0x00, 0x00, 0x01 };
        /* Fermat F6 = 2^64 + 1 = 274177 x 67280421310721 */
        const unsigned char fermat_F6[] =
                { 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01 };
        nn n;
        n.magic = WORD(0);

        FORCE_USED_VAR(argc);
        FORCE_USED_VAR(argv);

        ret = nn_init(&n, 0); EG(ret, err);
        /* Execute factorization on F5 */
        ret = nn_init_from_buf(&n, fermat_F5, sizeof(fermat_F5)); EG(ret, err);
        ret = find_divisors(&n); EG(ret, err);
        /* Execute factorization on F6 */
        ret = nn_init_from_buf(&n, fermat_F6, sizeof(fermat_F6)); EG(ret, err);
        ret = find_divisors(&n); EG(ret, err);
        /* Execute factorization on a random 80 bits number */
        ret = nn_one(&n); EG(ret, err);
        /* Compute 2**80 = 0x1 << 80 */
        ret = nn_lshift(&n, &n, 80); EG(ret, err);
        ret = nn_get_random_mod(&n, &n); EG(ret, err);
        ret = find_divisors(&n); EG(ret, err);

        return 0;
err:
        return -1;
}
#endif