#!/usr/bin/python3
import base64
import sys
import hashlib
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from jinja2 import Template
from jinja2 import Environment, FileSystemLoader
def encrypt_aes_cbc(file_path, key):
with open(file_path, 'rb') as file:
data = file.read()
iv = os.urandom(16)
cipher = AES.new(key, AES.MODE_CBC, iv)
encrypted_data = cipher.encrypt(pad(data, AES.block_size))
return base64.b64encode(iv + encrypted_data)
def encrypt_aes_cbc_bin(data, key):
iv = os.urandom(16)
cipher = AES.new(key, AES.MODE_CBC, iv)
assert len(data) == AES.block_size
encrypted_data = cipher.encrypt(data)
return base64.b64encode(iv + encrypted_data)
def render_template(template_file, output_file, context):
template_str = '''
Encrypted {{ filename }}
'''
template = Template(template_str)
with open(output_file, 'w') as file:
file.write(template.render(context))
def render_template_from_file(template_file, output_file, context):
env = Environment(loader=FileSystemLoader('.'))
template = env.get_template(template_file)
with open(output_file, 'w') as file:
file.write(template.render(context))
if __name__ == "__main__":
if len(sys.argv) not in [3, 4, 5]:
print("Usage: **.py [output_file] [output_template]")
sys.exit(1)
input_file = sys.argv[1]
key = sys.argv[2].encode()
if len(sys.argv) >= 4:
if sys.argv[3] == "%":
output_html = input_file + "-enc.html"
else:
output_html = sys.argv[3]
elif len(sys.argv) == 3:
output_html = input_file + "-enc.html"
# if len(key) not in (16, 24, 32): # AES key must be 128, 192, or 256 bits
# print("Error: The key must be 16, 24, or 32 characters long.")
# sys.exit(1)
try:
key = hashlib.sha256(key).digest()
except Exception as e:
print(f"Key hashing has failed.\n{e}")
sys.exit()
encrypted_base64 = encrypt_aes_cbc(input_file, key)
encrypted_data_url = f"data:application/octet-stream;base64,{encrypted_base64.decode()}"
verify_bit = os.urandom(AES.block_size)
verify_base64 = base64.b64encode(verify_bit)
encrypted_verify = encrypt_aes_cbc_bin(verify_bit, key)
verify_data = f"{verify_base64.decode()}:{encrypted_verify.decode()}"
filename = os.path.basename(input_file)
context = {
'encrypted_data': encrypted_data_url,
'filename': filename,
'verify_data': verify_data
}
if len(sys.argv) == 5:
render_template_from_file(sys.argv[4], output_html, context)
else:
render_template('', output_html, context)
print(f"Encryption complete {output_html}")