104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
import hashlib,os
|
|
from utility import read_in_chunks
|
|
|
|
|
|
try:
|
|
pycrypto = True
|
|
from Crypto.Cipher import AES
|
|
from Crypto import Random
|
|
print("Using PyCrypto for AES enc/dec")
|
|
except ImportError:
|
|
pycrypto = False
|
|
import pyaes
|
|
print("Using Pyaes for AES enc/dec")
|
|
|
|
|
|
|
|
class AESCipherFile:
|
|
def __init__( self ,f,d, size_limit_reading_os, key ):
|
|
h = hashlib.sha256()
|
|
h.update(key.encode("utf_8"))
|
|
self.key = h.digest()
|
|
self.f = f
|
|
self.d = d
|
|
self.size_limit_reading_os = size_limit_reading_os
|
|
|
|
def pad(self,b):
|
|
b = bytearray(b)
|
|
fill = 16 - len(b) % 16
|
|
for i in range(fill):
|
|
b.append(fill)
|
|
return bytes(b)
|
|
|
|
def unpad(self,b) :
|
|
return b[:-b[len(b)-1]]
|
|
|
|
|
|
def encrypt( self ):
|
|
if pycrypto:
|
|
iv = Random.new().read( 16 )
|
|
aes_encrypt_state = AES.new( self.key, AES.MODE_CBC, iv )
|
|
else:
|
|
iv = os.urandom(16)
|
|
aes_encrypt_state = pyaes.AESModeOfOperationCBC(self.key, iv = iv)
|
|
#check if the file is a multiple of 16 bytes for padding
|
|
self.f.seek(0,2)
|
|
if self.f.tell() % 16 == 0:
|
|
self.d.write(bytes(b"1"))
|
|
else:
|
|
self.d.write(bytes(b"0"))
|
|
self.f.seek(0,0)
|
|
#end check
|
|
self.d.write(iv)
|
|
hash = hashlib.md5()
|
|
for piece in read_in_chunks(self.f,self.size_limit_reading_os):
|
|
hash.update(piece)
|
|
if len(piece) % 16 > 0 :
|
|
piece = self.pad(piece)
|
|
steps = int(len(piece) / 16)
|
|
for i in range(steps):
|
|
self.d.write(aes_encrypt_state.encrypt( piece[16 * i : 16 * (i + 1)] ))
|
|
return hash.hexdigest()
|
|
|
|
def decrypt( self ):
|
|
perfect = False
|
|
perfectbyte = self.f.read(1)
|
|
if perfectbyte == bytes(b"1"):
|
|
perfect = True
|
|
if pycrypto:
|
|
aes_decrypt_state = AES.new( self.key, AES.MODE_CBC, self.f.read(16) )
|
|
else:
|
|
aes_decrypt_state = pyaes.AESModeOfOperationCBC(self.key, self.f.read(16))
|
|
last_16_bytes_in_piece = None
|
|
for piece in read_in_chunks(self.f,self.size_limit_reading_os):
|
|
if last_16_bytes_in_piece != None:
|
|
self.d.write(aes_decrypt_state.decrypt(last_16_bytes_in_piece))
|
|
steps = int(len(piece) / 16)
|
|
if steps > 1 :
|
|
for i in range(steps - 1):
|
|
self.d.write(aes_decrypt_state.decrypt( piece[16 * i : 16 * (i + 1)] ))
|
|
i = i + 1
|
|
else:
|
|
i = 0
|
|
last_16_bytes_in_piece = aes_decrypt_state.decrypt( piece[16 * i : 16 * (i + 1)] )
|
|
if last_16_bytes_in_piece != None:
|
|
if perfect:
|
|
self.d.write(last_16_bytes_in_piece)
|
|
else:
|
|
self.d.write(self.unpad(last_16_bytes_in_piece))
|
|
|
|
|
|
|
|
''' usage
|
|
with open("c:\\test\\file", 'rb') as f:
|
|
with open("c:\\test\\file_enc", 'wb') as d:
|
|
aes = AESCipherFile(f,d,134217728,"chiave")
|
|
aes.encrypt()
|
|
with open("c:\\test\\file_enc", 'rb') as f:
|
|
with open("c:\\test\\file", 'wb') as d:
|
|
aes = AESCipherFile(f,d,134217728,"chiave")
|
|
aes.decrypt()
|
|
'''
|
|
|
|
|