116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
import argparse
|
|
import hashlib
|
|
from Crypto.Cipher import AES
|
|
import sys
|
|
import os
|
|
|
|
# pip install voussoirkit
|
|
from voussoirkit import bytestring
|
|
|
|
|
|
# 128 bits
|
|
BLOCK_SIZE = 32
|
|
KEY_SIZE = 32
|
|
SEEK_END = 2
|
|
|
|
def decrypt_data(aes, data):
|
|
data = aes.decrypt(data)
|
|
pad_byte = data[-1:]
|
|
data = data.rstrip(pad_byte)
|
|
return data
|
|
|
|
def encrypt_data(aes, data):
|
|
pad_byte = (data[-1] + 1) % 256
|
|
pad_length = BLOCK_SIZE - (len(data) % BLOCK_SIZE)
|
|
data += (bytes([pad_byte]) * pad_length)
|
|
data = aes.encrypt(data)
|
|
return data
|
|
|
|
def decrypt_file(aes, input_handle, output_handle):
|
|
current_pos = input_handle.tell()
|
|
input_size = input_handle.seek(0, SEEK_END) - current_pos
|
|
input_handle.seek(current_pos)
|
|
bytes_read = 0
|
|
while True:
|
|
chunk = input_handle.read(BLOCK_SIZE)
|
|
if len(chunk) == 0:
|
|
break
|
|
bytes_read += len(chunk)
|
|
chunk = aes.decrypt(chunk)
|
|
if bytes_read == input_size:
|
|
last_byte = chunk[-1]
|
|
while chunk and chunk[-1] == last_byte:
|
|
chunk = chunk[:-1]
|
|
if bytes_read % bytestring.MIBIBYTE == 0:
|
|
print(bytestring.bytestring(bytes_read))
|
|
output_handle.write(chunk)
|
|
|
|
def encrypt_file(aes, input_handle, output_handle):
|
|
last_byte = 0
|
|
done = False
|
|
bytes_read = 0
|
|
while not done:
|
|
chunk = input_handle.read(BLOCK_SIZE)
|
|
if len(chunk) > 0:
|
|
last_byte = chunk[-1]
|
|
if len(chunk) < BLOCK_SIZE:
|
|
pad_byte = (last_byte + 1) % 256
|
|
pad_byte = bytes([pad_byte])
|
|
chunk += pad_byte * (BLOCK_SIZE - len(chunk))
|
|
done = True
|
|
bytes_read += len(chunk)
|
|
if bytes_read % bytestring.MIBIBYTE == 0:
|
|
print(bytestring.bytestring(bytes_read))
|
|
chunk = aes.encrypt(chunk)
|
|
output_handle.write(chunk)
|
|
#print(''.join((hex(x)[2:].rjust(2, '0') for x in chunk)))
|
|
|
|
def prepare_handles_argparse(args):
|
|
return (aes, input_handle, output_handle)
|
|
|
|
def encrypt_argparse(args):
|
|
input_handle = open(args.input, 'rb')
|
|
output_handle = open(args.output, 'wb')
|
|
|
|
password = hashit(args.password)
|
|
initialization_vector = os.urandom(16)
|
|
aes = AES.new(password, mode=3, IV=initialization_vector)
|
|
output_handle.write(initialization_vector)
|
|
|
|
encrypt_file(aes, input_handle, output_handle)
|
|
|
|
def decrypt_argparse(args):
|
|
input_handle = open(args.input, 'rb')
|
|
output_handle = open(args.output, 'wb')
|
|
|
|
password = hashit(args.password)
|
|
initialization_vector = input_handle.read(16)
|
|
aes = AES.new(password, mode=3, IV=initialization_vector)
|
|
decrypt_file(aes, input_handle, output_handle)
|
|
|
|
def hashit(text):
|
|
h = hashlib.sha512(text.encode('utf-8')).hexdigest()
|
|
h = h[:BLOCK_SIZE]
|
|
return h
|
|
|
|
def main(argv):
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers()
|
|
|
|
p_encrypt = subparsers.add_parser('encrypt')
|
|
p_encrypt.add_argument('-i', '--input', dest='input', required=True)
|
|
p_encrypt.add_argument('-o', '--output', dest='output', required=True)
|
|
p_encrypt.add_argument('-p', '--password', dest='password', required=True)
|
|
p_encrypt.set_defaults(func=encrypt_argparse)
|
|
|
|
p_decrypt = subparsers.add_parser('decrypt')
|
|
p_decrypt.add_argument('-i', '--input', dest='input', required=True)
|
|
p_decrypt.add_argument('-o', '--output', dest='output', required=True)
|
|
p_decrypt.add_argument('-p', '--password', dest='password', required=True)
|
|
p_decrypt.set_defaults(func=decrypt_argparse)
|
|
|
|
args = parser.parse_args(argv)
|
|
return args.func(args)
|
|
|
|
if __name__ == '__main__':
|
|
raise SystemExit(main(sys.argv[1:]))
|