EnigmaMachine/main.py

import random
import sys
import time

alphaBase = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz,. !="

def main():
    ### Argument validation ###
    if len(sys.argv) < 7:
        print("Usage: uv run main.py [0-58] [0-58] [0-58] '[message]' 'e|d' [int]")
        sys.exit(1)

    base_max = len(alphaBase)

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

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

    if int(sys.argv[1]) == 0 or int(sys.argv[1]) > base_max:
        seed1 = random.randint(1, base_max)
    else:
        seed1 = int(sys.argv[1])

    if int(sys.argv[2]) == 0 or int(sys.argv[2]) > base_max:
        seed2 = random.randint(1, base_max)
    else:
        seed2 = int(sys.argv[2])

    if int(sys.argv[3]) == 0 or int(sys.argv[3]) > base_max:
        seed3 = random.randint(1, base_max)
    else:
        seed3 = int(sys.argv[3])

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

    ### Processing ###
    # current_timestamp = time.time()
    # print(current_timestamp)

    seed_values = []

    if sys.argv[5] == "e":
        for iteration in range(0, iterations):
            random.seed(time.time())
            seed_values.append(random.randint(1, base_max))

    if sys.argv[5] == "d":
        seed_values.append(seed1)
        seed_values.append(seed2)
        seed_values.append(seed3)

    iteration_result = ""

    if sys.argv[5] == "e":
        for iteration in range(0, iterations):
            iteration_result = encode(seed_values[iteration], sys.argv[4])
            print("Seed:", iteration, "is", seed_values[iteration])

        print("Result:", iteration_result)
    elif sys.argv[5] == "d":
        for iteration in range(0, iterations):
            iteration_result = decode(seed_values[iteration], sys.argv[4])

        print("Seeds used:", seed1, seed2, seed3)
        print("Result:", iteration_result)
    else:
        print("Invalid option!")
        sys.exit(1)


def encode(offset, message):
    print("Encoding!")

    resultMessage = ""

    for indexA in range(len(message)):
        for indexB in range(len(alphaBase)):
            if message[indexA] == alphaBase[indexB]:
                location = indexB + offset
                if location > (len(alphaBase) - 1):
                    location = location - (len(alphaBase) - 1)
                resultMessage = resultMessage + alphaBase[location]
                break

    return resultMessage


def decode(offset, message):
    print("Decoding!")

    resultMessage = ""

    for indexA in range(len(message)):
        for indexB in range(len(alphaBase)):
            if message[indexA] == alphaBase[indexB]:
                location = indexB - offset
                if location < 0:
                    location = location + (len(alphaBase) - 1)
                resultMessage = resultMessage + alphaBase[location]
                break

    return resultMessage


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


if __name__ == "__main__":
    main()