fix(mypy): Fix annotations for 13 cipher algorithms (#4278)

* Initial fix for mypy errors in some cipher algorithms

* fix(mypy): Update type hints

* fix(mypy): Update type hints for enigma_machine2.py

* Update as per the suggestion

Co-authored-by: Christian Clauss <cclauss@me.com>

Co-authored-by: Christian Clauss <cclauss@me.com>

ciphers/a1z26.py

@@ -7,7 +7,7 @@ http://bestcodes.weebly.com/a1z26.html
 """
 
 
-def encode(plain: str) -> list:
+def encode(plain: str) -> list[int]:
     """
     >>> encode("myname")
     [13, 25, 14, 1, 13, 5]
@@ -15,7 +15,7 @@ def encode(plain: str) -> list:
     return [ord(elem) - 96 for elem in plain]
 
 
-def decode(encoded: list) -> str:
+def decode(encoded: list[int]) -> str:
     """
     >>> decode([13, 25, 14, 1, 13, 5])
     'myname'
@@ -23,8 +23,8 @@ def decode(encoded: list) -> str:
     return "".join(chr(elem + 96) for elem in encoded)
 
 
-def main():
+def main() -> None:
-    encoded = encode(input("->").strip().lower())
+    encoded = encode(input("-> ").strip().lower())
     print("Encoded: ", encoded)
     print("Decoded:", decode(encoded))
 

ciphers/affine_cipher.py

@@ -9,26 +9,6 @@ SYMBOLS = (
 )
 
 
-def main():
-    """
-    >>> key = get_random_key()
-    >>> msg = "This is a test!"
-    >>> decrypt_message(key, encrypt_message(key, msg)) == msg
-    True
-    """
-    message = input("Enter message: ").strip()
-    key = int(input("Enter key [2000 - 9000]: ").strip())
-    mode = input("Encrypt/Decrypt [E/D]: ").strip().lower()
-
-    if mode.startswith("e"):
-        mode = "encrypt"
-        translated = encrypt_message(key, message)
-    elif mode.startswith("d"):
-        mode = "decrypt"
-        translated = decrypt_message(key, message)
-    print(f"\n{mode.title()}ed text: \n{translated}")
-
-
 def check_keys(keyA: int, keyB: int, mode: str) -> None:
     if mode == "encrypt":
         if keyA == 1:
@@ -80,7 +60,7 @@ def decrypt_message(key: int, message: str) -> str:
     keyA, keyB = divmod(key, len(SYMBOLS))
     check_keys(keyA, keyB, "decrypt")
     plainText = ""
-    modInverseOfkeyA = cryptomath.findModInverse(keyA, len(SYMBOLS))
+    modInverseOfkeyA = cryptomath.find_mod_inverse(keyA, len(SYMBOLS))
     for symbol in message:
         if symbol in SYMBOLS:
             symIndex = SYMBOLS.find(symbol)
@@ -98,6 +78,26 @@ def get_random_key() -> int:
             return keyA * len(SYMBOLS) + keyB
 
 
+def main() -> None:
+    """
+    >>> key = get_random_key()
+    >>> msg = "This is a test!"
+    >>> decrypt_message(key, encrypt_message(key, msg)) == msg
+    True
+    """
+    message = input("Enter message: ").strip()
+    key = int(input("Enter key [2000 - 9000]: ").strip())
+    mode = input("Encrypt/Decrypt [E/D]: ").strip().lower()
+
+    if mode.startswith("e"):
+        mode = "encrypt"
+        translated = encrypt_message(key, message)
+    elif mode.startswith("d"):
+        mode = "decrypt"
+        translated = decrypt_message(key, message)
+    print(f"\n{mode.title()}ed text: \n{translated}")
+
+
 if __name__ == "__main__":
     import doctest
 

ciphers/atbash.py

@@ -61,6 +61,6 @@ def benchmark() -> None:
 
 
 if __name__ == "__main__":
-    for sequence in ("ABCDEFGH", "123GGjj", "testStringtest", "with space"):
+    for example in ("ABCDEFGH", "123GGjj", "testStringtest", "with space"):
-        print(f"{sequence} encrypted in atbash: {atbash(sequence)}")
+        print(f"{example} encrypted in atbash: {atbash(example)}")
     benchmark()

ciphers/base32.py

@@ -1,7 +1,7 @@
 import base64
 
 
-def main():
+def main() -> None:
     inp = input("->")
     encoded = inp.encode("utf-8")  # encoded the input (we need a bytes like object)
     b32encoded = base64.b32encode(encoded)  # b32encoded the encoded string

ciphers/base85.py

@@ -1,7 +1,7 @@
 import base64
 
 
-def main():
+def main() -> None:
     inp = input("->")
     encoded = inp.encode("utf-8")  # encoded the input (we need a bytes like object)
     a85encoded = base64.a85encode(encoded)  # a85encoded the encoded string

ciphers/beaufort_cipher.py

@@ -66,7 +66,7 @@ def original_text(cipher_text: str, key_new: str) -> str:
     return or_txt
 
 
-def main():
+def main() -> None:
     message = "THE GERMAN ATTACK"
     key = "SECRET"
     key_new = generate_key(message, key)

ciphers/brute_force_caesar_cipher.py

@@ -43,7 +43,7 @@ def decrypt(message: str) -> None:
         print(f"Decryption using Key #{key}: {translated}")
 
 
-def main():
+def main() -> None:
     message = input("Encrypted message: ")
     message = message.upper()
     decrypt(message)

ciphers/cryptomath_module.py

@@ -4,9 +4,9 @@ def gcd(a: int, b: int) -> int:
     return b
 
 
-def findModInverse(a: int, m: int) -> int:
+def find_mod_inverse(a: int, m: int) -> int:
     if gcd(a, m) != 1:
-        return None
+        raise ValueError(f"mod inverse of {a!r} and {m!r} does not exist")
     u1, u2, u3 = 1, 0, a
     v1, v2, v3 = 0, 1, m
     while v3 != 0:

ciphers/decrypt_caesar_with_chi_squared.py

@@ -1,14 +1,14 @@
 #!/usr/bin/env python3
 
-from typing import Tuple
+from typing import Optional
 
 
 def decrypt_caesar_with_chi_squared(
     ciphertext: str,
-    cipher_alphabet: str = None,
+    cipher_alphabet: Optional[list[str]] = None,
-    frequencies_dict: str = None,
+    frequencies_dict: Optional[dict[str, float]] = None,
     case_sensetive: bool = False,
-) -> Tuple[int, float, str]:
+) -> tuple[int, float, str]:
     """
     Basic Usage
     ===========
@@ -123,9 +123,9 @@ def decrypt_caesar_with_chi_squared(
     AttributeError: 'int' object has no attribute 'lower'
     """
     alphabet_letters = cipher_alphabet or [chr(i) for i in range(97, 123)]
-    frequencies_dict = frequencies_dict or {}
 
-    if frequencies_dict == {}:
+    # If the argument is None or the user provided an empty dictionary
+    if not frequencies_dict:
         # Frequencies of letters in the english language (how much they show up)
         frequencies = {
             "a": 0.08497,
@@ -163,7 +163,7 @@ def decrypt_caesar_with_chi_squared(
         ciphertext = ciphertext.lower()
 
     # Chi squared statistic values
-    chi_squared_statistic_values = {}
+    chi_squared_statistic_values: dict[int, tuple[float, str]] = {}
 
     # cycle through all of the shifts
     for shift in range(len(alphabet_letters)):
@@ -215,22 +215,22 @@ def decrypt_caesar_with_chi_squared(
                     chi_squared_statistic += chi_letter_value
 
         # Add the data to the chi_squared_statistic_values dictionary
-        chi_squared_statistic_values[shift] = [
+        chi_squared_statistic_values[shift] = (
             chi_squared_statistic,
             decrypted_with_shift,
-        ]
+        )
 
     # Get the most likely cipher by finding the cipher with the smallest chi squared
     # statistic
-    most_likely_cipher = min(
+    most_likely_cipher: int = min(
         chi_squared_statistic_values, key=chi_squared_statistic_values.get
-    )
+    )  # type: ignore # First argument to `min` is not optional
 
     # Get all the data from the most likely cipher (key, decoded message)
-    most_likely_cipher_chi_squared_value = chi_squared_statistic_values[
+    (
-        most_likely_cipher
+        most_likely_cipher_chi_squared_value,
-    ][0]
+        decoded_most_likely_cipher,
-    decoded_most_likely_cipher = chi_squared_statistic_values[most_likely_cipher][1]
+    ) = chi_squared_statistic_values[most_likely_cipher]
 
     # Return the data on the most likely shift
     return (

ciphers/diffie.py

@@ -1,4 +1,7 @@
-def find_primitive(n: int) -> int:
+from typing import Optional
+
+
+def find_primitive(n: int) -> Optional[int]:
     for r in range(1, n):
         li = []
         for x in range(n - 1):
@@ -8,18 +11,22 @@ def find_primitive(n: int) -> int:
             li.append(val)
         else:
             return r
+    return None
 
 
 if __name__ == "__main__":
     q = int(input("Enter a prime number q: "))
     a = find_primitive(q)
-    a_private = int(input("Enter private key of A: "))
+    if a is None:
-    a_public = pow(a, a_private, q)
+        print(f"Cannot find the primitive for the value: {a!r}")
-    b_private = int(input("Enter private key of B: "))
+    else:
-    b_public = pow(a, b_private, q)
+        a_private = int(input("Enter private key of A: "))
+        a_public = pow(a, a_private, q)
+        b_private = int(input("Enter private key of B: "))
+        b_public = pow(a, b_private, q)
 
-    a_secret = pow(b_public, a_private, q)
+        a_secret = pow(b_public, a_private, q)
-    b_secret = pow(a_public, b_private, q)
+        b_secret = pow(a_public, b_private, q)
 
-    print("The key value generated by A is: ", a_secret)
+        print("The key value generated by A is: ", a_secret)
-    print("The key value generated by B is: ", b_secret)
+        print("The key value generated by B is: ", b_secret)

ciphers/elgamal_key_generator.py

@@ -2,24 +2,18 @@ import os
 import random
 import sys
 
-from . import cryptomath_module as cryptoMath
+from . import cryptomath_module as cryptomath
-from . import rabin_miller as rabinMiller
+from . import rabin_miller
 
 min_primitive_root = 3
 
 
-def main():
-    print("Making key files...")
-    makeKeyFiles("elgamal", 2048)
-    print("Key files generation successful")
-
-
 # I have written my code naively same as definition of primitive root
 # however every time I run this program, memory exceeded...
 # so I used 4.80 Algorithm in
 # Handbook of Applied Cryptography(CRC Press, ISBN : 0-8493-8523-7, October 1996)
 # and it seems to run nicely!
-def primitiveRoot(p_val: int) -> int:
+def primitive_root(p_val: int) -> int:
     print("Generating primitive root of p")
     while True:
         g = random.randrange(3, p_val)
@@ -30,20 +24,20 @@ def primitiveRoot(p_val: int) -> int:
         return g
 
 
-def generateKey(keySize: int) -> ((int, int, int, int), (int, int)):
+def generate_key(key_size: int) -> tuple[tuple[int, int, int, int], tuple[int, int]]:
     print("Generating prime p...")
-    p = rabinMiller.generateLargePrime(keySize)  # select large prime number.
+    p = rabin_miller.generateLargePrime(key_size)  # select large prime number.
-    e_1 = primitiveRoot(p)  # one primitive root on modulo p.
+    e_1 = primitive_root(p)  # one primitive root on modulo p.
     d = random.randrange(3, p)  # private_key -> have to be greater than 2 for safety.
-    e_2 = cryptoMath.findModInverse(pow(e_1, d, p), p)
+    e_2 = cryptomath.find_mod_inverse(pow(e_1, d, p), p)
 
-    publicKey = (keySize, e_1, e_2, p)
+    public_key = (key_size, e_1, e_2, p)
-    privateKey = (keySize, d)
+    private_key = (key_size, d)
 
-    return publicKey, privateKey
+    return public_key, private_key
 
 
-def makeKeyFiles(name: str, keySize: int):
+def make_key_files(name: str, keySize: int) -> None:
     if os.path.exists("%s_pubkey.txt" % name) or os.path.exists(
         "%s_privkey.txt" % name
     ):
@@ -55,7 +49,7 @@ def makeKeyFiles(name: str, keySize: int):
         )
         sys.exit()
 
-    publicKey, privateKey = generateKey(keySize)
+    publicKey, privateKey = generate_key(keySize)
     print("\nWriting public key to file %s_pubkey.txt..." % name)
     with open("%s_pubkey.txt" % name, "w") as fo:
         fo.write(
@@ -67,5 +61,11 @@ def makeKeyFiles(name: str, keySize: int):
         fo.write("%d,%d" % (privateKey[0], privateKey[1]))
 
 
+def main() -> None:
+    print("Making key files...")
+    make_key_files("elgamal", 2048)
+    print("Key files generation successful")
+
+
 if __name__ == "__main__":
     main()

ciphers/enigma_machine2.py

@@ -15,6 +15,10 @@ Module includes:
 Created by TrapinchO
 """
 
+RotorPositionT = tuple[int, int, int]
+RotorSelectionT = tuple[str, str, str]
+
+
 # used alphabet --------------------------
 # from string.ascii_uppercase
 abc = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
@@ -63,7 +67,9 @@ rotor8 = "LFKIJODBEGAMQPXVUHYSTCZRWN"
 rotor9 = "KOAEGVDHXPQZMLFTYWJNBRCIUS"
 
 
-def _validator(rotpos: tuple, rotsel: tuple, pb: str) -> tuple:
+def _validator(
+    rotpos: RotorPositionT, rotsel: RotorSelectionT, pb: str
+) -> tuple[RotorPositionT, RotorSelectionT, dict[str, str]]:
     """
     Checks if the values can be used for the 'enigma' function
 
@@ -99,12 +105,12 @@ def _validator(rotpos: tuple, rotsel: tuple, pb: str) -> tuple:
         )
 
     # Validates string and returns dict
-    pb = _plugboard(pb)
+    pbdict = _plugboard(pb)
 
-    return rotpos, rotsel, pb
+    return rotpos, rotsel, pbdict
 
 
-def _plugboard(pbstring: str) -> dict:
+def _plugboard(pbstring: str) -> dict[str, str]:
     """
     https://en.wikipedia.org/wiki/Enigma_machine#Plugboard
 
@@ -145,17 +151,17 @@ def _plugboard(pbstring: str) -> dict:
 
     # Created the dictionary
     pb = {}
-    for i in range(0, len(pbstring) - 1, 2):
+    for j in range(0, len(pbstring) - 1, 2):
-        pb[pbstring[i]] = pbstring[i + 1]
+        pb[pbstring[j]] = pbstring[j + 1]
-        pb[pbstring[i + 1]] = pbstring[i]
+        pb[pbstring[j + 1]] = pbstring[j]
 
     return pb
 
 
 def enigma(
     text: str,
-    rotor_position: tuple,
+    rotor_position: RotorPositionT,
-    rotor_selection: tuple = (rotor1, rotor2, rotor3),
+    rotor_selection: RotorSelectionT = (rotor1, rotor2, rotor3),
     plugb: str = "",
 ) -> str:
     """

ciphers/rsa_key_generator.py

@@ -1,19 +1,18 @@
 import os
 import random
 import sys
-from typing import Tuple
 
 from . import cryptomath_module as cryptoMath
 from . import rabin_miller as rabinMiller
 
 
-def main():
+def main() -> None:
     print("Making key files...")
     makeKeyFiles("rsa", 1024)
     print("Key files generation successful.")
 
 
-def generateKey(keySize: int) -> Tuple[Tuple[int, int], Tuple[int, int]]:
+def generateKey(keySize: int) -> tuple[tuple[int, int], tuple[int, int]]:
     print("Generating prime p...")
     p = rabinMiller.generateLargePrime(keySize)
     print("Generating prime q...")
@@ -27,14 +26,14 @@ def generateKey(keySize: int) -> Tuple[Tuple[int, int], Tuple[int, int]]:
             break
 
     print("Calculating d that is mod inverse of e...")
-    d = cryptoMath.findModInverse(e, (p - 1) * (q - 1))
+    d = cryptoMath.find_mod_inverse(e, (p - 1) * (q - 1))
 
     publicKey = (n, e)
     privateKey = (n, d)
     return (publicKey, privateKey)
 
 
-def makeKeyFiles(name: int, keySize: int) -> None:
+def makeKeyFiles(name: str, keySize: int) -> None:
     if os.path.exists("%s_pubkey.txt" % (name)) or os.path.exists(
         "%s_privkey.txt" % (name)
     ):