11
11
POSTFIX_MODE = {
12
12
'cbc' : modes .CBC ,
13
13
'cfb' : modes .CFB ,
14
- 'ecb' : modes .ECB ,
15
14
}
16
15
17
16
AES_BLOCK_SIZE = int (algorithms .AES .block_size / 8 )
18
17
19
18
20
19
class AESCipher (object ):
21
- def __init__ (self , key , iv = None ):
20
+ def __init__ (self , key ):
22
21
"""
23
22
:param key: The encryption key
24
- :param iv: Init vector
25
23
:return: AESCipher instance
26
24
"""
27
25
self .key = key
28
- self .iv = iv
29
26
30
- def build_cipher (self , iv = None , alg = 'aes_128_cbc' ):
27
+ def build_cipher (self , alg = 'aes_128_cbc' ):
31
28
"""
32
- :param iv: init vector
33
29
:param alg: cipher algorithm
34
30
:return: A Cipher instance
35
31
"""
36
32
typ , bits , cmode = alg .lower ().split ('_' )
37
33
bits = int (bits )
38
-
39
- if not iv :
40
- if self .iv :
41
- iv = self .iv
42
- else :
43
- iv = os .urandom (AES_BLOCK_SIZE )
34
+ iv = os .urandom (AES_BLOCK_SIZE )
44
35
45
36
if len (iv ) != AES_BLOCK_SIZE :
46
37
raise Exception ('Wrong iv size: {}' .format (len (iv )))
@@ -63,11 +54,10 @@ def build_cipher(self, iv=None, alg='aes_128_cbc'):
63
54
64
55
return cipher , iv
65
56
66
- def encrypt (self , msg , iv = None , alg = 'aes_128_cbc' , padding = 'PKCS#7' ,
67
- b64enc = True , block_size = AES_BLOCK_SIZE ):
57
+ def encrypt (self , msg , alg = 'aes_128_cbc' , padding = 'PKCS#7' , b64enc = True ,
58
+ block_size = AES_BLOCK_SIZE ):
68
59
"""
69
60
:param key: The encryption key
70
- :param iv: init vector
71
61
:param msg: Message to be encrypted
72
62
:param padding: Which padding that should be used
73
63
:param b64enc: Whether the result should be base64encoded
@@ -87,7 +77,7 @@ def encrypt(self, msg, iv=None, alg='aes_128_cbc', padding='PKCS#7',
87
77
c = chr (plen ).encode ()
88
78
msg += c * plen
89
79
90
- cipher , iv = self .build_cipher (iv , alg )
80
+ cipher , iv = self .build_cipher (alg )
91
81
encryptor = cipher .encryptor ()
92
82
cmsg = iv + encryptor .update (msg ) + encryptor .finalize ()
93
83
@@ -98,48 +88,18 @@ def encrypt(self, msg, iv=None, alg='aes_128_cbc', padding='PKCS#7',
98
88
99
89
return enc_msg
100
90
101
- def decrypt (self , msg , iv = None , alg = 'aes_128_cbc' , padding = 'PKCS#7' ,
102
- b64dec = True ):
91
+ def decrypt (self , msg , alg = 'aes_128_cbc' , padding = 'PKCS#7' , b64dec = True ):
103
92
"""
104
93
:param key: The encryption key
105
- :param iv: init vector
106
94
:param msg: Base64 encoded message to be decrypted
107
95
:return: The decrypted message
108
96
"""
109
97
data = b64decode (msg ) if b64dec else msg
110
98
111
- _iv = data [:AES_BLOCK_SIZE ]
112
- if iv :
113
- assert iv == _iv
114
- cipher , iv = self .build_cipher (iv , alg = alg )
99
+ cipher , iv = self .build_cipher (alg = alg )
115
100
decryptor = cipher .decryptor ()
116
101
res = decryptor .update (data )[AES_BLOCK_SIZE :] + decryptor .finalize ()
117
102
if padding in ['PKCS#5' , 'PKCS#7' ]:
118
103
idx = bytearray (res )[- 1 ]
119
104
res = res [:- idx ]
120
105
return res
121
-
122
-
123
- def run_test ():
124
- key = b'1234523451234545' # 16 byte key
125
- iv = os .urandom (AES_BLOCK_SIZE )
126
- # Iff padded, the message doesn't have to be multiple of 16 in length
127
- original_msg = b'ToBeOrNotTobe W.S.'
128
- aes = AESCipher (key )
129
-
130
- encrypted_msg = aes .encrypt (original_msg , iv )
131
- decrypted_msg = aes .decrypt (encrypted_msg , iv )
132
- assert decrypted_msg == original_msg
133
-
134
- encrypted_msg = aes .encrypt (original_msg )
135
- decrypted_msg = aes .decrypt (encrypted_msg )
136
- assert decrypted_msg == original_msg
137
-
138
- aes = AESCipher (key , iv )
139
- encrypted_msg = aes .encrypt (original_msg )
140
- decrypted_msg = aes .decrypt (encrypted_msg )
141
- assert decrypted_msg == original_msg
142
-
143
-
144
- if __name__ == '__main__' :
145
- run_test ()
0 commit comments