This repository has been archived by the owner on Oct 18, 2019. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
enigma.py
executable file
·394 lines (335 loc) · 12.4 KB
/
enigma.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
#!/usr/bin/env python3
"""Entrypoint into the Enigma simulation, processes command line arguments
and evaluates launch mode based upon them. Can launch in cli or gui mode.
Can read message data from stdin pipes."""
import argparse
import logging
from json import JSONDecodeError
from sys import stdin, stdout
from pytest import main as pytest_main
from benchmark import benchmark
from enigma.api.enigma_api import \
EnigmaAPI # pylint: disable=no-name-in-module
from enigma.core.components import ( # pylint: disable=no-name-in-module
DEFAULT_LAYOUT, HISTORICAL)
from enigma.interface.cli import cli
from enigma.interface.gui import load_views
from enigma.interface.gui.gui import runtime
from enigma.utils.cfg_handler import load_config
DEFAULT_INIT = {"model": "Enigma I", "rotors": ["I", "II", "III"], "reflector": "UKW-A",
"position_buffer": 1000000}
logging.basicConfig(level=logging.CRITICAL)
def config_from_args(args, load_to=None):
"""Returns dictionary of options aquired from args
:param args: argparse object with config
:param load_to: {EnigmaAPI} EnigmaAPI to load the config to
"""
config = {}
model = args.model
if model:
model = model[0]
config["model"] = model
if load_to:
load_to.model(model)
else:
config["model"] = None
if args.reflector:
config["reflector"] = args.reflector[0]
if load_to:
load_to.reflector(config["reflector"])
if args.rotors:
config["rotors"] = args.rotors
if load_to:
load_to.rotors(config["rotors"])
if args.positions:
config["positions"] = args.positions
if load_to:
load_to.positions(config["positions"])
if args.ring_settings:
config["ring_settings"] = args.ring_settings
if load_to:
load_to.ring_settings(config["ring_settings"])
if args.reflector_position:
config["reflector_position"] = args.reflector_position[0]
if load_to:
load_to.reflector_position(config["reflector_position"])
if args.reflector_pairs:
config["reflector_pairs"] = args.reflector_pairs
if load_to:
load_to.reflector_pairs(config["reflector_pairs"])
if args.uhr:
config["uhr_position"] = args.uhr[0]
if load_to:
load_to.uhr("connect")
load_to.uhr_position(config["uhr_position"])
if args.plug_pairs:
config["plug_pairs"] = args.plug_pairs
if load_to:
load_to.plug_pairs(config["plug_pairs"])
return config
def resolve_conflicts(args):
"""Resolves conflicting cli options"""
if args.run_tests and args.only_run_tests:
logging.error("Cannot run both --test and -T at once!")
logging.shutdown()
exit(1)
if args.silent and args.verbose:
logging.error("Conflicting flags --verbose and --silent!")
print("Conflicting flags --verbose and --silent!")
logging.shutdown()
exit(-1)
if args.silent and not args.cli:
logging.error("Silent mode not available for graphical mode!")
print("Silent mode not available for graphical mode!")
logging.shutdown()
exit(-1)
def load_custom(custom):
try:
for model, data in custom.items():
if HISTORICAL.get(model, False):
print("Could not load custom config: model '%s' already exists!" % model)
exit(1)
if not data.get("layout", False):
data["layout"] = DEFAULT_LAYOUT
HISTORICAL[model] = data
except (TypeError, KeyError, ValueError) as err:
print("Invalid custom data, please fix 'config.json'! Message: %s" % str(err))
exit(1)
load_views(custom)
if __name__ == "__main__":
CUSTOM = None
try:
CONFIG_DATA = load_config("config.json")
DEFAULT_INIT = CONFIG_DATA["default"]
CUSTOM = CONFIG_DATA.get("custom")
except (FileNotFoundError, KeyError, ValueError, JSONDecodeError):
msg = "Failed to load default config, using builtin defaults instead..."
logging.info(msg)
if CUSTOM:
load_custom(CUSTOM)
# FLAG ARGS ====================================================
PARSER = argparse.ArgumentParser(
description="returns Enigma encrypted text base"
"on settings and provided text."
)
ARGUMENT_DATA = (
(
("-t", "--test"),
dict(
help="runs tests before launching the simulator to"
"ensure it works correctly",
dest="run_tests",
),
),
(
("-T",),
dict(
help="runs tests with detailed stack traces and quits",
dest="only_run_tests",
),
),
(
("-c", "--cli"),
dict(help="launches the simulator in the command line mode", dest="cli"),
),
(("-p", "--preview"), dict(help="Runs a sample cli command")),
(("-v", "--verbose"), dict(help="Turns on verbose logging messages")),
(("-s", "--silent"), dict(help="Turns off all prints except cli output")),
)
for arg in ARGUMENT_DATA:
PARSER.add_argument(*arg[0], **arg[1], action="store_true", default=False)
PARSER.add_argument(
"-b",
"--benchmark",
help="benchmarks encryption speed for N characters",
nargs=1,
dest="benchmark_n",
metavar="N",
)
# SETTINGS ARGS ====================================================
CLI_ARGS = PARSER.add_argument_group("startup settings")
CLI_DATA = (
(
("--from",),
dict(help="file to load Enigma settings from", nargs=1, dest="filename"),
),
(
("--model",),
dict(
help="available Enigma models: %s" % ", ".join(HISTORICAL.keys()),
nargs=1,
dest="model",
),
),
(
("--rotors",),
dict(help="rotors that will be used", nargs="+", metavar="rotor"),
),
(
("--positions",),
dict(
help="starting rotor positions",
nargs="+",
default=None,
metavar="position",
),
),
(
("--ring_settings",),
dict(
help="rotor ring settings",
nargs="+",
default=None,
metavar="ring_setting",
),
),
(("--reflector",), dict(help="reflector that will be used", nargs=1)),
(
("--reflector_position",),
dict(
help="reflector position (only available in EnigmaD, Enigma K,"
"SwissK, EnigmaG, Railway, Tirpitz)",
nargs=1,
default=None,
metavar="position",
),
),
(
("--reflector_pairs",),
dict(
help="reflector wiring pairs for UKW-D (pairs do not "
"correspond with real wiring!)",
nargs="*",
default=None,
metavar="pair",
),
),
(
("--plug_pairs",),
dict(
help="letter pairs to connect in the plugboard",
nargs="*",
default=None,
metavar="pair",
),
),
(
("--uhr",),
dict(
help="connects uhr to plugboard and sets position",
nargs=1,
default=None,
metavar="position",
),
),
(
("-m", "--message"),
dict(help="text for encryption in cli mode", nargs=1, dest="message"),
),
)
for arg in CLI_DATA:
CLI_ARGS.add_argument(*arg[0], **arg[1])
# ARG PARSE ====================================================
ARGS = PARSER.parse_args()
# PRE-LAUNCH ACTIONS ===========================================
if ARGS.verbose and not ARGS.silent: # Set verbose logs
logging.basicConfig(
level=logging.INFO,
stream=stdout,
format="%(levelname)s:%(module)s:%(funcName)s: %(message)s"
)
logging.getLogger().setLevel(logging.INFO)
logging.getLogger('PySide2').setLevel(logging.CRITICAL)
else: # Disable logs (only show critical)
logging.basicConfig(level=logging.CRITICAL)
resolve_conflicts(ARGS) # Checks for conflicting options
# TEST PHASE ==================================================
if ARGS.run_tests:
logging.info("Running pre-launch tests...")
# -x = stop at first failure
EXIT_CODE = pytest_main(["tests", "-x", "--tb=no", "-s"])
if EXIT_CODE != 1:
logging.error("Pre-launch tests failed! Aborting...")
logging.shutdown()
exit(EXIT_CODE)
logging.info("All pre-launch tests succeeded...")
elif ARGS.only_run_tests:
logging.info("Running tests with detailed feedback...")
logging.shutdown()
exit(pytest_main(["tests", "--tb=long", "--durations=3", "-s"]))
# BENCHMARK ===========================================================
if ARGS.benchmark_n:
try:
N_LETTERS = int(ARGS.benchmark_n[0])
except ValueError:
logging.error('Invalid number "%s" for benchmark, exiting...', str(N_LETTERS))
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)
try:
config_from_args(ARGS, ENIGMA_API)
except (KeyError, ValueError) as err:
print(err)
exit(1)
logging.info("Loading settings specified in cli arguments...")
elif FILENAME: # Load from file specified in --from
try:
ENIGMA_API.load_from(FILENAME)
except FileNotFoundError:
msg = "No configuration file '%s' found!" % FILENAME
logging.info(msg)
print(msg)
exit(1)
except JSONDecodeError:
msg = "Configuration file '%s' is not of JSON format!" % FILENAME
logging.info(msg)
print(msg)
exit(1)
except KeyError:
msg = "Configuration file '%s' loaded but did not contain required data!" % FILENAME
logging.info(msg)
print(msg)
exit(1)
# APPLICATION INIT ====================================================
logging.info("Starting Enigma...")
if ARGS.cli: # Command line mode
logging.info("Loading in CLI mode with settings:\n%s...", str(ENIGMA_API))
# If stdin exists, load text from it
MESSAGE = None if stdin.isatty() else str(stdin.readline()).strip()
if MESSAGE is not None:
logging.info("Loaded input '%s' from stdin...", MESSAGE)
cli(ENIGMA_API, ARGS, MESSAGE)
elif ARGS.preview: # Preview command only
logging.info("Printing preview...")
print(
"Copy the command below:\n\n./enigma.py --cli --model 'Enigma I' "
"--rotors II I III --reflector UKW-A "
"--message THISISANENIGMASAMPLEMESSAGE"
)
else: # Graphical mode
logging.info("Launching Enigma Qt Application...")
runtime(ENIGMA_API)
# APPLICATION SHUTDOWN =================================================
logging.info("Program terminated...")
logging.shutdown()