EnigmaMachine/main.py

import random
import sys
import time

# define the full range of characters used by this program
alpha_base: str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz,. !="


# MAIN
def main():
    ### Argument validation ###
    if len(sys.argv) < 4:
        print("Usage: uv run main.py \"[message]\" e|d [int]")
        sys.exit(1)

    base_max: int = len(alpha_base)

    if not is_number(sys.argv[3]):
        print("Invalid iterations value.")
        sys.exit(1)

    if int(sys.argv[3]) == 0:
        iterations = 3
    else:
        iterations = int(sys.argv[3])

    ### Processing ###
    seed_values: list[int] = []

    if sys.argv[2] == "e":
        with open("seeds.txt", "w") as f:
            for iteration in range(0, iterations):
                random.seed(time.time())
                seed_values.append(random.randint(1, base_max))
                f.write(str(seed_values[iteration]) + "\n")
        f.close()

    if sys.argv[2] == "d":
        with open("seeds.txt", "r") as f:
            for iteration in range(0, iterations):
                seed = f.readline()
                seed_values.append(int(seed))
        f.close()

    iteration_result: str = ""

    if sys.argv[2] == "e":
        print("Encoding!")
        for iteration in range(0, iterations):
            iteration_result = encode(seed_values[iteration], sys.argv[1])

        print("Result:", "\"" + iteration_result + "\"")
    elif sys.argv[2] == "d":
        print("Decoding!")
        for iteration in range(0, iterations):
            iteration_result = decode(seed_values[iteration], sys.argv[1])

        print("Result:", "\"" + iteration_result + "\"")
    else:
        print("Invalid option!")
        sys.exit(1)


def encode(offset: int, message: str) -> str:
    result_message: str = ""

    for index_a in range(len(message)):
        for index_b in range(len(alpha_base)):
            if message[index_a] == alpha_base[index_b]:
                location = index_b + offset
                if location > (len(alpha_base) - 1):
                    location = location - (len(alpha_base) - 1)
                result_message = result_message + alpha_base[location]
                break

    return result_message


def decode(offset: int, message: str) -> str:
    result_message: str = ""

    for index_a in range(len(message)):
        for index_b in range(len(alpha_base)):
            if message[index_a] == alpha_base[index_b]:
                location = index_b - offset
                if location < 0:
                    location = location + (len(alpha_base) - 1)
                result_message = result_message + alpha_base[location]
                break

    return result_message


def is_number(i) -> bool:
    try:
        int(i)
        return True
    except ValueError:
        return False


if __name__ == "__main__":
    main()