EnigmaMachine/main.py

import sys
import time

alphaBase = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz,. !="

def main():
    if len(sys.argv) < 6:
        print("Usage: python script.py <arg1> <arg2>")

    if sys.argv[5] == "e":
        # first iteration
        firstIterationResult = encode(int(sys.argv[1]), sys.argv[4])
        print(firstIterationResult)
        
        secondIterationResult = encode(int(sys.argv[2]), sys.argv[4])
        print(secondIterationResult)
        
        thirdIterationResult = encode(int(sys.argv[3]), sys.argv[4])
        print(thirdIterationResult)
    else:
        firstIterationResult = decode(int(sys.argv[1]), sys.argv[4])
        print(firstIterationResult)

        secondIterationResult = decode(int(sys.argv[2]), sys.argv[4])
        print(secondIterationResult)

        thirdIterationResult = decode(int(sys.argv[3]), sys.argv[4])
        print(thirdIterationResult)


def encode(offset, message):
    print("Encoding!")
    current_timestamp = time.time()
    print(current_timestamp)

    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


if __name__ == "__main__":
    main()