derive test skeleton for disaster rescue mode

[python-delta-tar] / file_crypt.py

#!/usr/bin/env python3

""" Encrypt a single file

low-cost quick-devel badly-documented

.. codeauthor:: Intra2net <info@intra2net.com>
"""

import sys
from deltatar import crypto
from deltatar import tarfile
from traceback import print_exc


CRYPTO_MODE_ENCRYPT = 0
CRYPTO_MODE_DECRYPT = 1


def initialize_encryption (mode, password=None, key=None, nacl=None):
    if key is None and password is None:
        raise Exception ("encryption requires either a key or a password")
    if mode == CRYPTO_MODE_ENCRYPT:
        return crypto.Encrypt (password=password,
                               key=key,
                               nacl=nacl,
                               version=crypto.PDTCRYPT_DEFAULT_VER,
                               paramversion=crypto.PDTCRYPT_DEFAULT_PVER)
    if mode == CRYPTO_MODE_DECRYPT:
        return crypto.Decrypt (password=password, key=key)


def main(do_encrypt, in_file, out_file, password, comptype='gz', enctype='aes',
         bufsize=tarfile.RECORDSIZE, encoding='UTF-8'):
    """ Main function, called when running file as script

    see module doc for more info
    """

    write_handle = read_handle = None
    return_code = 4

    try:
        # open file to read
        if do_encrypt:
            read_handle = open(in_file, 'rt')
        else:
            decryptor = initialize_encryption (CRYPTO_MODE_DECRYPT,
                                               password=password)
            read_handle = tarfile._Stream(name=in_file, mode="r",
                                          comptype=comptype, bufsize=bufsize,
                                          fileobj=None, encryption=decryptor)
        return_code = 3

        # open file to write
        if do_encrypt:
            encryptor = initialize_encryption (CRYPTO_MODE_ENCRYPT,
                                               password=password)
            write_handle = tarfile._Stream(name=out_file, mode='w',
                                           comptype=comptype, bufsize=bufsize,
                                           fileobj=None, encryption=encryptor)
        else:
            write_handle = open(out_file, 'wt')
        return_code = 1

        # convert
        while True:
            buf = read_handle.read(bufsize)
            print('.', end='')
            if do_encrypt:
                write_handle.write(buf.encode(encoding, errors='strict'))
            else:
                write_handle.write(buf.decode(encoding, errors='replace'))
            if len(buf) < bufsize:
                if do_encrypt:
                    print('successfully encrypted {} into {}'
                          .format(in_file, out_file))
                else:
                    print('successfully decrypted {} into {}'
                          .format(in_file, out_file))
                break   # reached EOF

        return_code = 0

    except Exception:
        print('error encrypting file')
        print_exc()

    finally:
        # close everything
        if write_handle:
            try:
                write_handle.close()
            except Exception:
                return_code += 8
                print('error closing out file')
                print_exc()

        if read_handle:
            try:
                read_handle.close()
            except Exception:
                return_code += 16
                print('error closing in file')
                print_exc()

    # done
    return return_code


if __name__ == '__main__':
    if len(sys.argv) < 4:
        print('file_crypt.py {enc|dec} [-c] INFILE OUTFILE PASSWORD')
        sys.exit(2)
    elif sys.argv[1] == '-h':
        print('encrypt_file.py {enc|dec} [-c] INFILE OUTFILE PASSWORD')
        sys.exit(0)

    comptype = 'tar'  # auto-detect, not sure whether this would work
    idx = 1
    if sys.argv[idx] == '-c':
        print('assuming gzip compression')
        comptype = 'gz'
        idx += 1

    do_encrypt = None
    if sys.argv[idx] == 'enc':
        do_encrypt = True
    elif sys.argv[idx] == 'dec':
        do_encrypt = False
    else:
        print('encrypt_file.py {enc|dec} [-c] INFILE OUTFILE PASSWORD')
        sys.exit(2)

    if len(sys.argv) != idx+4:
        print('encrypt_file.py {enc|dec} [-c] INFILE OUTFILE PASSWORD')
        sys.exit(2)

    sys.exit(main(do_encrypt, sys.argv[idx+1], sys.argv[idx+2],
                  sys.argv[idx+3], comptype))