crypto.py

#! /usr/bin/python

from Crypto.Cipher import AES
import hashlib
import struct
from random import randint
from getpass import getpass
import argparse
import os
import sys

#This is a basic script to encrypt/decrypt a file
#Happy hacking! -Defalt

def parsing_func():
    parser = argparse.ArgumentParser(description='Basic encryption/decryption tool')
    parser.add_argument('-e', '--encrypt', help='Encrypt specified file', action='store_true', default=False, dest='encrypt')
    parser.add_argument('-d', '--decrypt', help='Decrypt specified file', action='store_true', default=False, dest='decrypt')
    parser.add_argument('-f', '--file', help='Specify file to encrypt/decrypt', action='store', default=False, dest='file')
    parser.add_argument('-o', '--output', help='Name of output file', action='store', default=False, dest='outfile')
    args = parser.parse_args()
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    elif ((not args.encrypt) and (not args.decrypt)) or ((not not args.encrypt) and (not not args.decrypt)):
        parser.error('invalid action specification')
    elif not args.file:
        parser.error('no file specified')
    else:
        return args

args = parsing_func()

def encrypt(infile=args.file, outfile=args.outfile):
    if not os.path.isfile(infile):
        print '[!] Failed to Detect Input File'
        sys.exit(1)
    elif not outfile:
        try:
            overwrite = raw_input('[*] No Output File Given, Overwrite Input File? [y/N] ').strip().lower()[0]
        except KeyboardInterrupt:
            print '[!] Aborting Encryption'
            sys.exit(1)
        if overwrite == 'y':
            outfile = infile
        else:
            print '[!] Aborting Encryption'
            sys.exit(1)
    filesize = os.path.getsize(infile)
    print '[*] Reading File to Encrypt'
    try:
        with open(infile, 'rb') as file:
            to_encrypt = file.read()
    except IOError:
        print '[!] Failed to Read Input File'
        sys.exit(1)
    print '[*] Padding File Contents\n'
    to_encrypt += 'password_is_correct'
    while len(to_encrypt) % 16 != 0:
        to_encrypt += ' '
    try:
        password = getpass('[*] Enter Password: ')
        confirm_pass = getpass('[*] Confirm Password: ')
    except KeyboardInterrupt:
        print '[!] Aborting Encryption'
        sys.exit(1)
    if password != confirm_pass:
        print '[!] Password Confirmation Failed'
        sys.exit(1)
    print '\n[*] Encrypting Contents'
    ivector = ''.join(chr(randint(0, 0xFF)) for i in range(16))
    encrypted_file = AES.new(hashlib.sha256(password).digest(), AES.MODE_CBC, ivector).encrypt(to_encrypt)
    try:
        with open(outfile, 'wb') as file:
            file.write(struct.pack('<Q', filesize))
            file.write(ivector)
            file.write(encrypted_file)
    except IOError:
        print '[!] Failed to Encrypt File Contents'
        sys.exit(1)
    print '[*] Encryption Complete!'

def decrypt(infile=args.file, outfile=args.outfile):
    if not os.path.isfile(infile):
        print '[!] Failed to Detect Input File'
        sys.exit(1)
    elif not outfile:
        try:
            overwrite = raw_input('[*] No Output File Given, Overwrite Input File? [y/N] ').strip().lower()[0]
        except KeyboardInterrupt:
            print '[!] Aborting Encryption'
            sys.exit(1)
        if overwrite == 'y':
            outfile = infile
        else:
            print '[!] Aborting Encryption'
            sys.exit(1)
    print '[*] Parsing File to Decrypt\n'
    try:
        with open(infile, 'rb') as file:
            size = struct.unpack('<Q', file.read(8))
            ivector = file.read(16)
            to_decrypt = file.read()
    except IOError:
        print '[!] Failed to Read Input File'
        sys.exit(1)
    try:
        password = getpass('[*] Enter Password: ')
    except Exception:
        print '[!] Aborting Decryption'
        sys.exit(1)
    print '\n[*] Attempting to Decrypt File Contents'
    try:
        decrypted = AES.new(hashlib.sha256(password).digest(), AES.MODE_CBC, ivector).decrypt(to_decrypt)
        if not 'password_is_correct' in decrypted:
            print '[!] Incorrect Password'
            sys.exit(1)
    except Exception:
        print '[!] Failed to Decrypt File Contents'
        sys.exit(1)
    try:
        with open(outfile, 'wb') as file:
            file.write(decrypted)
            file.truncate(size[0])
    except IOError:
        print '[!] Failed to Write Decrypted File'
        sys.exit(1)
    print '[*] Decryption Complete!'

if args.encrypt:
    encrypt()
else:
    decrypt()