From 763be884137d37773632da327e73d214e33757e0 Mon Sep 17 00:00:00 2001
From: Faiz Rahiemy <frachiemy@gmail.com>
Date: Tue, 3 Sep 2024 21:46:02 +0700
Subject: [PATCH] Check digest function to prevent error on OTP Generation
 (#170)

## Summary
This pull request introduces a validation check in the `init` and
`generate_otp` functions to prevent errors caused by incompatible digest
functions and inadequate digest sizes. The changes aim to improve the
stability and reliability of OTP generation. While the `RFC4226` states
that the base function is `sha1`, the existing code will allow user to
use other hash functions such as `MD5` and `SHAKE-128` due to their
availability in the hashlib.

## Changes Made
### Digest Function Validation
Added checks to disallow the use of `hashlib.md5` and
`hashlib.shake_128` as digest functions in the `__init__` function of
`otp.py`, `hotp.py`, and `totp.py`. These functions are not suitable for
OTP generation due to their shorter hash sizes.
### Digest Size Check
Implemented a check to ensure that the digest size is not lower than `18
bytes`. This prevents an `IndexError` that could occur when the last
hash byte is `0xF`, causing the `generate_otp` function to set the
offset to `15`. The subsequent operation would attempt to access
`hmac_hash[offset + 1]` and so on, leading to an out-of-bounds error for
`md5` and `shake128` due to their shorter digest lengths.

## Impact
- Prevents potential errors and crashes during OTP generation by
ensuring only suitable digest functions and sizes are used.
- Improves overall code robustness by validating inputs before
proceeding with OTP generation.

## Testing
- Tested that `hashlib.md5` and `hashlib.shake128` are correctly
rejected as digest functions on `__init__`.
- Tested the digest size check to ensure that digest lengths below 18
bytes trigger the appropriate error handling on `generate_otp`.
- Added `DigestFunctionTest` unit test as a negative test case where it
will trigger an error if the selected digest function is `md5` or
`shake_128`.
---
 src/pyotp/hotp.py |  5 +++++
 src/pyotp/otp.py  |  7 +++++++
 src/pyotp/totp.py |  5 +++++
 test.py           | 10 ++++++++++
 4 files changed, 27 insertions(+)

diff --git a/src/pyotp/hotp.py b/src/pyotp/hotp.py
index be9d536..84b6e9b 100644
--- a/src/pyotp/hotp.py
+++ b/src/pyotp/hotp.py
@@ -29,6 +29,11 @@ def __init__(
         """
         if digest is None:
             digest = hashlib.sha1
+        elif digest in [
+            hashlib.md5,
+            hashlib.shake_128
+        ]:
+            raise ValueError("selected digest function must generate digest size greater than or equals to 18 bytes")
 
         self.initial_count = initial_count
         super().__init__(s=s, digits=digits, digest=digest, name=name, issuer=issuer)
diff --git a/src/pyotp/otp.py b/src/pyotp/otp.py
index 9f0d1af..b339b8f 100644
--- a/src/pyotp/otp.py
+++ b/src/pyotp/otp.py
@@ -21,6 +21,11 @@ def __init__(
         if digits > 10:
             raise ValueError("digits must be no greater than 10")
         self.digest = digest
+        if digest in [
+            hashlib.md5,
+            hashlib.shake_128
+        ]:
+            raise ValueError("selected digest function must generate digest size greater than or equals to 18 bytes")
         self.secret = s
         self.name = name or "Secret"
         self.issuer = issuer
@@ -33,6 +38,8 @@ def generate_otp(self, input: int) -> str:
         if input < 0:
             raise ValueError("input must be positive integer")
         hasher = hmac.new(self.byte_secret(), self.int_to_bytestring(input), self.digest)
+        if hasher.digest_size < 18:
+            raise ValueError("digest size is lower than 18 bytes, which will trigger error on otp generation")
         hmac_hash = bytearray(hasher.digest())
         offset = hmac_hash[-1] & 0xF
         code = (
diff --git a/src/pyotp/totp.py b/src/pyotp/totp.py
index fd49ed0..ae077bc 100644
--- a/src/pyotp/totp.py
+++ b/src/pyotp/totp.py
@@ -32,6 +32,11 @@ def __init__(
         """
         if digest is None:
             digest = hashlib.sha1
+        elif digest in [
+            hashlib.md5,
+            hashlib.shake_128
+        ]:
+            raise ValueError("selected digest function must generate digest size greater than or equals to 18 bytes")
 
         self.interval = interval
         super().__init__(s=s, digits=digits, digest=digest, name=name, issuer=issuer)
diff --git a/test.py b/test.py
index 04f7374..8c902af 100755
--- a/test.py
+++ b/test.py
@@ -334,6 +334,16 @@ def test_valid_window(self):
         self.assertTrue(totp.verify("681610", 200, 1))
         self.assertFalse(totp.verify("195979", 200, 1))
 
+class DigestFunctionTest(unittest.TestCase):
+    def test_md5(self):
+        with self.assertRaises(ValueError) as cm:
+            pyotp.OTP(s="secret", digest=hashlib.md5)
+        self.assertEqual("selected digest function must generate digest size greater than or equals to 18 bytes", str(cm.exception))
+
+    def test_shake128(self):
+        with self.assertRaises(ValueError) as cm:
+            pyotp.OTP(s="secret", digest=hashlib.shake_128)
+        self.assertEqual("selected digest function must generate digest size greater than or equals to 18 bytes", str(cm.exception))
 
 class ParseUriTest(unittest.TestCase):
     def test_invalids(self):