#
#   number.py : Number-theoretic functions
#
#  Part of the Python Cryptography Toolkit
#
#  Written by Andrew M. Kuchling, Barry A. Warsaw, and others
#
# ===================================================================
# The contents of this file are dedicated to the public domain.  To
# the extent that dedication to the public domain is not available,
# everyone is granted a worldwide, perpetual, royalty-free,
# non-exclusive license to exercise all rights associated with the
# contents of this file for any purpose whatsoever.
# No rights are reserved.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ===================================================================
#

import math
import sys
import struct
from Crypto import Random
from Crypto.Util.py3compat import _memoryview, iter_range

# Backward compatibility
_fastmath = None


def ceil_div(n, d):
    """Return ceil(n/d), that is, the smallest integer r such that r*d >= n"""

    if d == 0:
        raise ZeroDivisionError()
    if (n < 0) or (d < 0):
        raise ValueError("Non positive values")
    r, q = divmod(n, d)
    if (n != 0) and (q != 0):
        r += 1
    return r


def size (N):
    """Returns the size of the number N in bits."""

    bits = 0
    while N >> bits:
        bits += 1
    return bits


def getRandomInteger(N, randfunc=None):
    """Return a random number at most N bits long.

    If :data:`randfunc` is omitted, then :meth:`Random.get_random_bytes` is used.

    .. deprecated:: 3.0
        This function is for internal use only and may be renamed or removed in
        the future. Use :func:`Crypto.Random.random.getrandbits` instead.
    """

    if randfunc is None:
        randfunc = Random.get_random_bytes

    S = randfunc(N>>3)
    odd_bits = N % 8
    if odd_bits != 0:
        rand_bits = ord(randfunc(1)) >> (8-odd_bits)
        S = struct.pack('B', rand_bits) + S
    value = bytes_to_long(S)
    return value

def getRandomRange(a, b, randfunc=None):
    """Return a random number *n* so that *a <= n < b*.

    If :data:`randfunc` is omitted, then :meth:`Random.get_random_bytes` is used.

    .. deprecated:: 3.0
        This function is for internal use only and may be renamed or removed in
        the future. Use :func:`Crypto.Random.random.randrange` instead.
    """

    range_ = b - a - 1
    bits = size(range_)
    value = getRandomInteger(bits, randfunc)
    while value > range_:
        value = getRandomInteger(bits, randfunc)
    return a + value

def getRandomNBitInteger(N, randfunc=None):
    """Return a random number with exactly N-bits,
    i.e. a random number between 2**(N-1) and (2**N)-1.

    If :data:`randfunc` is omitted, then :meth:`Random.get_random_bytes` is used.

    .. deprecated:: 3.0
        This function is for internal use only and may be renamed or removed in
        the future.
    """

    value = getRandomInteger (N-1, randfunc)
    value |= 2 ** (N-1)                # Ensure high bit is set
    assert size(value) >= N
    return value

def GCD(x,y):
    """Greatest Common Denominator of :data:`x` and :data:`y`.
    """

    x = abs(x) ; y = abs(y)
    while x > 0:
        x, y = y % x, x
    return y

def inverse(u, v):
    """The inverse of :data:`u` *mod* :data:`v`."""

    u3, v3 = u, v
    u1, v1 = 1, 0
    while v3 > 0:
        q = u3 // v3
        u1, v1 = v1, u1 - v1*q
        u3, v3 = v3, u3 - v3*q
    while u1<0:
        u1 = u1 + v
    return u1

# Given a number of bits to generate and a random generation function,
# find a prime number of the appropriate size.

def getPrime(N, randfunc=None):
    """Return a random N-bit prime number.

    If randfunc is omitted, then :meth:`Random.get_random_bytes` is used.
    """
    if randfunc is None:
        randfunc = Random.get_random_bytes

    number=getRandomNBitInteger(N, randfunc) | 1
    while (not isPrime(number, randfunc=randfunc)):
        number=number+2
    return number


def _rabinMillerTest(n, rounds, randfunc=None):
    """_rabinMillerTest(n:long, rounds:int, randfunc:callable):int
    Tests if n is prime.
    Returns 0 when n is definitly composite.
    Returns 1 when n is probably prime.
    Returns 2 when n is definitly prime.

    If randfunc is omitted, then Random.new().read is used.

    This function is for internal use only and may be renamed or removed in
    the future.
    """
    # check special cases (n==2, n even, n < 2)
    if n < 3 or (n & 1) == 0:
        return n == 2
    # n might be very large so it might be beneficial to precalculate n-1
    n_1 = n - 1
    # determine m and b so that 2**b * m = n - 1 and b maximal
    b = 0
    m = n_1
    while (m & 1) == 0:
        b += 1
        m >>= 1

    tested = []
    # we need to do at most n-2 rounds.
    for i in iter_range (min (rounds, n-2)):
        # randomly choose a < n and make sure it hasn't been tested yet
        a = getRandomRange (2, n, randfunc)
        while a in tested:
            a = getRandomRange (2, n, randfunc)
        tested.append (a)
        # do the rabin-miller test
        z = pow (a, m, n) # (a**m) % n
        if z == 1 or z == n_1:
            continue
        composite = 1
        for r in iter_range(b):
            z = (z * z) % n
            if z == 1:
                return 0
            elif z == n_1:
                composite = 0
                break
        if composite:
            return 0
    return 1

