Enigma Machine

.py

'''


NOTE: Code is run in part1_test.py


'''

from copy import deepcopy
STURGEON_ALPHABET = "2T3O4HNM5LRGIPCVEZDBSYFXAWJ6UQK7"

'''
ORDER   = order to wheels
OFFSET  = offset of each wheel
WHEEL_SIZE  = number of bits on each wheel
'''
ORDER = [8, 7, 2, 4, 3, 5, 6, 1, 0, 9]
OFFSET = [44, 52, 35, 14, 19, 55, 6, 4, 3, 51]
WHEEL_SIZE = [47,53,59,61,64,65,67,69,71,73]

def getFile(fname):
    with open(fname) as f:
        line = f.readlines()
        return [l.replace('\n','') for l in line]

rotorsBits = getFile("./part_1/rotors.txt")
cipherText = getFile("./part_1/ciphertext.txt")
plainText = getFile("./part_1/plaintext.txt")

'''
Converts sturgeon alphabet into binary
'''
def tobits(c):
    results=[]
    alphabestPos = STURGEON_ALPHABET.index(c)
    bits = '{0:05b}'.format(alphabestPos)
    results.extend([int(b) for b in bits])
    return results

'''
Converts binary into sturgeon alphabet
@return String: return a letter from the STURGEON_ALPHABET
'''
def frombits(bits):
    # print(bits)
    if 2 in bits or 3 in bits:
        # print(bits)
        return "-"
    else:
        # print("K", bits)
        alphabetPos = 16*bits[0] + 8*bits[1] + 4*bits[2] + 2*bits[3] + bits[4]
        return STURGEON_ALPHABET[alphabetPos]

'''
Gets Rotor bits
@returns [] of length 10, with current rotor bit
'''
def getRotorsBits(order, offset, rotorsBits):
    bits = []
    for i in range(len(order)):
        rotor = rotorsBits[order[i]]
        if int(offset[i]) <= len(rotor):
            bits.append(rotor[offset[i]])
        else:
            raise ValueError('offset bigger then rotor length')
            break
    return [int(b) for b in bits]

'''
Increments wheel bits by 1
@returns [] of length 10, with new offset
'''
def incrementRotorsBit(order, offset):
    newOffset = []
    for i, v in enumerate(order):
        newBit = (offset[i] + 1) % WHEEL_SIZE[v]
        newOffset.append(newBit)
    return newOffset

'''
XOR bits with first 5 rotor rotor bits
@return [] length 5
'''
def xorText(textBits, rotorBitsFirst):
    return [(textBits[i]) ^ (rotorBitsFirst[i]) for i in range(5)]

'''
Utility function to switch bits
'''
def swapBits(a ,b, xoredTextBits):
    temp = xoredTextBits[a]
    xoredTextBits[a] = xoredTextBits[b]
    xoredTextBits[b] = temp
    return

def switchBits(xoredTextBits, rotorBitsSecond):
    if rotorBitsSecond[0] == 1: swapBits(0, 4, xoredTextBits)
    if rotorBitsSecond[1] == 1: swapBits(0, 1, xoredTextBits)
    if rotorBitsSecond[2] == 1: swapBits(1, 2, xoredTextBits)
    if rotorBitsSecond[3] == 1: swapBits(2, 3, xoredTextBits)
    if rotorBitsSecond[4] == 1: swapBits(3, 4, xoredTextBits)
    return xoredTextBits

# Encrypt each line
def encrypt(plainText, order, offset, rotorsBits=rotorsBits):
    encryptedBit = []
    for p in plainText:
        # 1. Convert text to bits based on STURGEON_ALPHABET
        textBits = tobits(p)

        # 2. Get the 10 rotor bits based on order of wheel and offset
        rotorBits = getRotorsBits(order, offset, rotorsBits)
        # 3. XOR bits with the first 5 bits in the rotor
        xoredTextBits = xorText(textBits, rotorBits[:5])
        # 4. Switch bits based on the last 5 bits of the rotor
        encryptedBit.append(frombits(switchBits(xoredTextBits, rotorBits[5:])))

        # 5. increment rotors by 1 bit rotiation
        offset = incrementRotorsBit(order, offset)

    encryptedText = ''.join(encryptedBit)
    return encryptedText, offset


# def decrementRotorsBit(order, offset):
#     newOffset = []
#     for i, v in enumerate(order):
#         newBit = (offset[i] - 1) % WHEEL_SIZE[v]
#         newOffset.append(newBit)
#     return newOffset


'''
Switch bits in reverse order from the instructions
@return [] length 5
'''
def switchBitsReverse(xoredTextBits, rotorBitsSecond):
    if rotorBitsSecond[4] == 1: swapBits(3, 4, xoredTextBits)
    if rotorBitsSecond[3] == 1: swapBits(2, 3, xoredTextBits)
    if rotorBitsSecond[2] == 1: swapBits(1, 2, xoredTextBits)
    if rotorBitsSecond[1] == 1: swapBits(0, 1, xoredTextBits)
    if rotorBitsSecond[0] == 1: swapBits(0, 4, xoredTextBits)
    return xoredTextBits

'''
Decryption for each line of cipher text
'''
def decrypt(cipherText, order, offset, rotorsBits=rotorsBits):
    decryptedBit = []
    for p in cipherText:
        rotorBits = getRotorsBits(order, offset, rotorsBits)

        textBits = tobits(p)
        reversedTextBits = switchBitsReverse(textBits, rotorBits[5:])
        decryptedBit.append(frombits(xorText(reversedTextBits, rotorBits[:5])))

        offset = incrementRotorsBit(order, offset)
        # print(offset)

    decryptedText = ''.join(decryptedBit)
    return decryptedText, offset