File's encryption tool in terminal

Question

I've created a wrapper for the SimpleCrypt library that allows to easily encrypt/decrypt a string.

The application is used this :

//To encrypt
python encryptor.py file_name -e

//To decrypt
python encryptor.py file_name -d

Basically, to encrypt, I read the file, encrypt it and create a new file with the extension .enc, so that it's obvious (to me) it's been encrypted.

I'm not looking for a super secure application, I'm tutoring and when I leave my laptop alone I wouldn't want students to check their future exams or stuff like that. So it's more of a small project for my own fun, that protects files from non IT people.

I think I've covered every cases (Permission denied, file doesn't exist, wrong password etc).

I think the code could look better, but I'm mostly interested in knowing if I missed some bug I'm not aware of (As I said, I tested the application to the best of my knowledge) or if there are some Python basics that I missed.

encryptor.py

import sys
import os.path
import getpass
import simplecrypt

ReadMode = "r"
WriteMode = "w"
EncryptArg = "-e"
DecryptArg = "-d"
PermissionDenied = 13
FileDoesNotExist = 2


def main():
    if not len(sys.argv) == 3:
        print "Incorrect number of arguments."
        print "Expected file-path encrypt/decrypt"
        exit(1)

    filepath = sys.argv[1]
    encrypt_or_decrypt = sys.argv[2]

    if not (encrypt_or_decrypt == EncryptArg or encrypt_or_decrypt == DecryptArg):
        print "2nd argument provided is invalid, expected -e or -d, received : " + encrypt_or_decrypt
        exit(1)

    requires_encryption = encrypt_or_decrypt == EncryptArg

    try:
        stream = open(filepath, ReadMode)
        file_content = stream.read()
        stream.close()

        if requires_encryption:
            password = getpass.getpass("Password to decrypt : ")
            encrypted_content = simplecrypt.encrypt(password, file_content)

            write_stream = open(filepath + ".enc", WriteMode)
            write_stream.write(encrypted_content)
            write_stream.close()
        else:
            if not filepath[-4:0] == ".enc":
                print "File was not encrypted using this application. Expected .enc format."
                exit(1)

            decrypted_content = None
            try:
                password = getpass.getpass("Password : ")
                decrypted_content = simplecrypt.decrypt(password, file_content)
            except simplecrypt.DecryptionException:
                print "Password is invalid."
                exit(1)

            write_stream = open(filepath[:-4], WriteMode)
            write_stream.write(decrypted_content)
            stream.close()

        os.remove(filepath)

    except IOError, ioError:
        if ioError.errno == PermissionDenied:
            print "You do not have permission to access this file. Consider using sudo"
        elif ioError.errno == FileDoesNotExist:
            print "File does not exist."
        else:
            raise ioError

        exit(1)

if __name__ == '__main__':
    main()

Answer 1

• Your arguments are pretty static, and hard to add to. Lets say, you sometimes don't want to cleanup after yourself. Adding a -c flag would be nice. With most *nix command line applications you could do this with something like encryptor -ec path. And so I'd highly recommend using argparse. With something like:

  def make_parser():
      parser = argparse.ArgumentParser(prog='encryptor')
      parser.add_argument('file_path', type=str,
                          help='Path to file to encrypt or decrypt.')
      parser.add_argument('-c', '--cleanup', action='store_false',
                          help='Stop program cleanup.')
      group = parser.add_mutually_exclusive_group(required=True)
      group.add_argument('-d', '--decrypt', action='store_true',
                         help='Specify program should run in decrypt mode.')
      group.add_argument('-e', '--encrypt', action='store_true',
                         help='Specify program should run in encrypt mode.')
      return parser
  

• Constants in Python are UPPER_SNAKE_CASE, and so I'd recommend you change ReadMode etc. to this format. Currently I'd read that variable to be a class.

• Everyone that uses Python knows the arguments to open, open(..., 'r') is much simpler to read than open(..., READ_MODE). And when not using 'r', the variable name will only get worse as time progresses - abstract to a common constants module, rename to better specify for files, etc.

• Please use with, I've previously written a pretty generic answer why you should do this, however it mostly comes down to, you're not using it correctly. Instead you need to use one of:

  # Option 1
  stream = open(filepath, 'r')
  try:
      content = stream.read()
  finally:
      stream.close()
  
  # Option 2
  with open(filepath, 'r') as stream:
      content = stream.read()
  

• I'm not a fan of exit(1). Python has pretty good errors, whilst they may be more targeted at developers, they're kinda alright to give to end users too. But the worst part is, if you later want to debug why you're getting an error, you have to manually change every exit(1) from your program, to get a useful stack trace. Instead, just wrap your main, print the errors, and then exit.

  try:
      main()
  except Exception as e:
      print(str(e))
      raise SystemExit(1)
  

---

In all I'd use something like:

import argparse
import getpass
import simplecrypt


def main(args):
    if args.decrypt and not args.filepath.endswith('.enc'):
        raise ValueError('Expected .enc format file to decrypt.')

    with open(args.filepath, 'r') as stream:
        file_content = stream.read()

    password = getpass.getpass('Password: ')
    if args.encrypt:
        file_path = args.filepath + '.enc'
        content = simplecrypt.encrypt(password, file_content)
    else:
        file_path = args.filepath[:-4]
        content = simplecrypt.decrypt(password, file_content)

    with open(file_path, 'w') as write_stream:
        write_stream.write(content)

    if args.cleanup:
        os.remove(args.filepath)


def make_parser():
    parser = argparse.ArgumentParser(prog='encryptor')
    parser.add_argument('file_path', type=str,
                        help='Path to file to encrypt or decrypt.')
    parser.add_argument('-c', '--cleanup', action='store_false',
                        help='Stop program cleanup.')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-d', '--decrypt', action='store_true',
                       help='Specify program should run in decrypt mode.')
    group.add_argument('-e', '--encrypt', action='store_true',
                       help='Specify program should run in encrypt mode.')
    return parser



if __name__ == '__main__':
    args = make_parser().parse_args()

    try:
        main(args)
    except Exception as e:
        print(str(e))
        exit(1)

Answer 2

If you crack argv using argparse you'll enjoy a couple of benefits, including you get --help for free.

Defining read & WRITE_MODE is maybe a little over the top, as open(f, 'w') is a standard idiom. Anyway, you have an opportunity to let with close files for you, if you like. For example:

    with open(filepath + ".enc", "wb") as write_stream:
        write_stream.write(encrypted_content)

You probably want binary mode: "wb" rather than "w".

Consider throwing an exception in place of calling exit(1).

For the filepath[-4:0] == ".enc" enforcement test you might write: filepath.endswith(".enc").

Here is a magic number, corresponding to len(".enc"): filepath[:-4]. Consider changing the enforcement test so it uses a regex, re.compile(r'^(.+)\.enc$'). Then you can use group(1) of the match object when specifying a file path.

The ioError identifier (or io_error according to PEP8) is accurate, but perhaps a bit long -- err would suffice. Better to use the "as" keyword:

except IOError as err: