47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
try:
|
|
import Crypto.Hash.SHA256
|
|
import Crypto.Cipher.AES
|
|
has_aes = True
|
|
except:
|
|
has_aes = False
|
|
|
|
class CrypterInitError(Exception):
|
|
pass
|
|
|
|
class Crypter(object):
|
|
'''Wraps AES functionality provided by pycrypto'''
|
|
def __init__(self, password, salt='aes.py', alignment=16):
|
|
'''Initializes a crypter object'''
|
|
if not has_aes:
|
|
raise CrypterInitError("pycrypto not installed")
|
|
self.password = password
|
|
self.salt = salt
|
|
pdata = (password + salt).encode('UTF-8')
|
|
self.hash = Crypto.Hash.SHA256.new(pdata)
|
|
self.cipher = Crypto.Cipher.AES.new(self.hash.digest())
|
|
self.alignment = alignment
|
|
|
|
def pad(self, s):
|
|
'''Add NULL padding to create aligned data'''
|
|
xtra = len(s) % self.alignment
|
|
if xtra:
|
|
return s + b'\x00' * (self.alignment - xtra) #FIX3
|
|
else:
|
|
return s
|
|
|
|
def unpad(self, s):
|
|
'''Remove NULL padding to restore original data'''
|
|
l = len(s)
|
|
assert type(s) == type(b'')
|
|
while l > 0 and s[l - 1] == 0: #FIX3
|
|
l -= 1
|
|
return s[:l]
|
|
|
|
def encrypt(self, data):
|
|
'''Creates encrypted data string from original input'''
|
|
return self.cipher.encrypt(self.pad(data))
|
|
|
|
def decrypt(self, data):
|
|
'''Creates original data string from encrypted input'''
|
|
return self.unpad(self.cipher.decrypt(self.pad(data)))
|