codekingpro/portable-devtools

1import base642import hashlib3import hmac4from typing import Any, Optional5 6 7class OTP(object):8    """9    Base class for OTP handlers.10    """11 12    def __init__(13        self,14        s: str,15        digits: int = 6,16        digest: Any = hashlib.sha1,17        name: Optional[str] = None,18        issuer: Optional[str] = None,19    ) -> None:20        self.digits = digits21        if digits > 10:22            raise ValueError("digits must be no greater than 10")23        self.digest = digest24        self.secret = s25        self.name = name or "Secret"26        self.issuer = issuer27 28    def generate_otp(self, input: int) -> str:29        """30        :param input: the HMAC counter value to use as the OTP input.31            Usually either the counter, or the computed integer based on the Unix timestamp32        """33        if input < 0:34            raise ValueError("input must be positive integer")35        hasher = hmac.new(self.byte_secret(), self.int_to_bytestring(input), self.digest)36        hmac_hash = bytearray(hasher.digest())37        offset = hmac_hash[-1] & 0xF38        code = (39            (hmac_hash[offset] & 0x7F) << 2440            | (hmac_hash[offset + 1] & 0xFF) << 1641            | (hmac_hash[offset + 2] & 0xFF) << 842            | (hmac_hash[offset + 3] & 0xFF)43        )44        str_code = str(10_000_000_000 + (code % 10**self.digits))45        return str_code[-self.digits :]46 47    def byte_secret(self) -> bytes:48        secret = self.secret49        missing_padding = len(secret) % 850        if missing_padding != 0:51            secret += "=" * (8 - missing_padding)52        return base64.b32decode(secret, casefold=True)53 54    @staticmethod55    def int_to_bytestring(i: int, padding: int = 8) -> bytes:56        """57        Turns an integer to the OATH specified58        bytestring, which is fed to the HMAC59        along with the secret60        """61        result = bytearray()62        while i != 0:63            result.append(i & 0xFF)64            i >>= 865        # It's necessary to convert the final result from bytearray to bytes66        # because the hmac functions in python 2.6 and 3.3 don't work with67        # bytearray68        return bytes(bytearray(reversed(result)).rjust(padding, b"\0"))69