From 0c6283729632334a4fda10fba1d1a8422fa33052 Mon Sep 17 00:00:00 2001
From: Felix Fontein <felix@fontein.de>
Date: Mon, 29 Apr 2024 08:50:28 +0200
Subject: [PATCH] crypto.math module utils: add some tests, fix
 quick_is_not_prime() for small primes (#733)

* Fix quick_is_not_prime() for small primes. Add some tests.

* Fix return value of convert_int_to_bytes(0, 0) on Python 2.

* Add some more test cases.

* Simplify the changelog and point out that these errors only happen for cases not happening in regular use.
---
 changelogs/fragments/733-math-prime.yml       |   2 +
 plugins/module_utils/crypto/math.py           |  15 ++-
 .../plugins/module_utils/crypto/test_math.py  | 102 ++++++++++++++++++
 3 files changed, 117 insertions(+), 2 deletions(-)
 create mode 100644 changelogs/fragments/733-math-prime.yml
 create mode 100644 tests/unit/plugins/module_utils/crypto/test_math.py

diff --git a/changelogs/fragments/733-math-prime.yml b/changelogs/fragments/733-math-prime.yml
new file mode 100644
index 00000000..a0d96a0b
--- /dev/null
+++ b/changelogs/fragments/733-math-prime.yml
@@ -0,0 +1,2 @@
+bugfixes:
+  - "crypto.math module utils - change return values for ``quick_is_not_prime()`` and ``convert_int_to_bytes(0, 0)`` for special cases that do not appear when using the collection (https://github.com/ansible-collections/community.crypto/pull/733)."
diff --git a/plugins/module_utils/crypto/math.py b/plugins/module_utils/crypto/math.py
index f56f22d3..b329dbe1 100644
--- a/plugins/module_utils/crypto/math.py
+++ b/plugins/module_utils/crypto/math.py
@@ -42,9 +42,18 @@ def quick_is_not_prime(n):
     that we could not detect quickly whether it is not prime.
     '''
     if n <= 2:
-        return True
+        return n < 2
     # The constant in the next line is the product of all primes < 200
-    if simple_gcd(n, 7799922041683461553249199106329813876687996789903550945093032474868511536164700810) > 1:
+    prime_product = 7799922041683461553249199106329813876687996789903550945093032474868511536164700810
+    gcd = simple_gcd(n, prime_product)
+    if gcd > 1:
+        if n < 200 and gcd == n:
+            # Explicitly check for all primes < 200
+            return n not in (
+                2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83,
+                89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179,
+                181, 191, 193, 197, 199,
+            )
         return True
     # TODO: maybe do some iterations of Miller-Rabin to increase confidence
     # (https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test)
@@ -106,6 +115,8 @@ if sys.version_info[0] >= 3:
 else:
     # Python 2
     def _convert_int_to_bytes(count, n):
+        if n == 0 and count == 0:
+            return ''
         h = '%x' % n
         if len(h) > 2 * count:
             raise Exception('Number {1} needs more than {0} bytes!'.format(count, n))
diff --git a/tests/unit/plugins/module_utils/crypto/test_math.py b/tests/unit/plugins/module_utils/crypto/test_math.py
new file mode 100644
index 00000000..31ccad2e
--- /dev/null
+++ b/tests/unit/plugins/module_utils/crypto/test_math.py
@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2024, Felix Fontein <felix@fontein.de>
+# GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+import pytest
+
+from ansible_collections.community.crypto.plugins.module_utils.crypto.math import (
+    binary_exp_mod,
+    simple_gcd,
+    quick_is_not_prime,
+    convert_int_to_bytes,
+    convert_int_to_hex,
+)
+
+
+@pytest.mark.parametrize('f, e, m, result', [
+    (0, 0, 5, 1),
+    (0, 1, 5, 0),
+    (2, 1, 5, 2),
+    (2, 2, 5, 4),
+    (2, 3, 5, 3),
+    (2, 10, 5, 4),
+])
+def test_binary_exp_mod(f, e, m, result):
+    value = binary_exp_mod(f, e, m)
+    print(value)
+    assert value == result
+
+
+@pytest.mark.parametrize('a, b, result', [
+    (0, -123, -123),
+    (0, 123, 123),
+    (-123, 0, -123),
+    (123, 0, 123),
+    (-123, 1, 1),
+    (123, 1, 1),
+    (1, -123, -1),
+    (1, 123, 1),
+    (1024, 10, 2),
+])
+def test_simple_gcd(a, b, result):
+    value = simple_gcd(a, b)
+    print(value)
+    assert value == result
+
+
+@pytest.mark.parametrize('n, result', [
+    (-2, True),
+    (0, True),
+    (1, True),
+    (2, False),
+    (3, False),
+    (4, True),
+    (5, False),
+    (6, True),
+    (7, False),
+    (8, True),
+    (9, True),
+    (10, True),
+    (211, False),  # the smallest prime number >= 200
+])
+def test_quick_is_not_prime(n, result):
+    value = quick_is_not_prime(n)
+    print(value)
+    assert value == result
+
+
+@pytest.mark.parametrize('no, count, result', [
+    (0, None, b''),
+    (0, 1, b'\x00'),
+    (0, 2, b'\x00\x00'),
+    (1, None, b'\x01'),
+    (1, 2, b'\x00\x01'),
+    (255, None, b'\xff'),
+    (256, None, b'\x01\x00'),
+])
+def test_convert_int_to_bytes(no, count, result):
+    value = convert_int_to_bytes(no, count=count)
+    print(value)
+    assert value == result
+
+
+@pytest.mark.parametrize('no, digits, result', [
+    (0, None, '0'),
+    (1, None, '1'),
+    (16, None, '10'),
+    (1, 3, '001'),
+    (255, None, 'ff'),
+    (256, None, '100'),
+    (256, 2, '100'),
+    (256, 3, '100'),
+    (256, 4, '0100'),
+])
+def test_convert_int_to_hex(no, digits, result):
+    value = convert_int_to_hex(no, digits=digits)
+    print(value)
+    assert value == result