# Encrypting files in a directory (AES)

I am quite new to Python (and programming in general). I was trying to create a command-line program for encrypting my files. Please tell me what I can improve (styling, crypto, etc.)

Disclaimer: I understand that trying to write crypto-code isn't a good idea unless you really know what you're doing. This piece of code is not going to be used in any serious application.

import os
import base64

from pbkdf2 import crypt
from fernet import Fernet

r = "\033[00m"
g = "\033[92m"

class DeriveKey():
"""Uses PBKDF2's crypt() function
to derive a 32-bit key from the password,
using a salt and 100000 iterations

self.password: User input. There are no restrictions

self.salt: A previously randomly generated string,
must be *at least* 16 bytes long;
according to http://www.ietf.org/rfc/rfc2898.txt

self.key: Whatever is returned by crypt(), we use this
in EncryptFiles() / DecryptFiles()

"""
def __init__(self):
self.key = None

self.salt = "lCqsDNHWCWIVZWrMZNSJnJUvltiYPell"

class EncryptFiles():
"""Encrypt all files in a list (obtained with os.listdir())
using Fernet. Fernet uses AES-128 in CBC MoO.

self.cipher: The cipher [Fernet()] which we use to
encrypt data

self.selection: List containing all unencrypted files in
the directory.

output_file: The file to which the encrypted data is written;
this is formed by taking (filename) + ".xCrypt"
"""

def __init__(self):
self.cipher = None

def look_for_files(self, directory):
self.selection = []
for filename in os.listdir(directory):
if os.path.isfile(filename):
if not filename[len(filename)-6:len(filename)] == ".xCrypt":
# filter out already encrypted files
self.selection.append(filename)

def set_fernet_cipher(self, key):
key = key[0:32]
key = bytes(key.encode("utf-8"))
self.cipher = Fernet(base64.b64encode(key))
# Fernet requires base64-encoded keys

def encrypt_selection(self):
chunk_size = 64*1024

for input_file in self.selection:
output_file = input_file + ".xCrypt"

with open(input_file, "rb") as infile:
with open(output_file, "wb") as outfile:

while True:
if len(chunk) == 0:
break

outfile.write(self.cipher.encrypt(chunk))

return True

class DecryptFiles():
"""Decrypt all files in a list (obtained with os.listdir())
using Fernet.
Fernet uses AES-128 in CBC MoO.

self.cipher: The cipher [Fernet()] which we use to
decrypt data

self.selection: List containing all encrypted files in
the directory.

output_file: The file to which the decrypted data is written;
this is formed by taking (filename) - ".xCrypt"
"""
def __init__(self):
self.cipher = None

def look_for_files(self, directory):
self.selection = []
for filename in os.listdir(directory):
if os.path.isfile(filename):
if filename[len(filename)-7:len(filename)] == ".xCrypt":
self.selection.append(filename)

def set_fernet_cipher(self, key):
key = key[0:32]
key = bytes(key.encode("utf-8"))
self.cipher = Fernet(base64.b64encode(key))

def decrypt_selection(self):
chunk_size = 64*1024

for input_file in self.selection:
output_file = input_file[0:len(input_file)-7] # Strip ".xCrypt"

with open(input_file, "rb") as infile:
with open(output_file, "wb") as outfile:

while True:
if len(chunk) == 0:
break

outfile.write(self.cipher.decrypt(chunk))

return True

def main():
print(reset)
print("{}[1]{} Encrypt files".format(g, r))
print("{}[2]{} Decrypt files".format(g, r))

while True:
menu = input("Make a choice: ")
break

print("\n")
derive = DeriveKey()

print("Deriving key...")
print("{}Key derived.{}".format(g, r))

key = derive.key

while True:
dir_ = input("Directory: ")
if os.path.isdir(dir_):
break

encrypt = EncryptFiles()
encrypt.seek_files(dir_)
encrypt.set_cipher(key)

