Updated Crypto Lab 3

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: John Samuels
# Date: 13/11/19
# Description: This is my solution for lab 3 of Applications of Cryptography. The main difference between this and the model solution is I've written it for python 3,
# and I'm using the pycryptodome library: https://pycryptodome.readthedocs.io/en/latest/

import argparse
from Crypto.Cipher import AES
from Crypto.Hash import SHA256

block_size = 16
key_size = 16

def LengthError():
    pass

def PaddingError():
    pass

def add_padding(message):
    padding_length = block_size - (len(message) % 16)

    if padding_length == 0:
        padding_length = 16

    padding = [padding_length - 1 for _ in range(padding_length)] # Removed chr()
    return message + bytes(padding)

def remove_padding(message):
    padding_byte = message[-1]
    padding_length = padding_byte + 1

    if padding_length > block_size:
        raise LengthError()

    for i in range(padding_length):
        if message[len(message)-i-1] != padding_byte:
            raise PaddingError()

    plaintext_length = len(message) - padding_length
    return message[:plaintext_length]

def password_to_key(password):
    hash = SHA256.new()
    hash.update(str.encode(password))
    key = hash.digest()
    return key[:16]


def encrypt_file(output_filename, input_filename, password):
    key = password_to_key(password)

    crypto_engine = AES.new(key, AES.MODE_CBC)
    
    with open(input_filename, "rb") as input_filehandle:
        with open(output_filename, "wb") as output_filehandle:
            plaintext = input_filehandle.read()
            ciphertext = crypto_engine.encrypt(add_padding(plaintext))
            output_filehandle.write(crypto_engine.iv + ciphertext)


def decrypt_file(output_filename, input_filename, password):
    key = password_to_key(password)
    with open(input_filename, "rb") as input_filehandle:
        with open(output_filename, "wb") as output_filehandle:
            iv_and_ciphertext = input_filehandle.read()
            iv = iv_and_ciphertext[:block_size]
            ciphertext = iv_and_ciphertext[block_size:]

            crypto_engine = AES.new(key, AES.MODE_CBC, IV=iv)
            plaintext_with_padding = crypto_engine.decrypt(ciphertext)
            plaintext = remove_padding(plaintext_with_padding)
            output_filehandle.write(plaintext)

def parse_commandline():
    parser = argparse.ArgumentParser(description='Encrypt all the files')
    parser.add_argument('--encrypt', '-e', action='store_true', help='encrypt')
    parser.add_argument('--decrypt', '-d', action='store_true', help='decrypt')
    parser.add_argument('--password', '-p', type=str, help='decrypt', required='true')
    parser.add_argument('--output', '-o', type=str, help='output filename')
    parser.add_argument('--input', '-i', type=str, help='input filename')

    args = parser.parse_args()

    if not (args.encrypt or args.decrypt):
        raise ValueError("Either encrypt with -e or decrypt with -d")

    return args

def main():
    args = parse_commandline()
    if args.encrypt:
        encrypt_file(args.output, args.input, args.password)
    else:
        decrypt_file(args.output, args.input, args.password)

if __name__ == '__main__':
    main()