file_encryption

encrypt_decrypt.py

#!/usr/bin/env python
import os
import fnmatch
from Crypto.PublicKey import RSA 
from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES, PKCS1_OAEP


def encrypt_file(data, enc_filename):
    try:
        data = data.encode('utf8')
    except:
        pass
    with open(enc_filename, 'wb') as enc_file:
        session_key = get_random_bytes(16)

        # Encrypt the session key with the public RSA key
        cipher_rsa = PKCS1_OAEP.new(key)
        enc_session_key = cipher_rsa.encrypt(session_key)

        # Encrypt the data with the AES session key
        cipher_aes = AES.new(session_key, AES.MODE_EAX)
        ciphertext, tag = cipher_aes.encrypt_and_digest(data)
        [ enc_file.write(x) for x in (enc_session_key, cipher_aes.nonce, tag, ciphertext) ]

def decrypt_file(enc_filename):
    with open(enc_filename, 'rb') as enc_file:
        enc_session_key, nonce, tag, ciphertext = [ enc_file.read(x) for x in (key.size_in_bytes(), 16, 16, -1) ]

        # Decrypt the session key with the private RSA key
        cipher_rsa = PKCS1_OAEP.new(key)
        session_key = cipher_rsa.decrypt(enc_session_key)

        # Decrypt the data with the AES session key
        cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
        data = cipher_aes.decrypt_and_verify(ciphertext, tag)
        try:
            data = data.decode('utf8')
        except:
            pass
        return data

if __name__ == '__main__':
    import sys
    key = RSA.import_key(open(os.environ['HOME'] + '/.ssh/id_rsa').read())
    doc_type = '.txt'
    enc_filename = 'encrypted_data.bin'
    try:
        target_dir = sys.argv[1]
    except:
        target_dir = os.environ['HOME']
    for rt, dr, fl in os.walk(target_dir):
      for filename in fnmatch.filter(fl, '*%s' % doc_type):
            doc = os.path.join(rt, filename)
            with open(doc, 'rb') as infile:
                data = infile.read()
            encrypt_file(data, enc_filename)
            print(decrypt_file(enc_filename))