test/toolkit.py: bugfix

* Add algorithm and mode detection functions.
* run_encrypt_tool & run_decrypt_tool now accept a list of inputs rather
than a single input.

1 files changed, 61 insertions, 5 deletions

diff --git a/test/toolkit.py b/test/toolkit.py
index a69ce73..fdbd51e 100644
--- a/test/toolkit.py
+++ b/test/toolkit.py
@@ -2,6 +2,7 @@
 # This file is licensed under the terms of the MIT License.
 # See LICENSE.txt for details.
 
+import collections
 import logging
 import os.path
 import subprocess
@@ -10,6 +11,32 @@ import sys
 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):
+    return mode != ECB
+
+def to_supported_algorithm(s):
+    s = s.lower()
+    if s in _supported_algorithms:
+        return s
+    return None
+
+def to_supported_mode(s):
+    s = s.lower()
+    if s in _supported_modes:
+        return s
+    if s == CFB + '128':
+        return CFB
+    return None
+
 class EncryptionInput:
     def __init__(self, key, plaintexts, iv=None):
         self.key = key
@@ -36,6 +63,9 @@ class DecryptionInput:
         args.extend(self.ciphertexts)
         return args
 
+class ToolkitError(RuntimeError):
+    pass
+
 class Tools:
     def __init__(self, root_dir_path, use_sde=False):
         self._root_dir_path = root_dir_path
@@ -55,12 +85,38 @@ class Tools:
         cmd_list = ['sde', '--', tool_path] if self._use_sde else [tool_path]
         cmd_list.extend(args)
         logging.info('Trying to execute: {0}'.format(subprocess.list2cmdline(cmd_list)))
-        output = subprocess.check_output(cmd_list, universal_newlines=True)
+        try:
+            output = subprocess.check_output(cmd_list, universal_newlines=True)
+        except subprocess.CalledProcessError as e:
+            logging.exception(e)
+            raise ToolkitError() from e
         logging.info('Output:\n' + output)
         return output.split()
 
-    def run_encrypt_tool(self, algo, mode, input):
-        return self.run_tool(self.get_encrypt_tool_path(algo, mode), input.to_args())
+    @staticmethod
+    def _inputs_to_args(inputs):
+        head = next(inputs, None)
+        if head is None:
+            return []
+        args = head.to_args()
+        while True:
+            tail = next(inputs, None)
+            if tail is None:
+                break
+            args.append('--')
+            args.extend(tail.to_args())
+        return args
+
+    def run_encrypt_tool(self, algo, mode, inputs):
+        if isinstance(inputs, collections.Iterable):
+            args = self._inputs_to_args(iter(inputs))
+        else:
+            args = inputs.to_args()
+        return self.run_tool(self.get_encrypt_tool_path(algo, mode), args)
 
-    def run_decrypt_tool(self, algo, mode, input):
-        return self.run_tool(self.get_decrypt_tool_path(algo, mode), input.to_args())
+    def run_decrypt_tool(self, algo, mode, inputs):
+        if isinstance(inputs, collections.Iterable):
+            args = self._inputs_to_args(iter(inputs))
+        else:
+            args = inputs.to_args()
+        return self.run_tool(self.get_decrypt_tool_path(algo, mode), args)