diff --git a/aesDrbg.py b/aesDrbg.py index 9fe0d41..b6bf9f1 100755 --- a/aesDrbg.py +++ b/aesDrbg.py @@ -26,13 +26,13 @@ class AES_CTR_DRBG: # 最大的熵输入长度 self.max_ectropy_input_length = 1 << 35 - 1 # 2^35比特 - self.seed_material = "" + self.seed_material = b"" self.aes = AES.new(b"\x00" * self.keylen, AES.MODE_ECB) self.AES_CTR_DRBG_Instantiate(personalization_string, nonce) def AES_CTR_DRBG_Instantiate( self, personalization_string: bytes = b"", nonce: bytes = b"" - ): + )-> None: self.min_entropy = self.seedlen self.entropy_input = secrets.token_bytes(self.min_entropy) self.seed_material = self.entropy_input + nonce + personalization_string @@ -42,11 +42,11 @@ class AES_CTR_DRBG: self.reseed_counter = 1 self.last_reseed_time = int(time.time()) - def AES_CTR_DRBG_Update(self, seed_material, Key, V): + def AES_CTR_DRBG_Update(self, seed_material:bytes, Key:bytes, V: bytes)->None: temp = b"" while len(temp) < len(seed_material): - self.V = (int.from_bytes(self.V, "big") + 1) % (1 << (8 * self.blocklen)) - self.V = self.V.to_bytes(self.blocklen, "big") + V_int = (int.from_bytes(self.V, "big") + 1) % (1 << (8 * self.blocklen)) + self.V = V_int.to_bytes(self.blocklen, "big") temp += self.aes.encrypt(self.V) temp = temp[: len(seed_material)] temp = int.from_bytes(temp, "big") ^ int.from_bytes(seed_material, "big") @@ -78,7 +78,7 @@ class AES_CTR_DRBG: requested_bits = tmp[:number_of_bits_to_return] return requested_bits - def CBC_MAC(self, Key, data_to_MAC): + def CBC_MAC(self, Key: bytes, data_to_MAC:bytes)->bytes: chaining_value = b"\x00" * self.blocklen for i in range(0, len(data_to_MAC), self.blocklen): block = data_to_MAC[i : i + self.blocklen] @@ -90,7 +90,7 @@ class AES_CTR_DRBG: ) return chaining_value - def AES_CTR_DRBG_Reseed(self, additional_input: bytes): + def AES_CTR_DRBG_Reseed(self, additional_input: bytes)-> None: self.min_entropy = self.seedlen self.entropy_input = secrets.token_bytes(self.min_entropy) self.seed_material = self.entropy_input + additional_input @@ -100,8 +100,8 @@ class AES_CTR_DRBG: self.last_reseed_time = int(time.time()) def AES_CTR_DRBG_Generate( - self, requested_number_of_bits, additional_input: bytes = b"" - ): + self, requested_number_of_bits: int, additional_input: bytes = b"" + )-> bytes: length = requested_number_of_bits // 8 returned_bits = b"" if ( @@ -115,8 +115,8 @@ class AES_CTR_DRBG: else: additional_input = b"\x00" * self.seedlen while len(returned_bits) < length: - self.V = (int.from_bytes(self.V, "big") + 1) % (1 << (8 * self.blocklen)) - self.V = self.V.to_bytes(self.blocklen, "big") + V_int = (int.from_bytes(self.V, "big") + 1) % (1 << (8 * self.blocklen)) + self.V = V_int.to_bytes(self.blocklen, "big") output_block = self.aes.encrypt(self.V) returned_bits += output_block self.AES_CTR_DRBG_Update(additional_input, self.Key, self.V)