def getStrongPrime(N, e=0, false_positive_prob=1e-6, randfunc=None):
    """
    Return a random strong *N*-bit prime number.
    In this context, *p* is a strong prime if *p-1* and *p+1* have at
    least one large prime factor.

    Args:
        N (integer): the exact length of the strong prime.
          It must be a multiple of 128 and > 512.
        e (integer): if provided, the returned prime (minus 1)
          will be coprime to *e* and thus suitable for RSA where
          *e* is the public exponent.
        false_positive_prob (float):
          The statistical probability for the result not to be actually a
          prime. It defaults to 10\ :sup:`-6`.
          Note that the real probability of a false-positive is far less. This is
          just the mathematically provable limit.
        randfunc (callable):
          A function that takes a parameter *N* and that returns
          a random byte string of such length.
          If omitted, :func:`Crypto.Random.get_random_bytes` is used.
    Return:
        The new strong prime.

    .. deprecated:: 3.0
        This function is for internal use only and may be renamed or removed in
        the future.
    """

    # This function was implemented following the
    # instructions found in the paper:
    #   "FAST GENERATION OF RANDOM, STRONG RSA PRIMES"
    #   by Robert D. Silverman
    #   RSA Laboratories
    #   May 17, 1997
    # which by the time of writing could be freely downloaded here:
    # http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.17.2713&rep=rep1&type=pdf

    if randfunc is None:
        randfunc = Random.get_random_bytes

    # Use the accelerator if available
    if _fastmath is not None:
        return _fastmath.getStrongPrime(long(N), long(e), false_positive_prob,
            randfunc)

    if (N < 512) or ((N % 128) != 0):
        raise ValueError ("bits must be multiple of 128 and > 512")

    rabin_miller_rounds = int(math.ceil(-math.log(false_positive_prob)/math.log(4)))

    # calculate range for X
    #   lower_bound = sqrt(2) * 2^{511 + 128*x}
    #   upper_bound = 2^{512 + 128*x} - 1
    x = (N - 512) >> 7;
    # We need to approximate the sqrt(2) in the lower_bound by an integer
    # expression because floating point math overflows with these numbers
    lower_bound = (14142135623730950489 * (2 ** (511 + 128*x))) //  10000000000000000000
    upper_bound = (1 << (512 + 128*x)) - 1
    # Randomly choose X in calculated range
    X = getRandomRange (lower_bound, upper_bound, randfunc)

    # generate p1 and p2
    p = [0, 0]
    for i in (0, 1):
        # randomly choose 101-bit y
        y = getRandomNBitInteger (101, randfunc)
        # initialize the field for sieving
        field = [0] * 5 * len (sieve_base)
        # sieve the field
        for prime in sieve_base:
            offset = y % prime
            for j in iter_range((prime - offset) % prime, len (field), prime):
                field[j] = 1

        # look for suitable p[i] starting at y
        result = 0
        for j in range(len(field)):
            composite = field[j]
            # look for next canidate
            if composite:
                continue
            tmp = y + j
            result = _rabinMillerTest (tmp, rabin_miller_rounds)
            if result > 0:
                p[i] = tmp
                break
        if result == 0:
            raise RuntimeError ("Couln't find prime in field. "
                                "Developer: Increase field_size")

    # Calculate R
    #     R = (p2^{-1} mod p1) * p2 - (p1^{-1} mod p2) * p1
    tmp1 = inverse (p[1], p[0]) * p[1]  # (p2^-1 mod p1)*p2
    tmp2 = inverse (p[0], p[1]) * p[0]  # (p1^-1 mod p2)*p1
    R = tmp1 - tmp2 # (p2^-1 mod p1)*p2 - (p1^-1 mod p2)*p1

    # search for final prime number starting by Y0
    #    Y0 = X + (R - X mod p1p2)
    increment = p[0] * p[1]
    X = X + (R - (X % increment))
    while 1:
        is_possible_prime = 1
        # first check candidate against sieve_base
        for prime in sieve_base:
            if (X % prime) == 0:
                is_possible_prime = 0
                break
        # if e is given make sure that e and X-1 are coprime
        # this is not necessarily a strong prime criterion but useful when
        # creating them for RSA where the p-1 and q-1 should be coprime to
        # the public exponent e
        if e and is_possible_prime:
            if e & 1:
                if GCD(e, X-1) != 1:
                    is_possible_prime = 0
            else:
                if GCD(e, (X-1) // 2) != 1:
                    is_possible_prime = 0

        # do some Rabin-Miller-Tests
        if is_possible_prime:
            result = _rabinMillerTest (X, rabin_miller_rounds)
            if result > 0:
                break
        X += increment
		# abort when X has more bits than requested
		# TODO: maybe we shouldn't abort but rather start over.
        if X >= 1 << N:
            raise RuntimeError ("Couln't find prime in field. "
                                "Developer: Increase field_size")
    return X

def isPrime(N, false_positive_prob=1e-6, randfunc=None):
    r"""Test if a number *N* is a prime.

    Args:
        false_positive_prob (float):
          The statistical probability for the result not to be actually a
          prime. It defaults to 10\ :sup:`-6`.
          Note that the real probability of a false-positive is far less.
          This is just the mathematically provable limit.
        randfunc (callable):
          A function that takes a parameter *N* and that returns
          a random byte string of such length.
          If omitted, :func:`Crypto.Random.get_random_bytes` is used.