From 320b5de76489bf4ca4e1d9f1ffaee8834be76df8 Mon Sep 17 00:00:00 2001
From: Daniel Regeci <536331+ovx@users.noreply.github.com>
Date: Sat, 21 Dec 2024 08:05:05 +0700
Subject: [PATCH] fix: create_challenge options.number [#10]

---
 altcha/altcha.py     |  2 +-
 tests/test_altcha.py | 48 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/altcha/altcha.py b/altcha/altcha.py
index 1214fab..2cfc28b 100644
--- a/altcha/altcha.py
+++ b/altcha/altcha.py
@@ -247,7 +247,7 @@ def create_challenge(options):
         options.salt
         or base64.b16encode(secrets.token_bytes(salt_length)).decode("utf-8").lower()
     )
-    number = options.number or secrets.randbelow(max_number)
+    number = options.number if options.number is not None else secrets.randbelow(max_number + 1)
 
     salt_params = {}
     if "?" in salt:
diff --git a/tests/test_altcha.py b/tests/test_altcha.py
index c6db5dd..5ff2611 100644
--- a/tests/test_altcha.py
+++ b/tests/test_altcha.py
@@ -81,6 +81,54 @@ def test_verify_solution_failure(self):
         result, _ = verify_solution(payload_encoded, self.hmac_key, check_expires=False)
         self.assertFalse(result)
 
+    # Test for number=0
+    def test_verify_solution_zero(self):
+        options = ChallengeOptions(
+            algorithm="SHA-256",
+            max_number=10,
+            salt_length=16,
+            hmac_key=self.hmac_key,
+            salt="somesalt",
+            number=0,
+        )
+        challenge = create_challenge(options)
+        payload = Payload(
+            algorithm="SHA-256",
+            challenge=challenge.challenge,
+            number=0,
+            salt="somesalt",
+            signature=challenge.signature,
+        )
+        payload_encoded = base64.b64encode(
+            json.dumps(payload.__dict__).encode()
+        ).decode()
+        result, _ = verify_solution(payload_encoded, self.hmac_key, check_expires=False)
+        self.assertTrue(result)
+
+    # Test for number being inclusive with max_number
+    def test_verify_solution_max_number(self):
+        options = ChallengeOptions(
+            algorithm="SHA-256",
+            max_number=10,
+            salt_length=16,
+            hmac_key=self.hmac_key,
+            salt="somesalt",
+            number=10,
+        )
+        challenge = create_challenge(options)
+        payload = Payload(
+            algorithm="SHA-256",
+            challenge=challenge.challenge,
+            number=10,
+            salt="somesalt",
+            signature=challenge.signature,
+        )
+        payload_encoded = base64.b64encode(
+            json.dumps(payload.__dict__).encode()
+        ).decode()
+        result, _ = verify_solution(payload_encoded, self.hmac_key, check_expires=False)
+        self.assertTrue(result)
+
     def test_verify_solution_not_expired(self):
         options = ChallengeOptions(
             algorithm="SHA-256",