add test/file.py

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))