path: root/frozen_deps/ecdsa/test_numbertheory.py

import operator
from six import print_
from functools import reduce
import operator

try:
    import unittest2 as unittest
except ImportError:
    import unittest
import hypothesis.strategies as st
import pytest
from hypothesis import given, settings, example

try:
    from hypothesis import HealthCheck

    HC_PRESENT = True
except ImportError:  # pragma: no cover
    HC_PRESENT = False
from .numbertheory import (
    SquareRootError,
    factorization,
    gcd,
    lcm,
    jacobi,
    inverse_mod,
    is_prime,
    next_prime,
    smallprimes,
    square_root_mod_prime,
)


BIGPRIMES = (
    999671,
    999683,
    999721,
    999727,
    999749,
    999763,
    999769,
    999773,
    999809,
    999853,
    999863,
    999883,
    999907,
    999917,
    999931,
    999953,
    999959,
    999961,
    999979,
    999983,
)


@pytest.mark.parametrize(
    "prime, next_p", [(p, q) for p, q in zip(BIGPRIMES[:-1], BIGPRIMES[1:])]
)
def test_next_prime(prime, next_p):
    assert next_prime(prime) == next_p


@pytest.mark.parametrize("val", [-1, 0, 1])
def test_next_prime_with_nums_less_2(val):
    assert next_prime(val) == 2


@pytest.mark.parametrize("prime", smallprimes)
def test_square_root_mod_prime_for_small_primes(prime):
    squares = set()
    for num in range(0, 1 + prime // 2):
        sq = num * num % prime
        squares.add(sq)
        root = square_root_mod_prime(sq, prime)
        # tested for real with TestNumbertheory.test_square_root_mod_prime
        assert root * root % prime == sq

    for nonsquare in range(0, prime):
        if nonsquare in squares:
            continue
        with pytest.raises(SquareRootError):
            square_root_mod_prime(nonsquare, prime)


@st.composite
def st_two_nums_rel_prime(draw):
    # 521-bit is the biggest curve we operate on, use 1024 for a bit
    # of breathing space
    mod = draw(st.integers(min_value=2, max_value=2 ** 1024))
    num = draw(
        st.integers(min_value=1, max_value=mod - 1).filter(
            lambda x: gcd(x, mod) == 1
        )
    )
    return num, mod


@st.composite
def st_primes(draw, *args, **kwargs):
    if "min_value" not in kwargs:  # pragma: no branch
        kwargs["min_value"] = 1
    prime = draw(
        st.sampled_from(smallprimes)
        | st.integers(*args, **kwargs).filter(is_prime)
    )
    return prime


@st.composite
def st_num_square_prime(draw):
    prime = draw(st_primes(max_value=2 ** 1024))
    num = draw(st.integers(min_value=0, max_value=1 + prime // 2))
    sq = num * num % prime
    return sq, prime


@st.composite
def st_comp_with_com_fac(draw):
    """
    Strategy that returns lists of numbers, all having a common factor.
    """
    primes = draw(
        st.lists(st_primes(max_value=2 ** 512), min_size=1, max_size=10)
    )
    # select random prime(s) that will make the common factor of composites
    com_fac_primes = draw(
        st.lists(st.sampled_from(primes), min_size=1, max_size=20)
    )
    com_fac = reduce(operator.mul, com_fac_primes, 1)

    # select at most 20 lists (returned numbers),
    # each having at most 30 primes (factors) including none (then the number
    # will be 1)
    comp_primes = draw(
        st.integers(min_value=1, max_value=20).flatmap(
            lambda n: st.lists(
                st.lists(st.sampled_from(primes), max_size=30),
                min_size=1,
                max_size=n,
            )
        )
    )

    return [reduce(operator.mul, nums, 1) * com_fac for nums in comp_primes]


@st.composite
def st_comp_no_com_fac(draw):
    """
    Strategy that returns lists of numbers that don't have a common factor.
    """
    primes = draw(
        st.lists(
            st_primes(max_value=2 ** 512), min_size=2, max_size=10, unique=True
        )
    )
    # first select the primes that will create the uncommon factor
    # between returned numbers
    uncom_fac_primes = draw(
        st.lists(
            st.sampled_from(primes),
            min_size=1,
            max_size=len(primes) - 1,
            unique=True,
        )
    )
    uncom_fac = reduce(operator.mul, uncom_fac_primes, 1)

    # then build composites from leftover primes
    leftover_primes = [i for i in primes if i not in uncom_fac_primes]

    assert leftover_primes
    assert uncom_fac_primes

    # select at most 20 lists, each having at most 30 primes
    # selected from the leftover_primes list
    number_primes = draw(
        st.integers(min_value=1, max_value=20).flatmap(
            lambda n: st.lists(
                st.lists(st.sampled_from(leftover_primes), max_size=30),
                min_size=1,
                max_size=n,
            )
        )
    )

    numbers = [reduce(operator.mul, nums, 1) for nums in number_primes]

    insert_at = draw(st.integers(min_value=0, max_value=len(numbers)))
    numbers.insert(insert_at, uncom_fac)
    return numbers


HYP_SETTINGS = {}
if HC_PRESENT:  # pragma: no branch
    HYP_SETTINGS["suppress_health_check"] = [
        HealthCheck.filter_too_much,
        HealthCheck.too_slow,
    ]
    # the fa