97 lines
2.8 KiB
Python
97 lines
2.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||
|
|
import os
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
def setup(key_file):
|
||
|
|
f = open(key_file,"wb")
|
||
|
|
key_bytes = os.urandom(32)
|
||
|
|
f.write(key_bytes)
|
||
|
|
f.close()
|
||
|
|
|
||
|
|
def encrypt(input_file, key_file):
|
||
|
|
fkey = open(key_file,"rb")
|
||
|
|
inp = open(input_file,"rb")
|
||
|
|
out = open(f"{input_file}.enc","wb")
|
||
|
|
|
||
|
|
key = fkey.read()
|
||
|
|
plaintext = inp.read()
|
||
|
|
print(f"plaintext len : {len(plaintext)}")
|
||
|
|
iv = os.urandom(16)
|
||
|
|
# FIX: block size for aes must be 16 bytes
|
||
|
|
# plaintext needs padding
|
||
|
|
cipher = Cipher(algorithms.AES(key),modes.CBC(iv))
|
||
|
|
encryptor = cipher.encryptor()
|
||
|
|
ciphertext = encryptor.update(plaintext)
|
||
|
|
print(f"plaintext len : {len(plaintext)}")
|
||
|
|
print(f"ciphertext len : {len(ciphertext)}")
|
||
|
|
print(f"iv len : {len(iv)}")
|
||
|
|
ciphertext = iv + ciphertext
|
||
|
|
|
||
|
|
out.write(ciphertext)
|
||
|
|
|
||
|
|
fkey.close()
|
||
|
|
inp.close()
|
||
|
|
out.close()
|
||
|
|
|
||
|
|
def decrypt(input_file,key_file):
|
||
|
|
fkey = open(key_file,"rb")
|
||
|
|
inp = open(f"{input_file}","rb")
|
||
|
|
out = open(f"{input_file}.dec","wb")
|
||
|
|
|
||
|
|
key = fkey.read()
|
||
|
|
input_bytes = inp.read()
|
||
|
|
iv = input_bytes[:16]
|
||
|
|
ciphertext = input_bytes[16:]
|
||
|
|
|
||
|
|
# FIX: block size for aes must be 16 bytes
|
||
|
|
# plaintext needs padding
|
||
|
|
cipher = Cipher(algorithms.AES(key=key),mode=modes.CBC(iv))
|
||
|
|
decryptor = cipher.decryptor()
|
||
|
|
plaintext = decryptor.update(ciphertext)
|
||
|
|
|
||
|
|
out.write(plaintext)
|
||
|
|
|
||
|
|
fkey.close()
|
||
|
|
inp.close()
|
||
|
|
out.close()
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Program to perform operations using AES cipher on files",
|
||
|
|
)
|
||
|
|
|
||
|
|
subparsers = parser.add_subparsers(dest="operation", help="Operation to perform")
|
||
|
|
|
||
|
|
# Setup subcommand
|
||
|
|
setup_parser = subparsers.add_parser("setup", help="Setup a key file")
|
||
|
|
setup_parser.add_argument("fkey", help="File to contain the appropriate key for the AES cipher")
|
||
|
|
|
||
|
|
# Encrypt subcommand
|
||
|
|
enc_parser = subparsers.add_parser("enc", help="Encrypt a file")
|
||
|
|
enc_parser.add_argument("fich", help="File to be encrypted")
|
||
|
|
enc_parser.add_argument("fkey", help="File containing the key for the AES cipher")
|
||
|
|
|
||
|
|
# Decrypt subcommand
|
||
|
|
dec_parser = subparsers.add_parser("dec", help="Decrypt a file")
|
||
|
|
dec_parser.add_argument("fich", help="File to be decrypted")
|
||
|
|
dec_parser.add_argument("fkey", help="File containing the key for the AES cipher")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
match args.operation:
|
||
|
|
case "setup":
|
||
|
|
key_file = args.fkey
|
||
|
|
setup(key_file)
|
||
|
|
case "enc":
|
||
|
|
input_file = args.fich
|
||
|
|
key_file = args.fkey
|
||
|
|
encrypt(input_file,key_file)
|
||
|
|
case "dec":
|
||
|
|
input_file = args.fich
|
||
|
|
key_file = args.fkey
|
||
|
|
decrypt(input_file,key_file)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|