feat: add type lint

This commit is contained in:
sangge 2024-01-08 05:39:37 +08:00
parent 72cdedfcdb
commit b4d68f1cba

View File

@ -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)