AES File encrypting tool

aes-util

#!/Library/Frameworks/Python.framework/Versions/3.5/bin/python3.5
"""
AES-256 CBC-MODE File encrypting tool
:author: rinzler
:date: 2018-01-27
:version: 1.0
:license: GPL
usage: aes-util [encrypt|decrypt] [input file] (output file)
"""

import os
import sys
import base64
import hashlib
import getpass

from Crypto.Cipher import AES

class AesCrypt(object):
    def __init__(self, key):
        self.bs = 32
        self.key = hashlib.sha256(AesCrypt.str_to_bytes(key)).digest()

    @staticmethod
    def str_to_bytes(data):
        u_type = type(b''.decode('utf8'))
        if isinstance(data, u_type):
            return data.encode('utf8')
        return data

    def _pad(self, s):
        return s + (self.bs - len(s) % self.bs) * AesCrypt.str_to_bytes(chr(self.bs - len(s) % self.bs))

    @staticmethod
    def _unpad(s):
        return s[:-ord(s[len(s)-1:])]

    def encrypt(self, raw):
        raw = self._pad(AesCrypt.str_to_bytes(raw))
        iv = os.urandom(AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return (iv + cipher.encrypt(raw))

    def decrypt(self, enc):
        enc = (enc)
        iv = enc[:AES.block_size]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return self._unpad(cipher.decrypt(enc[AES.block_size:]))


def read_file(file_path: str) -> bytes or False:
    data = False
    if os.path.isfile(file_path):
        with open(file_path, 'rb') as fr:
            data = fr.read()
            fr.close()
    return data


def write_file(file_path: str, data: bytes) -> True:
    with open(file_path, 'wb') as fw:
        fw.write(data)
        fw.close()
    return True

if __name__ == '__main__':
    try:
        action = sys.argv[1]
        file_in = read_file(sys.argv[2])
        file_out = os.path.realpath(sys.argv[3]) if len(sys.argv) > 3 else False

        if file_in is False:
            print("Input file doesn't exist or couldn't be read.")
            sys.exit(4)

        passwd = getpass.getpass("Passphrase: ")
        if len(passwd) < 1:
            print("Error: passphrase must be at least 1 char length...")
            sys.exit(3)

        if action == 'encrypt':
            message = file_in
            encrypted = AesCrypt(key=passwd).encrypt(message)
            if file_out:
                write_file(file_out, encrypted)
                print("Encrypted %d bytes." % len(encrypted))
                print("SHA512 %s" % hashlib.sha512(encrypted).hexdigest())
                print("SHA256 %s" % hashlib.sha256(encrypted).hexdigest())
                print("SHA224 %s" % hashlib.sha224(encrypted).hexdigest())
                print("SHA1 %s" % hashlib.sha1(encrypted).hexdigest())
                print("MD5 %s" % hashlib.md5(encrypted).hexdigest())
                print("File saved at %s" % file_out)
                sys.exit(0)
            else:
                print("Encrypted %d bytes: %s" % (len(encrypted), encrypted))
                sys.exit(0)
        elif action == 'decrypt':
            message = file_in
            decrypted = AesCrypt(key=passwd).decrypt(message)
            if file_out:
                write_file(file_out, decrypted)
                print("Decrypted %d bytes." % len(decrypted))
                print("SHA512 %s" % hashlib.sha512(decrypted).hexdigest())
                print("SHA256 %s" % hashlib.sha256(decrypted).hexdigest())
                print("SHA224 %s" % hashlib.sha224(decrypted).hexdigest())
                print("SHA1 %s" % hashlib.sha1(decrypted).hexdigest())
                print("MD5 %s" % hashlib.md5(decrypted).hexdigest())
                print("File saved at %s" % file_out)
                sys.exit(0)
            else:
                print("Decrypted %d bytes: %s" % (len(decrypted), decrypted.decode()))
                sys.exit(0)
        else:
            print("Error: usage aes-util [encrypt|decrypt] [input file] (output file)")
            sys.exit(6)
    except UnicodeDecodeError:
        print("Error: bad passphrase. =/")
        sys.exit(1)
    except IndexError:
        print("Error: usage aes-util [encrypt|decrypt] [input file] (output file)")
        sys.exit(1)
    except KeyboardInterrupt:
        print("User shutdown.")
        sys.exit(2)