def test_enigma_models(model, n_rotors, should_fail):
rotors = list(range(n_rotors))
# Fake values instead of rotors because they are not needed in this test
if should_fail:
with pytest.raises(ValueError):
EnigmaAPI.generate_enigma(model, "UKW-B", rotors)
else:
EnigmaAPI.generate_enigma(model, "UKW-B", rotors)
def test_reflector_pairs(pairs, should_fail):
enigma_api = EnigmaAPI("Enigma M4")
enigma_api.reflector("UKW-D")
if should_fail:
with pytest.raises(ValueError):
enigma_api.reflector_pairs(pairs)
else:
enigma_api.reflector_pairs(pairs)
assert isinstance(enigma_api.encrypt("ABCD"), str)
def benchmark(char_n=None):
"""Benchmarks Enigma encryption speed with
the (theoretically) most performance heavy settings.
:param char_n: {int} Number of characters to benchmark on
"""
enigma = EnigmaAPI.generate_enigma("Enigma M4", "UKW-b",
["Beta", "I", "II", "III"])
enigma.uhr("connect")
enigma.uhr_position(3)
enigma.plug_pairs(
["AB", "CD", "EF", "GH", "IJ", "KL", "MN", "OP", "QR", "ST"])
logging.info("Benchmarking encryption speed with Enigma M4 and Uhr...")
if isinstance(char_n, int):
char_n = (char_n, )
else:
char_n = [10**num for num in range(7)]
for num in char_n:
start = time.time()
for letter in repeat("A", num):
enigma.press_key(letter)
end = time.time()
output = "Total encryption time: %.3f seconds for %d letters..." % (
end - start, num)
logging.info(output)
print(output)
def test_turnover():
enigma = EnigmaAPI.generate_component("Enigma M4", "rotors", "VI")
for _ in range(50):
if enigma.in_turnover():
assert enigma.position() in enigma._turnover
enigma.rotate()
def test_uhr_addon():
enigma = EnigmaAPI.generate_enigma("Enigma M4", "UKW-b",
["Beta", "I", "II", "III"])
enigma.uhr("connect")
enigma.uhr_position(3)
enigma.plug_pairs(
["AB", "CD", "EF", "GH", "IJ", "KL", "MN", "OP", "QR", "ST"])
setting = [3, 1, 1, 1]
enigma.positions(setting)
original = "TESTINGHOWUHRENCRYPTSTHISMESSAGE"
output = ""
for letter in original:
output += enigma.press_key(letter)
enigma.positions(setting)
result = ""
for letter in output:
result += enigma.press_key(letter)
assert result == original
enigma.positions(setting)
enigma.uhr_position(35)
result = ""
for letter in output:
result += enigma.press_key(letter)
assert result != original
def test_position():
rotor = EnigmaAPI.generate_component("Enigma I", "rotors", "I")
rotor.rotate()
assert rotor.position(True) == "02"
assert rotor.position() == "B"
rotor.rotate(20)
assert rotor.position() == "V"
def test_rotors_set():
enigma_api = EnigmaAPI("Enigma I")
rotors = ["I", "II", "III", "IV", "V"]
for _ in range(100):
new_rotors = choices(rotors, k=3)
enigma_api.rotors(new_rotors)
assert enigma_api.rotors() == new_rotors
for _ in range(100):
new_rotors = choices(range(len(rotors)), k=3)
enigma_api.rotors(new_rotors)
assert enigma_api.rotors() == [rotors[i] for i in new_rotors]
for trash in TRASH_DATA:
with pytest.raises(ValueError):
enigma_api.rotors(trash)
def test_enigma():
enigma = EnigmaAPI.generate_enigma("Enigma M3", "UKW-B",
["I", "II", "III"])
result = ""
for _ in "BDZGOW":
result += enigma.press_key("A")
assert result == "BDZGOW"
def test_reflector_set():
enigma_api = EnigmaAPI("Enigma I")
for new_reflector in "UKW-B", "UKW-A", "UKW-D":
enigma_api.reflector(new_reflector)
assert enigma_api.reflector() == new_reflector
for trash in TRASH_DATA:
with pytest.raises(ValueError):
enigma_api.reflector(trash)
def test_single_encrypt():
enigma = EnigmaAPI.generate_component("Enigma I", "rotors", "I")
assert enigma.forward("A") == "E"
enigma.rotate()
assert enigma.forward("A") == "J"
enigma.rotate()
assert enigma.forward("A") == "K"
# "Looping back" to position 0 should produce the same result as in default position
enigma.rotate(24)
assert enigma.forward("A") == "E"
def test_reflector_position():
enigma_api = EnigmaAPI("Enigma K", "UKW")
for _ in range(100):
new_position = randint(-30, 30)
if new_position not in range(1, 27):
with pytest.raises(ValueError):
enigma_api.reflector_position(new_position)
else:
enigma_api.reflector_position(new_position)
assert enigma_api.reflector_position() == alphabet[new_position - 1]
def test_ring_settings_set():
enigma_api = EnigmaAPI("Enigma M4")
for _ in range(100):
new_ring_settings = [randint(-30, 30) for _ in range(4)]
if not all(map(lambda x: x in range(1, 27), new_ring_settings)):
with pytest.raises(ValueError):
enigma_api.ring_settings(new_ring_settings)
else:
enigma_api.ring_settings(new_ring_settings)
assert enigma_api.ring_settings() == new_ring_settings
def test_routing():
"""
Tests if the forward routing is being routed correctly in the opposite direction (taking the
relative rotor position into account)
"""
rotor = EnigmaAPI.generate_component("Enigma I", "rotors", "I")
for i in 1, 3, -2, 5, 7, 20:
for letter in alphabet:
assert letter == rotor.backward(
rotor.forward(letter)
), "Backwards routing doesn't return to the original location!"
rotor.rotate(i)
def test_model_set():
enigma_api = EnigmaAPI("Enigma I")
labels = list(HISTORICAL.keys())
for _ in range(100):
new_model = choice(labels)
enigma_api.model(new_model)
assert enigma_api.model() == new_model
for trash in TRASH_DATA:
with pytest.raises(ValueError):
enigma_api.model(trash)
def test_generate_component():
for _ in range(1000):
model = choice(list(HISTORICAL.keys()))
comp_type = choice(["rotors", "reflectors"])
component_data = choice(HISTORICAL[model][comp_type])
if component_data["label"] == "UKW-D":
continue
component = EnigmaAPI.generate_component(model, comp_type, label=component_data["label"])
assert component.label() == component_data["label"]
assert component._wiring == component_data["wiring"]
if comp_type == "rotors" and component_data["label"] not in ("Beta", "Gamma"):
assert component._turnover == component_data["turnover"]
def test_generate_enigma():
for _ in range(100):
model = choice(list(HISTORICAL.keys()))
model_data = HISTORICAL[model]
enigma = EnigmaAPI.generate_enigma(model)
assert enigma.model() == model
assert enigma.reflector_rotatable() == model_data["rotatable_ref"]
assert bool(enigma._plugboard) == model_data["plugboard"]
assert enigma._numeric == model_data["numeric"]
assert enigma._charset == model_data["charset"]
assert enigma.rotor_n() == model_data["rotor_n"]
assert enigma.charset() == model_data["charset"]
def test_generate_rotor_callback():
enigma_api = EnigmaAPI("Enigma I")
for i in list(range(3)) * 10:
rotate_by = randint(0, 25)
callback = enigma_api.generate_rotate_callback(i, rotate_by)
callback()
assert int(enigma_api.positions()[i]) - 1 == rotate_by
enigma_api.positions([1, 1, 1])
def test_positions_set():
enigma_api = EnigmaAPI("Enigma M4")
for _ in range(100):
new_positions = [randint(-30, 30) for _ in range(4)]
if not all(map(lambda x: x in range(1, 27), new_positions)):
with pytest.raises(ValueError):
enigma_api.positions(new_positions)
else:
enigma_api.positions(new_positions)
assert enigma_api.positions() == tuple(
map(lambda x: alphabet[x - 1], new_positions)
)
def test_implementation():
"""
Tests the implementation by encrypting each letter of the alphabet
4000 times and checking if the encrypted message does not contain the
letter (this must always be true because the Enigma worked this way)
The number of iterations is arbitrary (not to slow down the testing process)
"""
enigma = EnigmaAPI.generate_enigma("Enigma M3", "UKW-B",
["I", "II", "III"])
for letter in enigma.charset():
for _ in range(1000):
assert enigma.press_key(
letter) != letter, "Enigma implementation wrong!"
def test_historical_messages(
model,
rotors,
reflector,
positions,
ring_settings,
plug_pairs,
message,
correct_result,
):
# Enigma instruction manual message
enigma = EnigmaAPI.generate_enigma(model, reflector, rotors)
enigma.positions(positions)
enigma.ring_settings(ring_settings)
enigma.plug_pairs(plug_pairs)
result = ""
for letter in message:
result += enigma.press_key(letter)
assert result == correct_result
def test_uhr():
enigma_api = EnigmaAPI("Enigma M3")
pairs = generate_pairs(10)
enigma_api.uhr("connect")
enigma_api.plug_pairs(pairs)
correct, other = sample(range(0, 40), 2)
enigma_api.uhr_position(correct)
message = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
enigma_api.set_checkpoint()
encrypted = enigma_api.encrypt(message)
enigma_api.load_checkpoint()
assert enigma_api.encrypt(encrypted) == message
enigma_api.uhr_position(other)
assert enigma_api.encrypt(encrypted) != message
# Test Uhr position 00 compatibility mode
enigma_api.load_checkpoint()
enigma_api.uhr_position(0)
with_uhr = enigma_api.encrypt(message)
enigma_api.load_checkpoint()
enigma_api.uhr("disconnect")
enigma_api.plug_pairs(pairs)
without_uhr = enigma_api.encrypt(message)
assert with_uhr == without_uhr
def test_reflector():
reflector = EnigmaAPI.generate_component("Enigma M3", "reflectors",
"UKW-B")
assert reflector.reflect(reflector.reflect("A")) == "A"
def test_reflector_rotatable():
for model, data in HISTORICAL.items():
assert EnigmaAPI(model).reflector_rotatable() == data["rotatable_ref"]
def test_rotor_n(settings, rotor_n):
enigma_api = EnigmaAPI(*settings)
assert enigma_api.rotor_n() == rotor_n, "Incorrect rotor_n value!"
def test_data():
enigma_api = EnigmaAPI("Enigma M3", "UKW-B", ["I", "II", "III"])
assert enigma_api.data() == HISTORICAL["Enigma M3"]
enigma_api.model("Enigma M4")
assert enigma_api.data() == HISTORICAL["Enigma M4"]
print('Invalid number "%s", choose a valid number greater than 0!' % str(N_LETTERS))
exit(1)
if N_LETTERS <= 0:
logging.error("Benchmark character count is not greater than 0, exiting...")
print("Benchmark character count must be greater than 0!")
exit(1)
benchmark(N_LETTERS)
logging.shutdown()
exit()
# CONFIG LOAD =========================================================
logging.info("Loading config...")
ENIGMA_API = EnigmaAPI(**DEFAULT_INIT) # Fallback default configuration
CONFIG = config_from_args(ARGS)
FILENAME = None
if ARGS.filename:
FILENAME = ARGS.filename[0]
HAS_CONFIG = any(CONFIG.values())
if HAS_CONFIG and FILENAME:
print("Cannot load settings both from arguments and file!")
exit(1)
if HAS_CONFIG: # Load from settings arguments
if not CONFIG["model"] and len(CONFIG) > 1:
print("Enigma model must be specified when specifying settings!")
exit(1)
def test_letter_group(model, letter_group):
enigma_api = EnigmaAPI(model)
assert enigma_api.letter_group() == letter_group, "Incorrect letter group length!"
def test_plug_pairs():
pairs = generate_pairs(randint(0, 13))
enigma_api = EnigmaAPI("Enigma M3")
enigma_api.plug_pairs(pairs)
assert enigma_api.plug_pairs() == pairs
def test_rotation(offset_by, result):
enigma = EnigmaAPI.generate_component("Enigma I", "rotors", "I")
enigma.rotate(offset_by)
assert enigma.offset(
) == result, "Rotor offset is not being calculated correctly"
def test_model_labels(model, labels):
enigma_api = EnigmaAPI(model)
assert (
enigma_api.model_labels(model) == labels
), "Incorrect (or incomplete) list of labels for selected model!"