1# The MIT License (MIT)
2#
3# Copyright (c) 2014 Richard Moore
4#
5# Permission is hereby granted, free of charge, to any person obtaining a copy
6# of this software and associated documentation files (the "Software"), to deal
7# in the Software without restriction, including without limitation the rights
8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9# copies of the Software, and to permit persons to whom the Software is
10# furnished to do so, subject to the following conditions:
11#
12# The above copyright notice and this permission notice shall be included in
13# all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
21# THE SOFTWARE.
22
23
24import sys
25sys.path.append('../pyaes')
26
27from pyaes import *
28
29import os, time
30
31# Python 3 doesn't have xrange and returns bytes from urandom
32try:
33    xrange
34except NameError:
35    xrange = range
36else:
37    pass
38
39# compare against a known working implementation
40from Crypto.Cipher import AES as KAES
41from Crypto.Util import Counter as KCounter
42for mode in [ 'CBC', 'CTR',  'CFB', 'ECB', 'OFB' ]:
43
44    (tt_ksetup, tt_kencrypt, tt_kdecrypt) = (0.0, 0.0, 0.0)
45    (tt_setup, tt_encrypt, tt_decrypt) = (0.0, 0.0, 0.0)
46    count = 0
47
48    for key_size in (128, 192, 256):
49
50        for test in xrange(1, 8):
51            key = os.urandom(key_size // 8)
52
53            if mode == 'CBC':
54                iv = os.urandom(16)
55                plaintext = [ os.urandom(16) for x in xrange(0, test) ]
56
57                t0 = time.time()
58                kaes = KAES.new(key, KAES.MODE_CBC, IV = iv)
59                kaes2 = KAES.new(key, KAES.MODE_CBC, IV = iv)
60                tt_ksetup += time.time() - t0
61
62                t0 = time.time()
63                aes = AESModeOfOperationCBC(key, iv = iv)
64                aes2 = AESModeOfOperationCBC(key, iv = iv)
65                tt_setup += time.time() - t0
66
67            elif mode == 'CFB':
68                iv = os.urandom(16)
69                plaintext = [ os.urandom(test * 5) for x in xrange(0, test) ]
70
71                t0 = time.time()
72                kaes = KAES.new(key, KAES.MODE_CFB, IV = iv, segment_size = test * 8)
73                kaes2 = KAES.new(key, KAES.MODE_CFB, IV = iv, segment_size = test * 8)
74                tt_ksetup += time.time() - t0
75
76                t0 = time.time()
77                aes = AESModeOfOperationCFB(key, iv = iv, segment_size = test)
78                aes2 = AESModeOfOperationCFB(key, iv = iv, segment_size = test)
79                tt_setup += time.time() - t0
80
81            elif mode == 'ECB':
82                plaintext = [ os.urandom(16) for x in xrange(0, test) ]
83
84                t0 = time.time()
85                kaes = KAES.new(key, KAES.MODE_ECB)
86                kaes2 = KAES.new(key, KAES.MODE_ECB)
87                tt_ksetup += time.time() - t0
88
89                t0 = time.time()
90                aes = AESModeOfOperationECB(key)
91                aes2 = AESModeOfOperationECB(key)
92                tt_setup += time.time() - t0
93
94            elif mode == 'OFB':
95                iv = os.urandom(16)
96                plaintext = [ os.urandom(16) for x in xrange(0, test) ]
97
98                t0 = time.time()
99                kaes = KAES.new(key, KAES.MODE_OFB, IV = iv)
100                kaes2 = KAES.new(key, KAES.MODE_OFB, IV = iv)
101                tt_ksetup += time.time() - t0
102
103                t0 = time.time()
104                aes = AESModeOfOperationOFB(key, iv = iv)
105                aes2 = AESModeOfOperationOFB(key, iv = iv)
106                tt_setup += time.time() - t0
107
108            elif mode == 'CTR':
109                text_length = [None, 3, 16, 127, 128, 129, 1500, 10000, 100000, 10001, 10002, 10003, 10004, 10005, 10006, 10007, 10008][test]
110                if test < 6:
111                    plaintext = [ os.urandom(text_length) ]
112                else:
113                    plaintext = [ os.urandom(text_length) for x in xrange(0, test) ]
114
115                t0 = time.time()
116                kaes = KAES.new(key, KAES.MODE_CTR, counter = KCounter.new(128, initial_value = 0))
117                kaes2 = KAES.new(key, KAES.MODE_CTR, counter = KCounter.new(128, initial_value = 0))
118                tt_ksetup += time.time() - t0
119
120                t0 = time.time()
121                aes = AESModeOfOperationCTR(key, counter = Counter(initial_value = 0))
122                aes2 = AESModeOfOperationCTR(key, counter = Counter(initial_value = 0))
123                tt_setup += time.time() - t0
124
125            count += 1
126
127            t0 = time.time()
128            kenc = [kaes.encrypt(p) for p in plaintext]
129            tt_kencrypt += time.time() - t0
130
131            t0 = time.time()
132            enc = [aes.encrypt(p) for p in plaintext]
133            tt_encrypt += time.time() - t0
134
135            if kenc != enc:
136                print("Test: mode=%s operation=encrypt key_size=%d text_length=%d trial=%d" % (mode, key_size, len(plaintext), test))
137                raise Exception('Failed encypt test case (%s)' % mode)
138
139            t0 = time.time()
140            dt1 = [kaes2.decrypt(k) for k in kenc]
141            tt_kdecrypt += time.time() - t0
142
143            t0 = time.time()
144            dt2 = [aes2.decrypt(k) for k in kenc]
145            tt_decrypt += time.time() - t0
146
147            if plaintext != dt2:
148                print("Test: mode=%s operation=decrypt key_size=%d text_length=%d trial=%d" % (mode, key_size, len(plaintext), test))
149                raise Exception('Failed decypt test case (%s)' % mode)
150
151    better = (tt_setup + tt_encrypt + tt_decrypt) / (tt_ksetup + tt_kencrypt + tt_kdecrypt)
152    print("Mode: %s" % mode)
153    print("  Average time: PyCrypto: encrypt=%fs decrypt=%fs setup=%f" % (tt_kencrypt / count, tt_kdecrypt / count, tt_ksetup / count))
154    print("  Average time: pyaes:    encrypt=%fs decrypt=%fs setup=%f" % (tt_encrypt / count, tt_decrypt / count, tt_setup / count))
155    print("  Native better by: %dx" % better)
156
157print("All test cases passes!")
158
159