test: refactoring

4 files changed, 38 insertions, 63 deletions

diff --git a/test/cavp.py b/test/cavp.py
index 01d1e8d..500a7ca 100644
--- a/test/cavp.py
+++ b/test/cavp.py
@@ -18,17 +18,11 @@ class _MultiOrderedDict(OrderedDict):
         else:
             super(OrderedDict, self).__setitem__(key, value)
 
-def _gen_inputs(cls, keys, plaintexts, init_vectors):
+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 cls(key, [plaintext], iv)
-
-def _gen_encryption_inputs(keys, plaintexts, init_vectors):
-    return _gen_inputs(toolkit.EncryptionInput, keys, plaintexts, init_vectors)
-
-def _gen_decryption_inputs(keys, ciphertexts, init_vectors):
-    return _gen_inputs(toolkit.DecryptionInput, keys, ciphertexts, init_vectors)
+        yield toolkit.BlockInput(key, [plaintext], iv)
 
 def _split_into_chunks(expected_output, inputs, max_len=100):
     for i in range(0, len(inputs), max_len):
@@ -90,13 +84,13 @@ class _TestVectorsFile:
     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)
+        inputs = _gen_inputs(keys, plaintexts, init_vectors)
         return self._run_tests(tools.run_encrypt_block, inputs, ciphertexts, use_boxes)
 
     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)
+        inputs = _gen_inputs(keys, ciphertexts, init_vectors)
         return self._run_tests(tools.run_decrypt_block, inputs, plaintexts, use_boxes)
 
     def _parse(self):
diff --git a/test/file.py b/test/file.py
index fdb2d3a..d1abc3d 100644
--- a/test/file.py
+++ b/test/file.py
@@ -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, use_boxes=False)
+    tools = toolkit.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 a096a42..eb882b7 100644
--- a/test/nist-sp-800-38a.py
+++ b/test/nist-sp-800-38a.py
@@ -153,7 +153,7 @@ def _run_encryption_tests(tools, algo, mode, use_boxes=False):
     if algo in _init_vectors and mode in _init_vectors[algo]:
         iv = _init_vectors[algo][mode]
     ciphertexts = _ciphertexts[algo][mode]
-    _input = toolkit.EncryptionInput(key, _plaintexts, iv=iv)
+    _input = toolkit.BlockInput(key, _plaintexts, iv=iv)
     actual_output = tools.run_encrypt_block(algo, mode, _input, use_boxes)
     if _assert_output(actual_output, ciphertexts):
         return _TestExitCode.SUCCESS
@@ -167,7 +167,7 @@ def _run_decryption_tests(tools, algo, mode, use_boxes=False):
     if algo in _init_vectors and mode in _init_vectors[algo]:
         iv = _init_vectors[algo][mode]
     ciphertexts = _ciphertexts[algo][mode]
-    _input = toolkit.DecryptionInput(key, ciphertexts, iv=iv)
+    _input = toolkit.BlockInput(key, ciphertexts, iv=iv)
     actual_output = tools.run_decrypt_block(algo, mode, _input, use_boxes)
     if _assert_output(actual_output, _plaintexts):
         return _TestExitCode.SUCCESS
diff --git a/test/toolkit.py b/test/toolkit.py
index 8798e6d..e3c7daf 100644
--- a/test/toolkit.py
+++ b/test/toolkit.py
@@ -50,7 +50,7 @@ def is_mode_supported(s):
         return CFB
     return None
 
-class EncryptionInput:
+class BlockInput:
     def __init__(self, key, plaintexts, iv=None):
         self.key = key
         self.plaintexts = plaintexts
@@ -63,19 +63,6 @@ class EncryptionInput:
         args.extend(self.plaintexts)
         return args
 
-class DecryptionInput:
-    def __init__(self, key, ciphertexts, iv=None):
-        self.key = key
-        self.ciphertexts = ciphertexts
-        self.iv = iv
-
-    def to_args(self):
-        args = [self.key]
-        if self.iv is not None:
-            args.append(self.iv)
-        args.extend(self.ciphertexts)
-        return args
-
 class Tools:
     def __init__(self, search_dirs, use_sde=False):
         if search_dirs:
@@ -109,67 +96,61 @@ class Tools:
 
     @staticmethod
     def _block_inputs_to_args(inputs):
-        head = next(inputs, None)
-        if head is None:
-            return ['--']
-        args = ['--']
-        args.extend(head.to_args())
+        args = []
         while True:
-            tail = next(inputs, None)
-            if tail is None:
+            head = next(inputs, None)
+            if head is None:
                 break
             args.append('--')
-            args.extend(tail.to_args())
+            args.extend(head.to_args())
         return args
 
-    def run_encrypt_block(self, algorithm, mode, inputs, use_boxes=False):
+    @staticmethod
+    def _block_settings_to_args(algorithm, mode, use_boxes=False):
         args = [
             '--algorithm', algorithm,
             '--mode', mode,
         ]
         if use_boxes:
             args.append('--use-boxes')
+        return args
+
+    @staticmethod
+    def _build_block_args(algorithm, mode, inputs, use_boxes=False):
+        args = Tools._block_settings_to_args(algorithm, mode, use_boxes)
         if isinstance(inputs, collections.Iterable):
-            args.extend(self._block_inputs_to_args(iter(inputs)))
+            args.extend(Tools._block_inputs_to_args(iter(inputs)))
         else:
             args.extend(inputs.to_args())
-        return self.run(self._ENCRYPT_BLOCK, args)
+        return args
+
+    def run_encrypt_block(self, algorithm, mode, inputs, use_boxes=False):
+        return self.run(self._ENCRYPT_BLOCK,
+                        self._build_block_args(algorithm, mode, inputs, use_boxes))
 
     def run_decrypt_block(self, algorithm, mode, inputs, use_boxes=False):
-        args = [
-            '--algorithm', algorithm,
-            '--mode', mode,
-        ]
-        if use_boxes:
-            args.append('--use-boxes')
-        if isinstance(inputs, collections.Iterable):
-            args.extend(self._block_inputs_to_args(iter(inputs)))
-        else:
-            args.extend(inputs.to_args())
-        return self.run(self._DECRYPT_BLOCK, args)
+        return self.run(self._DECRYPT_BLOCK,
+                        self._build_block_args(algorithm, mode, inputs, use_boxes))
 
-    def run_encrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None):
+    @staticmethod
+    def _file_settings_to_args(algorithm, mode, key, input_path, output_path, iv=None):
         args = [
             '--algorithm', algorithm,
             '--mode', mode,
             '--key', key,
             '--input-path', input_path,
-            '--output-path', output_path]
+            '--output-path', output_path
+        ]
         if mode_requires_init_vector(mode):
             if not iv:
                 raise ValueError('mode \'{}\' requires initialization vector'.format(mode))
             args.extend(('--iv', iv))
-        return self.run(self._ENCRYPT_FILE, args)
+        return args
+
+    def run_encrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None):
+        return self.run(self._ENCRYPT_FILE,
+                        self._file_settings_to_args(algorithm, mode, key, input_path, output_path, iv))
 
     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 initialization vector'.format(mode))
-            args.extend(('--iv', iv))
-        return self.run(self._DECRYPT_FILE, args)
+        return self.run(self._DECRYPT_FILE,
+                        self._file_settings_to_args(algorithm, mode, key, input_path, output_path, iv))