totp.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. # This file is dual licensed under the terms of the Apache License, Version
  2. # 2.0, and the BSD License. See the LICENSE file in the root of this repository
  3. # for complete details.
  4. from __future__ import annotations
  5. import typing
  6. from cryptography.hazmat.primitives import constant_time
  7. from cryptography.hazmat.primitives.twofactor import InvalidToken
  8. from cryptography.hazmat.primitives.twofactor.hotp import (
  9. HOTP,
  10. HOTPHashTypes,
  11. _generate_uri,
  12. )
  13. class TOTP:
  14. def __init__(
  15. self,
  16. key: bytes,
  17. length: int,
  18. algorithm: HOTPHashTypes,
  19. time_step: int,
  20. backend: typing.Any = None,
  21. enforce_key_length: bool = True,
  22. ):
  23. self._time_step = time_step
  24. self._hotp = HOTP(
  25. key, length, algorithm, enforce_key_length=enforce_key_length
  26. )
  27. def generate(self, time: int | float) -> bytes:
  28. counter = int(time / self._time_step)
  29. return self._hotp.generate(counter)
  30. def verify(self, totp: bytes, time: int) -> None:
  31. if not constant_time.bytes_eq(self.generate(time), totp):
  32. raise InvalidToken("Supplied TOTP value does not match.")
  33. def get_provisioning_uri(
  34. self, account_name: str, issuer: str | None
  35. ) -> str:
  36. return _generate_uri(
  37. self._hotp,
  38. "totp",
  39. account_name,
  40. issuer,
  41. [("period", int(self._time_step))],
  42. )