diff options
Diffstat (limited to 'test/toolkit.py')
| -rw-r--r-- | test/toolkit.py | 54 |
1 files changed, 39 insertions, 15 deletions
diff --git a/test/toolkit.py b/test/toolkit.py index 1cf26b6..8890269 100644 --- a/test/toolkit.py +++ b/test/toolkit.py @@ -6,36 +6,37 @@ import collections import logging import os.path import subprocess -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 +_SUPPORTED_ALGORITHMS = AES128, AES192, AES256 +_SUPPORTED_MODES = ECB, CBC, CFB, OFB, CTR def get_supported_algorithms(): - return _supported_algorithms + return _SUPPORTED_ALGORITHMS def get_supported_modes(): - return _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): s = s.lower() - if s in _supported_algorithms: + if s in _SUPPORTED_ALGORITHMS: return s - return None + raise NotImplementedError('unsupported algorithm ' + s) def to_supported_mode(s): s = s.lower() - if s in _supported_modes: + if s in _SUPPORTED_MODES: return s if s == CFB + '128': return CFB - return None + raise NotImplementedError('unsupported algorithm ' + s) class EncryptionInput: def __init__(self, key, plaintexts, iv=None): @@ -81,12 +82,14 @@ class Tools: _ENCRYPT_BLOCK = 'encrypt_block.exe' _DECRYPT_BLOCK = 'decrypt_block.exe' + _ENCRYPT_FILE = 'encrypt_file.exe' + _DECRYPT_FILE = 'decrypt_file.exe' def run(self, tool_path, algo, mode, 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(('-a', algo, '-m', mode)) cmd_list.extend(args) logging.info('Trying to execute: {0}'.format(subprocess.list2cmdline(cmd_list))) try: @@ -100,11 +103,12 @@ class Tools: return output.split() @staticmethod - def _inputs_to_args(inputs): + def _block_inputs_to_args(inputs): head = next(inputs, None) if head is None: - return [] - args = head.to_args() + return ['--'] + args = ['--'] + args.extend(head.to_args()) while True: tail = next(inputs, None) if tail is None: @@ -115,14 +119,34 @@ class Tools: def run_encrypt_block(self, algo, mode, inputs): if isinstance(inputs, collections.Iterable): - args = self._inputs_to_args(iter(inputs)) + args = 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): if isinstance(inputs, collections.Iterable): - args = self._inputs_to_args(iter(inputs)) + args = 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): + if mode_requires_init_vector(mode): + if not iv: + raise ToolkitError('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): + if mode_requires_init_vector(mode): + if not iv: + raise ToolkitError('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)) |