6
\$\begingroup\$

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
                    on password complexity.

     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


  def get_password(self):
    self.password = input("Password: ")


  def derive_key_from_password(self):
    self.salt = "lCqsDNHWCWIVZWrMZNSJnJUvltiYPell"
    self.key = crypt(self.password, self.salt, 100000) 


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:
            chunk = infile.read(chunk_size)
            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:
            chunk = infile.read(chunk_size)
            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: ")
    if menu in ["1", "2"]:
      break

  if menu == "1":
    print("\n")
    derive = DeriveKey()
    derive.get_password()

    print("Deriving key...")
    derive.derive_key_from_password()
    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()
    derive.get_password()

    print("Deriving key...")
    derive.derive_key_from_password()
    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?
\$\endgroup\$
5
\$\begingroup\$

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.
\$\endgroup\$
  • \$\begingroup\$ 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? \$\endgroup\$ – Daniel May 27 '17 at 9:11
2
\$\begingroup\$

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).

\$\endgroup\$
  • \$\begingroup\$ Thanks for helping me with that. Is curses always available as module? Platform-dependent? \$\endgroup\$ – Daniel May 25 '17 at 17:12
  • 1
    \$\begingroup\$ 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. \$\endgroup\$ – Toby Speight May 25 '17 at 17:26
  • 1
    \$\begingroup\$ What about termcolor ? \$\endgroup\$ – ChatterOne May 26 '17 at 14:56
  • \$\begingroup\$ @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. \$\endgroup\$ – Toby Speight May 26 '17 at 15:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.