diff --git a/encryptedpickle/encryptedpickle.py b/encryptedpickle/encryptedpickle.py index efb1701..99afa6c 100644 --- a/encryptedpickle/encryptedpickle.py +++ b/encryptedpickle/encryptedpickle.py @@ -12,11 +12,11 @@ from struct import pack, unpack from collections import namedtuple -import simplejson as json +import json from pbkdf2 import PBKDF2 -from Crypto.Cipher import AES -from Crypto.Random import get_random_bytes -from Crypto.Hash import HMAC, SHA, SHA256, SHA384, SHA512 +from Cryptodome.Cipher import AES +from Cryptodome.Random import get_random_bytes +from Cryptodome.Hash import HMAC, SHA, SHA256, SHA384, SHA512 from .utils import ( const_equal, @@ -345,7 +345,7 @@ def _encode(self, data, algorithm, key=None): '''Encode data with specific algorithm''' if algorithm['type'] == 'hmac': - return data + self._hmac_generate(data, algorithm, key) + return data + str(self._hmac_generate(data, algorithm, key)) elif algorithm['type'] == 'aes': return self._aes_encrypt(data, algorithm, key) elif algorithm['type'] == 'no-serialization': @@ -400,7 +400,7 @@ def _sign_data(self, data, options): data = self._encode(data, algorithm, key) - return data + key_salt + return data + str(key_salt) def _unsign_data(self, data, options): '''Verify and remove signature''' @@ -561,8 +561,8 @@ def _add_magic(self, data): '''Add magic''' if self.magic: - return self.magic + data - + return self.magic + str(data) + return data def _add_header(self, data, options): @@ -575,7 +575,7 @@ def _add_header(self, data, options): flags = options['flags'] header_flags = dict( - (i, str(int(j))) for i, j in options['flags'].iteritems()) + (i, str(int(j))) for i, j in options['flags'].items()) header_flags = ''.join(version_info['flags'](**header_flags)) header_flags = int(header_flags, 2) options['flags'] = header_flags @@ -722,7 +722,7 @@ def _hmac_generate(data, algorithm, key): digestmod = EncryptedPickle._get_hashlib(algorithm['subtype']) - return HMAC.new(key, data, digestmod).digest() + return HMAC.new(key, data.encode('utf-8'), digestmod).digest() @staticmethod def _aes_encrypt(data, algorithm, key): @@ -749,7 +749,7 @@ def _aes_encrypt(data, algorithm, key): numpad = block_size - (len(data) % block_size) data = data + numpad * chr(numpad) - enc = AES.new(key, mode, iv_value).encrypt(data) + enc = AES.new(key, mode, iv_value).encrypt(data.encode('utf-8')) if include_iv: enc = iv_value + enc diff --git a/encryptedpickle/utils.py b/encryptedpickle/utils.py index bab96ad..817ad30 100644 --- a/encryptedpickle/utils.py +++ b/encryptedpickle/utils.py @@ -7,12 +7,14 @@ from __future__ import absolute_import from base64 import urlsafe_b64encode, urlsafe_b64decode +import warnings def urlsafe_nopadding_b64encode(data): '''URL safe Base64 encode without padding (=)''' - - return urlsafe_b64encode(data).rstrip('=') + + warnings.warn(type(data)) + return urlsafe_b64encode(data.encode('utf-8')).rstrip('=') def urlsafe_nopadding_b64decode(data): '''URL safe Base64 decode without padding (=)''' diff --git a/requirements.txt b/requirements.txt index a4b9885..9f8d82b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ pbkdf2>=1.3 -pycrypto>=2.6 +pycryptodomex>=3.9.8 simplejson