diff options
| -rw-r--r-- | test/toolkit.py | 66 |
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) |