From 1db94ae5d2201edb5f9421c2c30be049efc678f6 Mon Sep 17 00:00:00 2001 From: Egor Tensin Date: Sun, 27 Dec 2015 12:56:13 +0300 Subject: utils: refactor command line parsing --- test/cavp.py | 22 +++++++------- test/nist-sp-800-38a.py | 14 ++++----- test/toolkit.py | 78 ++++++++++++++++++++++++++++++------------------- 3 files changed, 66 insertions(+), 48 deletions(-) (limited to 'test') diff --git a/test/cavp.py b/test/cavp.py index 2508bda..fb719f4 100644 --- a/test/cavp.py +++ b/test/cavp.py @@ -80,24 +80,24 @@ class _TestVectorsFile: init_vectors = self._parser.get(section, 'iv') return keys, plaintexts, ciphertexts, init_vectors - def _run_tests(self, tool, inputs, expected_output): + def _run_tests(self, tool, inputs, expected_output, use_boxes=False): for expected_output_chunk, input_chunk in _split_into_chunks(expected_output, list(inputs)): - actual_output = tool(self.algorithm(), self.mode(), input_chunk) + actual_output = tool(self.algorithm(), self.mode(), input_chunk, use_boxes=use_boxes) if not _assert_output(actual_output, expected_output_chunk): return _TestExitCode.FAILURE return _TestExitCode.SUCCESS - def run_encryption_tests(self, tools): + def run_encryption_tests(self, tools, use_boxes=False): logging.info('Running encryption tests...') keys, plaintexts, ciphertexts, init_vectors = self._extract_test_data('ENCRYPT') inputs = _gen_encryption_inputs(keys, plaintexts, init_vectors) - return self._run_tests(tools.run_encrypt_block, inputs, ciphertexts) + return self._run_tests(tools.run_encrypt_block, inputs, ciphertexts, use_boxes) - def run_decryption_tests(self, tools): + def run_decryption_tests(self, tools, use_boxes=False): logging.info('Running decryption tests...') keys, plaintexts, ciphertexts, init_vectors = self._extract_test_data('DECRYPT') inputs = _gen_decryption_inputs(keys, ciphertexts, init_vectors) - return self._run_tests(tools.run_decrypt_block, inputs, plaintexts) + return self._run_tests(tools.run_decrypt_block, inputs, plaintexts, use_boxes) def _parse(self): logging.info('Trying to parse test vectors file name \'{0}\'...'.format(self._fn)) @@ -145,7 +145,7 @@ class _TestVectorsFile: logging.warn('Unknown or unsupported mode: ' + self._fn) return None -def _parse_archive_and_run_tests(tools, archive_path): +def _parse_archive_and_run_tests(tools, archive_path, use_boxes=False): archive = zipfile.ZipFile(archive_path) exit_codes = [] for fn in archive.namelist(): @@ -153,13 +153,13 @@ def _parse_archive_and_run_tests(tools, archive_path): if member.recognized(): member.parse() try: - exit_codes.append(member.run_encryption_tests(tools)) + exit_codes.append(member.run_encryption_tests(tools, use_boxes)) except Exception as e: logging.error('Encountered an exception!') logging.exception(e) exit_codes.append(_TestExitCode.ERROR) try: - exit_codes.append(member.run_decryption_tests(tools)) + exit_codes.append(member.run_decryption_tests(tools, use_boxes)) except Exception as e: logging.error('Encountered an exception!') logging.exception(e) @@ -201,5 +201,5 @@ if __name__ == '__main__': logging_options['filename'] = args.log logging.basicConfig(**logging_options) - tools = toolkit.Tools(args.path, use_sde=args.sde, use_boxes=args.box) - _parse_archive_and_run_tests(tools, args.archive) + tools = toolkit.Tools(args.path, use_sde=args.sde) + _parse_archive_and_run_tests(tools, args.archive, use_boxes=args.box) diff --git a/test/nist-sp-800-38a.py b/test/nist-sp-800-38a.py index 7189dde..05f0a39 100644 --- a/test/nist-sp-800-38a.py +++ b/test/nist-sp-800-38a.py @@ -146,7 +146,7 @@ def _assert_output(actual, expected): class _TestExitCode: SUCCESS, FAILURE, ERROR, SKIPPED = range(4) -def _run_encryption_tests(tools, algo, mode): +def _run_encryption_tests(tools, algo, mode, use_boxes=False): logging.info('Running encryption tests...') key = _keys[algo] iv = None @@ -154,13 +154,13 @@ def _run_encryption_tests(tools, algo, mode): iv = _init_vectors[algo][mode] ciphertexts = _ciphertexts[algo][mode] _input = toolkit.EncryptionInput(key, _plaintexts, iv=iv) - actual_output = tools.run_encrypt_block(algo, mode, _input) + actual_output = tools.run_encrypt_block(algo, mode, _input, use_boxes) if _assert_output(actual_output, ciphertexts): return _TestExitCode.SUCCESS else: return _TestExitCode.FAILURE -def _run_decryption_tests(tools, algo, mode): +def _run_decryption_tests(tools, algo, mode, use_boxes=False): logging.info('Running decryption tests...') key = _keys[algo] iv = None @@ -168,7 +168,7 @@ def _run_decryption_tests(tools, algo, mode): iv = _init_vectors[algo][mode] ciphertexts = _ciphertexts[algo][mode] _input = toolkit.DecryptionInput(key, ciphertexts, iv=iv) - actual_output = tools.run_decrypt_block(algo, mode, _input) + actual_output = tools.run_decrypt_block(algo, mode, _input, use_boxes) if _assert_output(actual_output, _plaintexts): return _TestExitCode.SUCCESS else: @@ -186,7 +186,7 @@ if __name__ == '__main__': parser.add_argument('--log', '-l', help='set log file path') args = parser.parse_args() - tools = toolkit.Tools(args.path, use_sde=args.sde, use_boxes=args.box) + tools = toolkit.Tools(args.path, use_sde=args.sde) logging_options = { 'format': '%(asctime)s | %(module)s | %(levelname)s | %(message)s', @@ -216,13 +216,13 @@ if __name__ == '__main__': mode = maybe_mode logging.info('Mode: ' + mode) try: - exit_codes.append(_run_encryption_tests(tools, algo, mode)) + exit_codes.append(_run_encryption_tests(tools, algo, mode, use_boxes=args.box)) except Exception as e: logging.error('Encountered an exception!') logging.exception(e) exit_codes.append(_TestExitCode.ERROR) try: - exit_codes.append(_run_decryption_tests(tools, algo, mode)) + exit_codes.append(_run_decryption_tests(tools, algo, mode, use_boxes=args.box)) except Exception as e: logging.error('Encountered an exception!') logging.exception(e) diff --git a/test/toolkit.py b/test/toolkit.py index 50a42c9..170bccf 100644 --- a/test/toolkit.py +++ b/test/toolkit.py @@ -25,10 +25,10 @@ def mode_requires_init_vector(mode): return mode != ECB def to_supported_algorithm(s): - algo = is_algorithm_supported(s) - if algo is None: + algorithm = is_algorithm_supported(s) + if algorithm is None: raise NotImplementedError('unsupported algorithm ' + s) - return algo + return algorithm def is_algorithm_supported(s): s = s.lower() @@ -77,7 +77,7 @@ class DecryptionInput: return args class Tools: - def __init__(self, search_dirs, use_sde=False, use_boxes=False): + def __init__(self, search_dirs, use_sde=False): if search_dirs: if isinstance(search_dirs, str): os.environ['PATH'] += os.pathsep + search_dirs @@ -86,7 +86,6 @@ class Tools: else: os.environ['PATH'] += os.pathsep + str(search_dirs) self._use_sde = use_sde - self._use_boxes = use_boxes self._logger = logging.getLogger(__name__) _ENCRYPT_BLOCK = 'encrypt_block.exe' @@ -94,11 +93,8 @@ class Tools: _ENCRYPT_FILE = 'encrypt_file.exe' _DECRYPT_FILE = 'decrypt_file.exe' - def run(self, tool_path, algo, mode, args): + def run(self, tool_path, args): cmd_list = ['sde', '--', tool_path] if self._use_sde else [tool_path] - if self._use_boxes: - cmd_list.append('-b') - cmd_list.extend(('-a', algo, '-m', mode)) cmd_list.extend(args) logging.info('Trying to execute: {0}'.format( subprocess.list2cmdline(cmd_list))) @@ -126,32 +122,54 @@ class Tools: args.extend(tail.to_args()) return args - def run_encrypt_block(self, algo, mode, inputs): + def run_encrypt_block(self, algorithm, mode, inputs, use_boxes=False): + args = [ + '--algorithm', algorithm, + '--mode', mode, + ] + if use_boxes: + args.append('--box') if isinstance(inputs, collections.Iterable): - args = self._block_inputs_to_args(iter(inputs)) + args.extend(self._block_inputs_to_args(iter(inputs))) else: - args = inputs.to_args() - return self.run(self._ENCRYPT_BLOCK, algo, mode, args) - - def run_decrypt_block(self, algo, mode, inputs): + args.extend(inputs.to_args()) + return self.run(self._ENCRYPT_BLOCK, args) + + def run_decrypt_block(self, algorithm, mode, inputs, use_boxes=False): + args = [ + '--algorithm', algorithm, + '--mode', mode, + ] + if use_boxes: + args.append('--box') if isinstance(inputs, collections.Iterable): - args = self._block_inputs_to_args(iter(inputs)) + args.extend(self._block_inputs_to_args(iter(inputs))) else: - args = inputs.to_args() - return self.run(self._DECRYPT_BLOCK, algo, mode, args) - - def run_encrypt_file(self, algo, mode, key, input_path, output_path, iv=None): + args.extend(inputs.to_args()) + return self.run(self._DECRYPT_BLOCK, args) + + def run_encrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None): + args = [ + '--algorithm', algorithm, + '--mode', mode, + '--key', key, + '--input-path', input_path, + '--output-path', output_path] if mode_requires_init_vector(mode): if not iv: - raise ValueError('mode \'{}\' requires init vector'.format(mode)) - return self.run(self._ENCRYPT_FILE, algo, mode, (key, iv, input_path, output_path)) - else: - return self.run(self._ENCRYPT_FILE, algo, mode, (key, input_path, output_path)) - - def run_decrypt_file(self, algo, mode, key, input_path, output_path, iv=None): + raise ValueError('mode \'{}\' requires initialization vector'.format(mode)) + args.extend(('--iv', iv)) + return self.run(self._ENCRYPT_FILE, args) + + def run_decrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None): + args = [ + '--algorithm', algorithm, + '--mode', mode, + '--key', key, + '--input-path', input_path, + '--output-path', output_path] if mode_requires_init_vector(mode): if not iv: - raise ValueError('mode \'{}\' requires init vector'.format(mode)) - return self.run(self._DECRYPT_FILE, algo, mode, (key, iv, input_path, output_path)) - else: - return self.run(self._DECRYPT_FILE, algo, mode, (key, input_path, output_path)) + raise ValueError('mode \'{}\' requires initialization vector'.format(mode)) + args.extend(('--iv', iv)) + return self.run(self._DECRYPT_FILE, args) -- cgit v1.2.3