First, a warm welcome to the python language! None of this is in any particular order.
argparse
There is a really nice argument parsing library available out of the box which makes parsing command line args much less clunky. I'd set up your parser like this:
# Maybe put this into cli.py
from argparse import ArgumentParser
def setup_parser():
parser = ArgumentParser(desc="Enigma written in python")
parser.add_argument("mode", type=str, choices=['enc', 'dec', 'keygen'], help="Encode, Decode, or generate key")
parser.add_argument("-i", "--input-file", dest="input_file", type=str, help="File to be encrypted/decrypted")
parser.add_argument("key_file", type=str, help="Key file to be generated or to be used to encode/decode")
return parser
parser = setup_parser()
Then your arguments can be used like:
args = parser.parse_args()
args.mode
'enc'
args.key_file
'mykey.rsa'
args.input_file
'to_be_encrypted.txt'
Variable Names
To go along with the argparse suggestion, variables should be given descriptive names. In run_main, you have names like arg<i>, which make reading and debugging the code a little tricky. Name them how you intend to use them:
def main(mode, key_file, input_file=None, num_rolls=0):
...
Variable names and function names are recommended to be snake_case, where classes are recommended to be PascalCase. The Style Guide (also known as PEP-8) can be found here. Note however, that it's a guide. You don't have to follow it at every step, but it will make your code look nicer most of the time.
Enums and type-checking
Enums do a nice job of handling type-checking for you. You can re-define your enum like this:
class Mode(Enum):
ENC = 'enc'
DEC = 'dec'
Then, not only will the argument parser catch an invalid enum entry, so will the class:
mode = Mode('enc')
mode is Mode.ENC
True
mode = Mode('dec')
mode is Mode.DEC
True
mode = Mode('bad')
TypeError:
'bad' is not a valid Mode
Creating rolls
There are a few ways to clean this code up. First, I think defining your slice points outside of the slice notation itself on the key is a bit easier to read:
~snip~
for roll_number in range(roll_count):
# For readability, I would just compute these values here. It
# makes the slices look better
start = roll_number * roll_key_length
middle = roll_number * roll_key_length + 256
end = roll_number * roll_key_length + 256 + transition_count
# now it's easier to see what the key slices are
roll = Roll(key[start:middle], key[middle:end]))
Second, you iterate over all of your rolls twice, when you don't need to. Just do the check during the iteration:
for roll_number in range(roll_count):
~snip~
roll.check_input(transition_count)
rolls.append(roll)
File Handling
It's better practice to open files using the context manager with, as it will automatically close the file handle for you when you are done using it, no matter what happens:
# go from this
fh = open(myfile, 'rb')
content = fh.read()
fh.close()
# to this
with open(myfile, 'rb') as fh:
content = fh.read()
# and to open multiple files:
with open('file1.txt') as fh1, open('file2.txt') as fh2, ..., open('fileN.txt') as fhN:
# do things
There is also a spot in your code where you remove a file, do some calculations, and then create it again by writing to it. The w and wb modes will truncate an existing file for you. So go from:
if os.path.exists(self._keyFilename):
os.remove(file)
# skipping code
file = open(self._keyFilename, 'wb')
file.write(key)
file.close()
Do this instead:
# get rid of the `if file exists` check
# skipping code
with open(self.key_file_name, 'wb') as fh:
fh.write(key)
Filling a bytearray
You have a loop where you fill a bytearray with numbers:
for i in range(roll_count):
transform = bytearray(256)
for j in range(256):
transform[j] = j
You can skip the other loop and simply have bytearray consume the range:
for i in range(roll_count):
transform = bytearray(range(256))
Swapping variables
You can use tuple-style assignment to quickly and easily swap variable values:
# go from this
temp = transform[rand1]
transform[rand1] = transform[rand2]
transform[rand2] = temp
# to this
transform[rand1], transform[rand2] = transform[rand2], transform[rand1]
Generating a list of pseudo-random numbers (no dupes)
You have the following code to do this:
transitions = bytearray()
while len(transitions) < self._transitionCount:
rand = randint(0, 255)
if not transitions.count(rand):
transitions.append(rand)
This can be simplified down to
transitions = bytearray(
random.sample(range(255), transition_count)
)
This is much, much faster. transitions.count(x) will be O(N), where N is growing for each iteration. Not ideal.
is_twisted
While 0 and 1 can be interpreted as False and True, respectively through their truthiness, I think it would be better to actually return a boolean here. Also, this could be more succinctly expressed as an any statement:
# go from this
def is_twisted(self, trans):
for i in range(256):
if trans[i] == i:
return 0
return 1
# to this
def is_twisted(self, trans):
return any(item == i for i, item in enumerate(trans))
Where enumerate will produce index, item pairs for each item in an iterable/sequence/generator. Last, you don't actually use self here, so is_twisted could be a staticmethod:
@staticmethod
def is_twisted(trans):
return any(item == i for i, item in enumerate(trans))
for...else?
In Rolls.py, you have the following snippet:
for i in range(256):
found = 0
for j in self._transitions:
if self._transitions[j] == i:
found = 1
continue
if not found:
raise ValueError("Transitions not 1-1 complete")
This can be reduced down by using an else statement on a for loop. Yes, there is an else for for. You can think of it checking to see if a StopIteration has occurred:
for i in range(10):
if i == 5:
break
else:
print("Didn't exit early!")
for i in range(10:
if i == 11:
break # this won't happen
else:
print("Didn't exit early!")
Didn't exit early!
So your code snippet can reduce down to:
for i in range(256):
for j in self._transitions:
if self._transitions[j] == i:
break
else:
# is not found
raise ValueError("Transitions not 1-1 complete")
But even then, this loop could be simplified further. You can use set comparisons to be even quicker:
if not set(range(256)).symmetric_difference(self.transitions):
# there are leftovers!
raise ValueError("Transitions not 1-1 complete")
Where set.symmetric_difference will return a set containing all of the values exclusively contained in either set.
Checking Corresponding Elements
Any time you have code comparing corresponding elements of a collection, you can probably use zip:
for i in range(len(self._turn_over_indices) - 1):
if self._turn_over_indices[i] == self._turn_over_indices[i + 1]:
raise ValueError("turn_over_indices has doubles")
Can be turned into:
for left, right in zip(
self._turn_over_indices[:-1],
self._turn_over_indices[1:]
):
if left == right:
raise ValueError("doubles!")
This could even be cleaned up in your unit tests:
def init_business_logic(self, transitions, turnovers):
rolls_encrypt = []
rolls_decrypt = []
for index in range(len(transitions)):
rolls_encrypt.append(Roll(transitions[index], turnovers[index]))
rolls_decrypt.append(Roll(transitions[index], turnovers[index]))
# can now be
def init_business_logic(self, transitions, turnovers):
rolls_encrypt, rolls_decrypt = [], []
for transition, turnover in zip(transitions, turnovers):
rolls_encrypt.append(Roll(transition, turnover))
rolls_decrypt.append(Roll(transition, turnover))
transform_buffer
This could be refactored to be more DRY:
def transform_buffer(self, buffer, mode):
if mode == Mode.ENC:
# almost all of this code is repeated except roll.encrypt/decrypt
for i in range(len(buffer)):
for roll in self._rolls:
roll.encrypt(buffer, i)
self.roll_on()
if mode == Mode.DEC:
for i in range(len(buffer)):
for roll in self._rolls_reverse:
roll.decrypt(buffer, i)
self.roll_on()
The only difference is roll.encrypt, roll.decrypt, and the direction in which we iterate over self._rolls:
def transform_buffer(self, buffer, mode):
if mode is Mode.ENC:
rolls = self._rolls
func_name = 'encrypt'
else:
rolls = reversed(self._rolls)
func_name = 'decrypt'
for i in range(len(buffer)):
for roll in rolls:
f = getattr(roll, func_name)
f(buffer, i)
self.roll_on()
Speaking of which...
Reversing iterators
Instead of copying an object and reversing its order, you can just do:
for item in reversed(collection):
print(item)
So instead of doing:
class Rolls:
def __init__(self, rolls):
self._rolls = rolls
self._rolls_reverse = rolls.copy()
self._rolls_reverse.reverse()
Just keep one copy of self._rolls around.
class Rolls:
def __init__(self, rolls):
self._rolls = rolls
# no reverse copy
Rolls.roll_on?
I think this method could be named better, since you have two implementations of the roll_on function for two different classes. However, Rolls.roll_on has a different problem. It's using bytearray.count to effectively do a membership test:
def roll_on(self):
self._position = (self._position + 1) & 0xff
return self._turn_over_indices.count(self._position)
Now, normally with a membership test, I'd recommend a dict or set, but for this example, I'd just use in, since you are either creating a whole new object in memory that you need to keep in sync with the bytearray or sacrificing time to check for membership in a 256 element array. I'd choose the latter. However, bytearray.count (or any collection.count, for that matter) will always iterate through the whole thing. in will short-circuit:
def roll_on(self):
self._position = (self._position + 1) & 0xff
return self._position in self._turn_over_indices
I'll add more items as I find them, but I think this is more than enough to start with.