aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/test/file.py
blob: a9189841ecd49d325f281f438b86be57bf7ca7b7 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# Copyright 2015 Egor Tensin <Egor.Tensin@gmail.com>
# This file is licensed under the terms of the MIT License.
# See LICENSE.txt for details.

from datetime import datetime
from enum import Enum
from glob import iglob as glob
import filecmp
import logging
import os
import shutil
import sys
from tempfile import TemporaryDirectory

from toolkit import *

class TestExitCode(Enum):
    SUCCESS, FAILURE, ERROR, SKIPPED = range(1, 5)

_KEY_EXT = 'key'
_IV_EXT = 'iv'
_PLAIN_EXT = 'plain'
_CIPHER_EXT = 'cipher'

def _run_encryption_test(tools, tmp_dir, algorithm, mode, key, plain_path, cipher_path, iv=None, force=False):
    logging.info('Running encryption test...')
    logging.info('\tPlaintext file path: ' + plain_path)
    logging.info('\tExpected ciphertext file path: ' + cipher_path)
    tmp_path = os.path.join(tmp_dir, os.path.basename(cipher_path))
    logging.info('\tEncrypted file path: ' + tmp_path)
    tools.run_encrypt_file(algorithm, mode, key, plain_path, tmp_path, iv)
    if force:
        logging.warn('Overwriting expected ciphertext file')
        shutil.copy(tmp_path, cipher_path)
        return TestExitCode.SKIPPED
    if filecmp.cmp(cipher_path, tmp_path):
        return TestExitCode.SUCCESS
    else:
        logging.error('The encrypted file doesn\'t match the ciphertext file')
        return TestExitCode.FAILURE

def _run_decryption_test(tools, tmp_dir, algorithm, mode, key, cipher_path, plain_path, iv=None):
    logging.info('Running decryption test...')
    logging.info('\tCiphertext file path: ' + cipher_path)
    logging.info('\tExpected plaintext file path: ' + plain_path)
    tmp_path = os.path.join(tmp_dir, os.path.basename(plain_path))
    logging.info('\tDecrypted file path: ' + tmp_path)
    tools.run_decrypt_file(algorithm, mode, key, cipher_path, tmp_path, iv)
    if filecmp.cmp(tmp_path, plain_path):
        return TestExitCode.SUCCESS
    else:
        logging.error('The decrypted file doesn\'t match the plaintext file')
        return TestExitCode.FAILURE

def _list_dirs(root_path):
    xs = map(lambda x: os.path.join(root_path, x), os.listdir(root_path))
    return filter(os.path.isdir, xs)

def _list_files(root_path, ext):
    xs = glob(os.path.join(root_path, '*.{}'.format(ext)))
    return filter(os.path.isfile, xs)

def _list_keys(root_path):
    return _list_files(root_path, _KEY_EXT)

def _read_line(path):
    with open(path) as f:
        return f.readline()

def _read_key(key_path):
    return _read_line(key_path)

def _read_iv(iv_path):
    return _read_line(iv_path)

def _extract_test_name(key_path):
    return os.path.splitext(os.path.basename(key_path))[0]

def _replace_ext(path, new_ext):
    return '{}.{}'.format(os.path.splitext(path)[0], new_ext)

def _build_iv_path(key_path):
    return _replace_ext(key_path, _IV_EXT)

def _build_plain_path(key_path):
    return _replace_ext(key_path, _PLAIN_EXT)

def _build_cipher_path(key_path):
    return _replace_ext(key_path, _CIPHER_EXT)

def _run_tests(tools, suite_dir, force=False):
    exit_codes = []
    suite_dir = os.path.abspath(suite_dir)
    logging.info('Suite directory path: ' + suite_dir)
    with TemporaryDirectory() as tmp_dir:
        for algorithm_dir in _list_dirs(suite_dir):
            algorithm = os.path.basename(algorithm_dir)
            maybe_algorithm = Algorithm.try_parse(algorithm)
            if maybe_algorithm is None:
                logging.warn('Unknown or unsupported algorithm: ' + algorithm)
                exit_codes.append(TestExitCode.SKIPPED)
                continue
            algorithm = maybe_algorithm
            logging.info('Algorithm: {}'.format(algorithm))
            for mode_dir in _list_dirs(algorithm_dir):
                mode = os.path.basename(mode_dir)
                maybe_mode = Mode.try_parse(mode)
                if maybe_mode is None:
                    logging.warn('Unknown or unsupported mode: ' + mode)
                    exit_codes.append(TestExitCode.SKIPPED)
                    continue
                mode = maybe_mode
                logging.info('Mode: {}'.format(mode))
                for key_path in _list_keys(mode_dir):
                    key = _read_key(key_path)
                    logging.info('Key: ' + key)
                    test_name = _extract_test_name(key_path)
                    logging.info('Test name: ' + test_name)
                    iv = None
                    if mode.requires_init_vector():
                        iv_path = _build_iv_path(key_path)
                        iv = _read_iv(iv_path)
                    plain_path = _build_plain_path(key_path)
                    cipher_path = _build_cipher_path(key_path)
                    os.makedirs(os.path.join(tmp_dir, str(algorithm), str(mode)), 0o777, True)
                    try:
                        exit_codes.append(_run_encryption_test(
                            tools, os.path.join(tmp_dir, str(algorithm), str(mode)),
                            algorithm, mode, key, plain_path, cipher_path, iv, force))
                    except Exception as e:
                        logging.error('Encountered an exception!')
                        logging.exception(e)
                        exit_codes.append(TestExitCode.ERROR)
                    if not force:
                        try:
                            exit_codes.append(_run_decryption_test(
                                tools, os.path.join(tmp_dir, str(algorithm), str(mode)),
                                algorithm, mode, key, cipher_path, plain_path, iv))
                        except Exception as e:
                            logging.error('Encountered an exception!')
                            logging.exception(e)
                            exit_codes.append(TestExitCode.ERROR)
    logging.info('Test exit codes:')
    logging.info('\tSkipped:   {}'.format(exit_codes.count(TestExitCode.SKIPPED)))
    logging.info('\tError(s):  {}'.format(exit_codes.count(TestExitCode.ERROR)))
    logging.info('\tSucceeded: {}'.format(exit_codes.count(TestExitCode.SUCCESS)))
    logging.info('\tFailed:    {}'.format(exit_codes.count(TestExitCode.FAILURE)))
    if (exit_codes.count(TestExitCode.ERROR) == 0 and
            exit_codes.count(TestExitCode.FAILURE) == 0):
        sys.exit()
    else:
        sys.exit(1)

def _build_default_log_path():
    return datetime.now().strftime('{}_%Y-%m-%d_%H-%M-%S.log').format(
        os.path.splitext(os.path.basename(__file__))[0])

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--path', '-p', nargs='*',
                        help='set path to file encryption utilities')
    parser.add_argument('--sde', '-e', action='store_true',
                        help='use Intel SDE to run *.exe files')
    parser.add_argument('--log', '-l', default=_build_default_log_path(),
                        help='set log file path')
    parser.add_argument('--force', '-f', action='store_true',
                        help='overwrite ciphertext files')
    parser.add_argument('--suite', '-s', default='file',
                        help='set test suite directory path')
    args = parser.parse_args()

    logging.basicConfig(filename=args.log,
                        format='%(asctime)s | %(module)s | %(levelname)s | %(message)s',
                        level=logging.DEBUG)

    tools = Tools(args.path, use_sde=args.sde)
    _run_tests(tools, args.suite, args.force)