From 12c5e38691b24d12a8fc9605e3b93fac86d461ae Mon Sep 17 00:00:00 2001 From: Simeon Simeonov Date: Mon, 21 Mar 2022 13:25:26 +0100 Subject: Reformat the code with annotations and new code style --- otp2289/generator.py | 142 +++++++++++++++++++++++++++++---------------------- 1 file changed, 82 insertions(+), 60 deletions(-) (limited to 'otp2289/generator.py') diff --git a/otp2289/generator.py b/otp2289/generator.py index 5ea738b..7b86e1b 100644 --- a/otp2289/generator.py +++ b/otp2289/generator.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- # SPDX-License-Identifier: BSD-2-Clause-FreeBSD # -# Copyright (c) 2020, Simeon Simeonov +# Copyright (c) 2020-2022 Simeon Simeonov # All rights reserved. # # Redistribution and use in source and binary forms, with or without @@ -28,7 +28,6 @@ import binascii import hashlib import string - OTP_ALGO_MD5 = 1 OTP_ALGO_SHA1 = 2 @@ -306,7 +305,12 @@ class OTPChallengeException(Exception): class OTPGenerator: """OTPGenerator class""" - def __init__(self, password, seed='', hash_algo=OTP_ALGO_MD5): + def __init__( + self, + password: bytes, + seed: str = '', + hash_algo=OTP_ALGO_MD5, + ): """ Constructs an OTPGenerator object with a given password and seed. @@ -330,16 +334,19 @@ class OTPGenerator: raise OTPGeneratorException('Password must be a byte-string') if len(password) < 10: raise OTPGeneratorException( - 'Password must be longer than 10 bytes') + 'Password must be longer than 10 bytes' + ) self._password = password def __repr__(self): """repr implementation""" - return (f'{self.__class__} at {id(self)} (seed={self._seed}, ' - f'hash_algo={self._hash_algo})') + return ( + f'{self.__class__} at {id(self)} (seed={self._seed}, ' + f'hash_algo={self._hash_algo})' + ) @staticmethod - def bit_pair_sum(bit_stream): + def bit_pair_sum(bit_stream: str) -> int: """ Split bit_stream in bit-pairs and sum them all together. @@ -359,7 +366,7 @@ class OTPGenerator: return value @staticmethod - def bytes_to_tokens(hash_bytes): + def bytes_to_tokens(hash_bytes: bytes) -> str: """ Returns a 6 words token from bytes as specified by RFC-2289. @@ -369,8 +376,7 @@ class OTPGenerator: :return: 6 words tokens :rtype: str """ - bit_stream = ''.join( - ['{0:0>8b}'.format(byte) for byte in hash_bytes]) + bit_stream = ''.join([f'{byte:0>8b}' for byte in hash_bytes]) bit_pair_sum = OTPGenerator.bit_pair_sum(bit_stream) tokens = [] tokens.append(RFC1760_TOKENS[int(bit_stream[:11], 2)]) @@ -379,12 +385,17 @@ class OTPGenerator: tokens.append(RFC1760_TOKENS[int(bit_stream[33:44], 2)]) tokens.append(RFC1760_TOKENS[int(bit_stream[44:55], 2)]) tokens.append( - RFC1760_TOKENS[int( - bit_stream[55:64] + '{0:0>8b}'.format(bit_pair_sum)[-2:], 2)]) + RFC1760_TOKENS[ + int( + bit_stream[55:64] + f'{bit_pair_sum:0>8b}'[-2:], + 2, + ) + ] + ) return ' '.join(tokens) @staticmethod - def get_tokens_from_challenge(challenge): + def get_tokens_from_challenge(challenge: str) -> tuple: """ Returns tokens (seed, hash_algo and step) from a challenge string. @@ -410,7 +421,7 @@ class OTPGenerator: raise OTPChallengeException('Invalid challenge') from None @staticmethod - def sha1_digest_folding(sha1_digest): + def sha1_digest_folding(sha1_digest: bytes) -> bytes: """ Implementation of the 160bit -> 64bit folding algorithm for sha1 digest. @@ -425,14 +436,17 @@ class OTPGenerator: raise OTPGeneratorException('sha1_digest must be of type bytes') if len(sha1_digest) != 20: raise OTPGeneratorException( - 'sha1_digest must be 160 bits (20 bytes) long') + 'sha1_digest must be 160 bits (20 bytes) long' + ) digested = list(5 * b'i') # 5 bytes (40 bits) result = list(8 * b'x') # 8 bytes (64 bits) for i in range(5): - digested[i] = (((sha1_digest[i * 4 + 0] & 0xff) << 24) | - ((sha1_digest[i * 4 + 1] & 0xff) << 16) | - ((sha1_digest[i * 4 + 2] & 0xff) << 8) | - (sha1_digest[i * 4 + 3] & 0xff)) + digested[i] = ( + ((sha1_digest[i * 4 + 0] & 0xFF) << 24) + | ((sha1_digest[i * 4 + 1] & 0xFF) << 16) + | ((sha1_digest[i * 4 + 2] & 0xFF) << 8) + | (sha1_digest[i * 4 + 3] & 0xFF) + ) # sha.digest[0] ^= sha.digest[2]; # sha.digest[1] ^= sha.digest[3]; # sha.digest[0] ^= sha.digest[4]; @@ -446,18 +460,18 @@ class OTPGenerator: # result[j+3] = (unsigned char)((sha.digest[i] >> 24) & 0xff); # } # just hardcoding the two iterations for better efficiency - result[0] = digested[0] & 0xff - result[1] = (digested[0] >> 8) & 0xff - result[2] = (digested[0] >> 16) & 0xff - result[3] = (digested[0] >> 24) & 0xff - result[4] = digested[1] & 0xff - result[5] = (digested[1] >> 8) & 0xff - result[6] = (digested[1] >> 16) & 0xff - result[7] = (digested[1] >> 24) & 0xff + result[0] = digested[0] & 0xFF + result[1] = (digested[0] >> 8) & 0xFF + result[2] = (digested[0] >> 16) & 0xFF + result[3] = (digested[0] >> 24) & 0xFF + result[4] = digested[1] & 0xFF + result[5] = (digested[1] >> 8) & 0xFF + result[6] = (digested[1] >> 16) & 0xFF + result[7] = (digested[1] >> 24) & 0xFF return bytes(result) @staticmethod - def strxor(byte_str1, byte_str2): + def strxor(byte_str1: bytes, byte_str2: bytes) -> bytes: """ Implementation of strxor similar to the one provided by pycrypto. @@ -472,17 +486,17 @@ class OTPGenerator: """ if not (isinstance(byte_str1, bytes) and isinstance(byte_str2, bytes)): raise OTPGeneratorException( - 'byte_str1 and byte_str2 must be of type bytes') + 'byte_str1 and byte_str2 must be of type bytes' + ) length = len(byte_str1) if length != len(byte_str2) or length < 1: raise OTPGeneratorException( - 'byte_str1 and byte_str2 must be of the same length > 0') - return bytes( - [byte_str1[i] ^ byte_str2[i] for i in range(length)] - ) + 'byte_str1 and byte_str2 must be of the same length > 0' + ) + return bytes([byte_str1[i] ^ byte_str2[i] for i in range(length)]) @staticmethod - def tokens_to_bytes(tokens_str): + def tokens_to_bytes(tokens_str: str) -> bytes: """ Returns bytes from a 6 words token as specified by RFC-2289. @@ -499,14 +513,17 @@ class OTPGenerator: tokens = tokens_str.split() if len(tokens) != 6: raise OTPGeneratorException( - 'Tokens-string does not contain 6 tokens') + 'Tokens-string does not contain 6 tokens' + ) token_ints = [] try: - token_ints = [RFC1760_TOKENS.index(token.upper()) for token in - tokens] + token_ints = [ + RFC1760_TOKENS.index(token.upper()) for token in tokens + ] except ValueError: raise OTPGeneratorException( - 'One or more words not present in RFC1760') from None + 'One or more words not present in RFC1760' + ) from None # now we build a string of bits bit_stream = format(token_ints[0], '011b') bit_stream += format(token_ints[1], '011b') @@ -519,14 +536,14 @@ class OTPGenerator: # OTP servers MUST verify this checksum explicitly as part of the # operation of decoding this representation of the one-time password. if ( - '{0:0>8b}'.format(OTPGenerator.bit_pair_sum( - bit_stream[:64]))[-2:] != bit_stream[-2:] + f'{OTPGenerator.bit_pair_sum(bit_stream[:64]):0>8b}'[-2:] + != bit_stream[-2:] ): raise OTPGeneratorException('Invalid bit checksum') return int(bit_stream[:64], 2).to_bytes(8, 'big') @staticmethod - def validate_hash_algo(hash_algo): + def validate_hash_algo(hash_algo) -> str: """ Validates the provided hash-algorithm. @@ -541,19 +558,20 @@ class OTPGenerator: if isinstance(hash_algo, int): if hash_algo not in _ALGO_DICT: raise OTPGeneratorException( - 'hash_algo is not among the known algorithms') + 'hash_algo is not among the known algorithms' + ) hash_algo = _ALGO_DICT.get(hash_algo) if not isinstance(hash_algo, str): - raise OTPGeneratorException( - 'hash_algo must be an int or a str') + raise OTPGeneratorException('hash_algo must be an int or a str') if hash_algo not in hashlib.algorithms_available: raise OTPGeneratorException( f'{hash_algo} is not supported by this version of the ' - 'hashlib module') + 'hashlib module' + ) return hash_algo @staticmethod - def validate_seed(seed): + def validate_seed(seed: str) -> str: """ Validates the provided seed as defined by RFC-2289. @@ -569,15 +587,17 @@ class OTPGenerator: raise OTPGeneratorException('Seed must be a string') if not seed or len(seed) > 16: raise OTPGeneratorException( - 'The seed MUST be of 1 to 16 characters in length') + 'The seed MUST be of 1 to 16 characters in length' + ) for char in seed: if char not in string.ascii_letters + string.digits: raise OTPGeneratorException( - 'The seed MUST consist of purely alphanumeric characters') + 'The seed MUST consist of purely alphanumeric characters' + ) return seed @staticmethod - def validate_step(step): + def validate_step(step: int) -> int: """ Validates the provided step as defined by RFC-2289. @@ -595,7 +615,7 @@ class OTPGenerator: raise OTPGeneratorException('Step value MUST be >= 0') return step - def generate_otp_hexdigest(self, step): + def generate_otp_hexdigest(self, step: int) -> str: """ Generates the OTP hexdigest for the given step. @@ -607,7 +627,7 @@ class OTPGenerator: """ return '0x' + binascii.hexlify(self._generate_otp_bytes(step)).decode() - def generate_otp_hexdigest_from_challenge(self, challenge): + def generate_otp_hexdigest_from_challenge(self, challenge: str) -> str: """ Same as generate_otp_hexdigest, but it generates hex. from a challenge. @@ -628,7 +648,7 @@ class OTPGenerator: self._hash_algo = self.validate_hash_algo(hash_algo) return self.generate_otp_hexdigest(step) - def generate_otp_words(self, step): + def generate_otp_words(self, step: int) -> str: """ Generates the OTP six words token for the given step. @@ -640,7 +660,7 @@ class OTPGenerator: """ return self.bytes_to_tokens(self._generate_otp_bytes(step)) - def generate_otp_words_from_challenge(self, challenge): + def generate_otp_words_from_challenge(self, challenge: str) -> str: """ Same as generate_otp_words, but it generates words from a challenge. @@ -661,7 +681,7 @@ class OTPGenerator: self._hash_algo = self.validate_hash_algo(hash_algo) return self.generate_otp_words(step) - def hexdigest_range(self, start=499, stop=0): + def hexdigest_range(self, start: int = 499, stop: int = 0): """ Returns an iterator that providing hexdigests corresponding to steps from `start` to and including `stop`. @@ -679,11 +699,12 @@ class OTPGenerator: raise OTPGeneratorException('Step value MUST be an int') if start < stop: raise OTPGeneratorException( - 'Start value can not be lower than stop') + 'Start value can not be lower than stop' + ) for step in range(start, stop - 1, -1): yield self.generate_otp_hexdigest(step) - def words_range(self, start=499, stop=0): + def words_range(self, start: int = 499, stop: int = 0): """ Returns an iterator that providing the words corresponding to steps from `start` to and including `stop`. @@ -701,11 +722,12 @@ class OTPGenerator: raise OTPGeneratorException('Step value MUST be an int') if start < stop: raise OTPGeneratorException( - 'Start value can not be lower than stop') + 'Start value can not be lower than stop' + ) for step in range(start, stop - 1, -1): yield self.generate_otp_words(step) - def _generate_otp_bytes(self, step): + def _generate_otp_bytes(self, step: int) -> bytes: """ Generates the OTP bytes for the given step. @@ -733,6 +755,6 @@ class OTPGenerator: digest = self.sha1_digest_folding(large_digest) else: raise OTPGeneratorException( - '{hash_algo} is not supported by this module'.format( - hash_algo=self._hash_algo)) + f'{self._hash_algo} is not supported by this module' + ) return digest -- cgit v1.3