print("Encrypting files...")
encrypt.encrypt_selection()
print("{}Files encrypted.{}".format(g, r))

else:
print("\n")
derive = DeriveKey()

print("Deriving key...")
print("{}Key derived.{}".format(g, r))

key = derive.key

while True:
dir_ = input("Directory: ")
if os.path.isdir(dir_):
break

decrypt = DecryptFiles()
decrypt.seek_files(dir_)
decrypt.set_cipher(key)

print("Decrypting files...")
decrypt.decrypt_selection()
print("{}Files decrypted.{}".format(g, r))

if __name__=="__main__":
main()


Specific questions:

1. Should I be using an IV?
2. Is it okay to have such a messy set_fernet_cipher() function?
3. Is it okay to use the same salt multiple times? If not, how should I pass it to pbkdf2's crypt() function?

These are some crypto recommendations which I believe apply to your case (I am not familiar with Python crypto details):

• The salt should not be hard coded, but should be passed to the code or generated randomly. Reusing the salt makes it possible to generate a rainbow table to recover passwords faster. The salt appears after the third dollar sign in the response from crypt, and it appears that crypt generates a salt for you if you don't specify one (just make sure to name the iterations parameter in the call). This is how for example *nix shadow files work.
• The number of iterations should be adapted to the speed of the computer (and should be an optional argument). More iterations for faster computers (and as computers evolve to be faster), fewer iterations for slower ones.
• Use getpass instead of input to get the password. This avoids shoulder surfing.

General recommendations:

• I would initialise fields at the class level rather than in the constructor, if only to save a line.
• Use os.path.splitext to get the file extension.
• The only part of this tool that should be interactive is the bit that gets the password. Everything else should be passed on the command line and parsed using argparse. One advantage of this approach is that it makes parameter re-use easy: if you want to repeat an operation with almost or exactly the same arguments, you simply bring it up in your terminal history and change as needed, rather than having to re-type everything. Another advantage is that non-interactive methods can be trivially re-used in other code without "polluting" that code with interactive prompts. Some things *nix tools typically do:
• Files are passed to the command as unnamed arguments at the end of the argument list, for example cat foo.txt bar.txt.
• Anything else is passed as either positional arguments (as in grep pattern file-name.txt) or named arguments (as in grep --expression=pattern file-name.txt). I recommend using only named arguments because they are self-documenting and they can be provided in whatever sequence is convenient.
• Thanks for your answers. If you don't mind, I have two questions. First, how would I make the setup with a randomly generated salt work? If I generate a salt randomly [ a ] for encryption, and then generate another salt randomly [ b ], a will not equal b. So the only solution is to write the salt to a file and then load it into memory later by reading that file. Isn't that practically the same as hardcoding? My second question is: Why should the parameters such as infile / outfile, be command-line arguments? Is there any clear advantage? – Daniel May 27 '17 at 9:11

Not a full review, but this part hit me:

r = "\033[00m"
g = "\033[92m"


Those look like terminal escape codes of some sort. But you can't know what kind of terminal your output is going to, nor even whether it's going to a terminal at all. If you feel that you need them, then the portable way to do so is to import curses and use its facilities (or an external program such as tput, but really, no).

• Thanks for helping me with that. Is curses always available as module? Platform-dependent? – Daniel May 25 '17 at 17:12
• All I know is what the reference page says: "* While curses is most widely used in the Unix environment, versions are available for DOS, OS/2, and possibly other systems as well. This extension module is designed to match the API of ncurses, an open-source curses library hosted on Linux and the BSD variants of Unix.*" I've linked to it, for convenience. – Toby Speight May 25 '17 at 17:26
• What about termcolor ? – ChatterOne May 26 '17 at 14:56
• @ChatterOne, I downloaded and inspected that library, and it seems to suffer the same problem as the code posted here. So I can't recommend it as a portable interface. – Toby Speight May 26 '17 at 15:04