Skip to content

Commit

Permalink
Check digest function to prevent error on OTP Generation (#170)
Browse files Browse the repository at this point in the history
## 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`.
  • Loading branch information
FaizRahiemy authored Sep 3, 2024
1 parent 5445bfe commit 763be88
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/pyotp/hotp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 7 additions & 0 deletions src/pyotp/otp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 = (
Expand Down
5 changes: 5 additions & 0 deletions src/pyotp/totp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 763be88

Please sign in to comment.