From 7c0467718e15f9bc1296d82325133bbba7cd2388 Mon Sep 17 00:00:00 2001 From: Egor Tensin Date: Sat, 13 Feb 2016 01:39:46 +0300 Subject: test: use Enum --- test/cavp.py | 13 +++---- test/file.py | 22 ++++++------ test/nist-sp-800-38a.py | 78 +++++++++++++++++----------------------- test/toolkit.py | 95 ++++++++++++++++++++++++------------------------- 4 files changed, 97 insertions(+), 111 deletions(-) diff --git a/test/cavp.py b/test/cavp.py index 64caf4a..00fe715 100644 --- a/test/cavp.py +++ b/test/cavp.py @@ -8,9 +8,10 @@ from datetime import datetime import logging import os.path import sys -import toolkit import zipfile +from toolkit import * + class _MultiOrderedDict(OrderedDict): def __setitem__(self, key, value): if isinstance(value, list) and key in self: @@ -22,7 +23,7 @@ def _gen_inputs(keys, plaintexts, init_vectors): if init_vectors is None: init_vectors = [None for key in keys] for key, plaintext, iv in zip(keys, plaintexts, init_vectors): - yield toolkit.BlockInput(key, [plaintext], iv) + yield BlockInput(key, [plaintext], iv) def _split_into_chunks(expected_output, inputs, max_len=100): for i in range(0, len(inputs), max_len): @@ -70,7 +71,7 @@ class _TestVectorsFile: plaintexts = self._parser.get(section, 'plaintext') ciphertexts = self._parser.get(section, 'ciphertext') init_vectors = None - if toolkit.mode_requires_init_vector(self.mode()): + if self.mode().requires_init_vector(): init_vectors = self._parser.get(section, 'iv') return keys, plaintexts, ciphertexts, init_vectors @@ -115,7 +116,7 @@ class _TestVectorsFile: def _strip_algorithm(self, stub): key_size = stub[-3:] maybe_algorithm = 'aes{0}'.format(key_size) - self._algorithm = toolkit.is_algorithm_supported(maybe_algorithm) + self._algorithm = Algorithm.try_parse(maybe_algorithm) if self._algorithm: logging.info('\tAlgorithm: {0}'.format(self._algorithm)) return stub[0:-3] @@ -131,7 +132,7 @@ class _TestVectorsFile: logging.warn('Unknown or unsupported method: ' + self._fn) def _strip_mode(self, stub): - self._mode = toolkit.is_mode_supported(stub) + self._mode = Mode.try_parse(stub) if self._mode: logging.info('\tMode: {0}'.format(self._mode)) return self._mode @@ -195,5 +196,5 @@ if __name__ == '__main__': logging_options['filename'] = args.log logging.basicConfig(**logging_options) - tools = toolkit.Tools(args.path, use_sde=args.sde) + tools = Tools(args.path, use_sde=args.sde) _parse_archive_and_run_tests(tools, args.archive, use_boxes=args.use_boxes) diff --git a/test/file.py b/test/file.py index 615132c..3f07aaa 100644 --- a/test/file.py +++ b/test/file.py @@ -2,8 +2,6 @@ # This file is licensed under the terms of the MIT License. # See LICENSE.txt for details. -import toolkit - from datetime import datetime from glob import iglob as glob import filecmp @@ -13,6 +11,8 @@ import shutil import sys from tempfile import TemporaryDirectory +from toolkit import * + class _TestExitCode: SUCCESS, FAILURE, ERROR, SKIPPED = range(4) @@ -94,37 +94,37 @@ def _run_tests(tools, suite_dir, force=False): with TemporaryDirectory() as tmp_dir: for algorithm_dir in _list_dirs(suite_dir): algorithm = os.path.basename(algorithm_dir) - maybe_algorithm = toolkit.is_algorithm_supported(algorithm) + maybe_algorithm = Algorithm.try_parse(algorithm) if maybe_algorithm is None: logging.warn('Unknown or unsupported algorithm: ' + algorithm) exit_codes.append(_TestExitCode.SKIPPED) continue algorithm = maybe_algorithm - logging.info('Algorithm: ' + algorithm) + logging.info('Algorithm: {}'.format(algorithm)) for mode_dir in _list_dirs(algorithm_dir): mode = os.path.basename(mode_dir) - maybe_mode = toolkit.is_mode_supported(mode) + maybe_mode = Mode.try_parse(mode) if maybe_mode is None: logging.warn('Unknown or unsupported mode: ' + mode) exit_codes.append(_TestExitCode.SKIPPED) continue mode = maybe_mode - logging.info('Mode: ' + mode) + logging.info('Mode: {}'.format(mode)) for key_path in _list_keys(mode_dir): key = _read_key(key_path) logging.info('Key: ' + key) test_name = _extract_test_name(key_path) logging.info('Test name: ' + test_name) iv = None - if toolkit.mode_requires_init_vector(mode): + if mode.requires_init_vector(): iv_path = _build_iv_path(key_path) iv = _read_iv(iv_path) plain_path = _build_plain_path(key_path) cipher_path = _build_cipher_path(key_path) - os.makedirs(os.path.join(tmp_dir, algorithm, mode), 0o777, True) + os.makedirs(os.path.join(tmp_dir, str(algorithm), str(mode)), 0o777, True) try: exit_codes.append(_run_encryption_test( - tools, os.path.join(tmp_dir, algorithm, mode), + tools, os.path.join(tmp_dir, str(algorithm), str(mode)), algorithm, mode, key, plain_path, cipher_path, iv, force)) except Exception as e: logging.error('Encountered an exception!') @@ -133,7 +133,7 @@ def _run_tests(tools, suite_dir, force=False): if not force: try: exit_codes.append(_run_decryption_test( - tools, os.path.join(tmp_dir, algorithm, mode), + tools, os.path.join(tmp_dir, str(algorithm), str(mode)), algorithm, mode, key, cipher_path, plain_path, iv)) except Exception as e: logging.error('Encountered an exception!') @@ -173,5 +173,5 @@ if __name__ == '__main__': logging_options['filename'] = args.log logging.basicConfig(**logging_options) - tools = toolkit.Tools(args.path, use_sde=args.sde) + tools = Tools(args.path, use_sde=args.sde) _run_tests(tools, args.suite, args.force) diff --git a/test/nist-sp-800-38a.py b/test/nist-sp-800-38a.py index a8fb57e..c0ad9e6 100644 --- a/test/nist-sp-800-38a.py +++ b/test/nist-sp-800-38a.py @@ -8,7 +8,7 @@ import logging import os.path import sys -import toolkit +from toolkit import * _TEST_PLAINTEXTS = [ '6bc1bee22e409f96e93d7e117393172a', @@ -18,117 +18,117 @@ _TEST_PLAINTEXTS = [ ] _TEST_KEYS = { - toolkit.AES128: '2b7e151628aed2a6abf7158809cf4f3c', - toolkit.AES192: '8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b', - toolkit.AES256: '603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4' + Algorithm.AES128: '2b7e151628aed2a6abf7158809cf4f3c', + Algorithm.AES192: '8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b', + Algorithm.AES256: '603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4' } _TEST_INIT_VECTOR = '000102030405060708090a0b0c0d0e0f' _TEST_INIT_VECTOR_CTR = 'f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff' _TEST_INIT_VECTORS = { - toolkit.AES128: { - toolkit.CBC: _TEST_INIT_VECTOR, - toolkit.CFB: _TEST_INIT_VECTOR, - toolkit.OFB: _TEST_INIT_VECTOR, - toolkit.CTR: _TEST_INIT_VECTOR_CTR + Algorithm.AES128: { + Mode.CBC: _TEST_INIT_VECTOR, + Mode.CFB: _TEST_INIT_VECTOR, + Mode.OFB: _TEST_INIT_VECTOR, + Mode.CTR: _TEST_INIT_VECTOR_CTR } } -_TEST_INIT_VECTORS[toolkit.AES192] = _TEST_INIT_VECTORS[toolkit.AES128] -_TEST_INIT_VECTORS[toolkit.AES256] = _TEST_INIT_VECTORS[toolkit.AES128] +_TEST_INIT_VECTORS[Algorithm.AES192] = _TEST_INIT_VECTORS[Algorithm.AES128] +_TEST_INIT_VECTORS[Algorithm.AES256] = _TEST_INIT_VECTORS[Algorithm.AES128] _TEST_CIPHERTEXTS = { - toolkit.AES128: { - toolkit.ECB: [ + Algorithm.AES128: { + Mode.ECB: [ '3ad77bb40d7a3660a89ecaf32466ef97', 'f5d3d58503b9699de785895a96fdbaaf', '43b1cd7f598ece23881b00e3ed030688', '7b0c785e27e8ad3f8223207104725dd4' ], - toolkit.CBC: [ + Mode.CBC: [ '7649abac8119b246cee98e9b12e9197d', '5086cb9b507219ee95db113a917678b2', '73bed6b8e3c1743b7116e69e22229516', '3ff1caa1681fac09120eca307586e1a7' ], - toolkit.CFB: [ + Mode.CFB: [ '3b3fd92eb72dad20333449f8e83cfb4a', 'c8a64537a0b3a93fcde3cdad9f1ce58b', '26751f67a3cbb140b1808cf187a4f4df', 'c04b05357c5d1c0eeac4c66f9ff7f2e6' ], - toolkit.OFB: [ + Mode.OFB: [ '3b3fd92eb72dad20333449f8e83cfb4a', '7789508d16918f03f53c52dac54ed825', '9740051e9c5fecf64344f7a82260edcc', '304c6528f659c77866a510d9c1d6ae5e' ], - toolkit.CTR: [ + Mode.CTR: [ '874d6191b620e3261bef6864990db6ce', '9806f66b7970fdff8617187bb9fffdff', '5ae4df3edbd5d35e5b4f09020db03eab', '1e031dda2fbe03d1792170a0f3009cee' ] }, - toolkit.AES192: { - toolkit.ECB: [ + Algorithm.AES192: { + Mode.ECB: [ 'bd334f1d6e45f25ff712a214571fa5cc', '974104846d0ad3ad7734ecb3ecee4eef', 'ef7afd2270e2e60adce0ba2face6444e', '9a4b41ba738d6c72fb16691603c18e0e' ], - toolkit.CBC: [ + Mode.CBC: [ '4f021db243bc633d7178183a9fa071e8', 'b4d9ada9ad7dedf4e5e738763f69145a', '571b242012fb7ae07fa9baac3df102e0', '08b0e27988598881d920a9e64f5615cd' ], - toolkit.CFB: [ + Mode.CFB: [ 'cdc80d6fddf18cab34c25909c99a4174', '67ce7f7f81173621961a2b70171d3d7a', '2e1e8a1dd59b88b1c8e60fed1efac4c9', 'c05f9f9ca9834fa042ae8fba584b09ff' ], - toolkit.OFB: [ + Mode.OFB: [ 'cdc80d6fddf18cab34c25909c99a4174', 'fcc28b8d4c63837c09e81700c1100401', '8d9a9aeac0f6596f559c6d4daf59a5f2', '6d9f200857ca6c3e9cac524bd9acc92a' ], - toolkit.CTR: [ + Mode.CTR: [ '1abc932417521ca24f2b0459fe7e6e0b', '090339ec0aa6faefd5ccc2c6f4ce8e94', '1e36b26bd1ebc670d1bd1d665620abf7', '4f78a7f6d29809585a97daec58c6b050' ] }, - toolkit.AES256: { - toolkit.ECB: [ + Algorithm.AES256: { + Mode.ECB: [ 'f3eed1bdb5d2a03c064b5a7e3db181f8', '591ccb10d410ed26dc5ba74a31362870', 'b6ed21b99ca6f4f9f153e7b1beafed1d', '23304b7a39f9f3ff067d8d8f9e24ecc7' ], - toolkit.CBC: [ + Mode.CBC: [ 'f58c4c04d6e5f1ba779eabfb5f7bfbd6', '9cfc4e967edb808d679f777bc6702c7d', '39f23369a9d9bacfa530e26304231461', 'b2eb05e2c39be9fcda6c19078c6a9d1b' ], - toolkit.CFB: [ + Mode.CFB: [ 'dc7e84bfda79164b7ecd8486985d3860', '39ffed143b28b1c832113c6331e5407b', 'df10132415e54b92a13ed0a8267ae2f9', '75a385741ab9cef82031623d55b1e471' ], - toolkit.OFB: [ + Mode.OFB: [ 'dc7e84bfda79164b7ecd8486985d3860', '4febdc6740d20b3ac88f6ad82a4fb08d', '71ab47a086e86eedf39d1c5bba97c408', '0126141d67f37be8538f5a8be740e484' ], - toolkit.CTR: [ + Mode.CTR: [ '601ec313775789a5b7a7f504bbf3d228', 'f443e3ca4d62b59aca84e990cacaf5c5', '2b0930daa23de94ce87017ba2d84988d', @@ -144,7 +144,7 @@ def get_test_key(algorithm, mode=None): return _TEST_KEYS[algorithm] def get_test_iv(algorithm, mode): - if not toolkit.mode_requires_init_vector(mode): + if not mode.requires_init_vector(): return None return _TEST_INIT_VECTORS[algorithm][mode] @@ -176,7 +176,7 @@ def run_encryption_test(tools, algorithm, mode, use_boxes=False): key = get_test_key(algorithm, mode) iv = get_test_iv(algorithm, mode) expected_ciphertexts = get_test_ciphertexts(algorithm, mode) - input_ = toolkit.BlockInput(key, plaintexts, iv=iv) + input_ = BlockInput(key, plaintexts, iv=iv) actual_ciphertexts = tools.run_encrypt_block(algorithm, mode, input_, use_boxes) if verify_test_output(actual_ciphertexts, expected_ciphertexts): return TestExitCode.SUCCESS @@ -191,7 +191,7 @@ def run_decryption_test(tools, algorithm, mode, use_boxes=False): key = get_test_key(algorithm, mode) iv = get_test_iv(algorithm, mode) expected_plaintexts = get_test_plaintexts(algorithm, mode) - input_ = toolkit.BlockInput(key, ciphertexts, iv=iv) + input_ = BlockInput(key, ciphertexts, iv=iv) actual_plaintexts = tools.run_decrypt_block(algorithm, mode, input_, use_boxes) if verify_test_output(actual_plaintexts, expected_plaintexts): return TestExitCode.SUCCESS @@ -199,18 +199,6 @@ def run_decryption_test(tools, algorithm, mode, use_boxes=False): return TestExitCode.FAILURE def _run_tests(tools, algorithm, mode, use_boxes=False): - maybe_algorithm = toolkit.is_algorithm_supported(algorithm) - if maybe_algorithm is None: - logging.warn('Unknown or unsupported algorithm: {}'.format(maybe_algorithm)) - yield TestExitCode.SKIPPED - raise StopIteration - algorithm = maybe_algorithm - maybe_mode = toolkit.is_mode_supported(mode) - if maybe_mode is None: - logging.warn('Unknown or unsupported node: {}'.format(maybe_mode)) - yield TestExitCode.SKIPPED - raise StopIteration - mode = maybe_mode try: yield run_encryption_test(tools, algorithm, mode, use_boxes=args.use_boxes) except Exception as e: @@ -241,7 +229,7 @@ if __name__ == '__main__': help='set log file path') args = parser.parse_args() - tools = toolkit.Tools(args.path, use_sde=args.sde) + tools = Tools(args.path, use_sde=args.sde) logging.basicConfig(filename=args.log, format='%(asctime)s | %(module)s | %(levelname)s | %(message)s', diff --git a/test/toolkit.py b/test/toolkit.py index e3c7daf..e94fe29 100644 --- a/test/toolkit.py +++ b/test/toolkit.py @@ -3,52 +3,51 @@ # See LICENSE.txt for details. import collections +from enum import Enum import logging import os.path import subprocess -AES128, AES192, AES256 = 'aes128', 'aes192', 'aes256' -ECB, CBC, CFB, OFB, CTR = 'ecb', 'cbc', 'cfb', 'ofb', 'ctr' - -_SUPPORTED_ALGORITHMS = AES128, AES192, AES256 -_SUPPORTED_MODES = ECB, CBC, CFB, OFB, CTR - -def get_supported_algorithms(): - return _SUPPORTED_ALGORITHMS - -def get_supported_modes(): - return _SUPPORTED_MODES - -def mode_requires_init_vector(mode): - if mode not in _SUPPORTED_MODES: - raise NotImplementedError('unsupported mode of operation ' + s) - return mode != ECB - -def to_supported_algorithm(s): - algorithm = is_algorithm_supported(s) - if algorithm is None: - raise NotImplementedError('unsupported algorithm ' + s) - return algorithm - -def is_algorithm_supported(s): - s = s.lower() - if s in _SUPPORTED_ALGORITHMS: - return s - return None - -def to_supported_mode(s): - mode = is_mode_supported(s) - if mode is None: - raise NotImplementedError('unsupported mode ' + s) - return mode - -def is_mode_supported(s): - s = s.lower() - if s in _SUPPORTED_MODES: - return s - if s == CFB + '128': - return CFB - return None +class Algorithm(Enum): + @staticmethod + def parse(s): + return Algorithm(s.lower()) + + @staticmethod + def try_parse(s): + try: + return Algorithm.parse(s) + except ValueError: + return None + + AES128, AES192, AES256 = 'aes128', 'aes192', 'aes256' + + def __str__(self): + return self.value + + +class Mode(Enum): + @staticmethod + def parse(s): + s = s.lower() + if '{}128'.format(Mode.CFB) == s: + return Mode.CFB + return Mode(s) + + @staticmethod + def try_parse(s): + try: + return Mode.parse(s) + except ValueError: + return None + + ECB, CBC, CFB, OFB, CTR = 'ecb', 'cbc', 'cfb', 'ofb', 'ctr' + + def requires_init_vector(self): + return self != Mode.ECB + + def __str__(self): + return self.value class BlockInput: def __init__(self, key, plaintexts, iv=None): @@ -108,8 +107,8 @@ class Tools: @staticmethod def _block_settings_to_args(algorithm, mode, use_boxes=False): args = [ - '--algorithm', algorithm, - '--mode', mode, + '--algorithm', str(algorithm), + '--mode', str(mode), ] if use_boxes: args.append('--use-boxes') @@ -135,15 +134,13 @@ class Tools: @staticmethod def _file_settings_to_args(algorithm, mode, key, input_path, output_path, iv=None): args = [ - '--algorithm', algorithm, - '--mode', mode, + '--algorithm', str(algorithm), + '--mode', str(mode), '--key', key, '--input-path', input_path, '--output-path', output_path ] - if mode_requires_init_vector(mode): - if not iv: - raise ValueError('mode \'{}\' requires initialization vector'.format(mode)) + if iv is not None: args.extend(('--iv', iv)) return args -- cgit v1.2.3