utils: refactor command line parsing

3 files changed, 66 insertions, 48 deletions

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)