non-fancy symmetric key encryption

crypt.py

#!/usr/bin/env python3
# coding: ascii

import os
import sys
import gzip
import base64
import getpass
import threading

import argon2

from cryptography.fernet import Fernet


__requirements__ = [
    'argon2-cffi==19.1.0',
    'cryptography'
    'cffi!=1.11.3,>=1.8',
    'six>=1.4.1',
    'asn1crypto>=0.21.0',
    'pycparser',
]


HASH_LEN = 16
KEY_LEN = 32
SALT_LEN = 16
ENCODED_VERIFY_LEN = ((HASH_LEN+SALT_LEN)*4//3+3) & -4
ARGON2_PARAMS = dict(
    type        = argon2.Type.ID,
    hash_len    = HASH_LEN + KEY_LEN,
    time_cost   = 2,
    memory_cost = 1000000,
    parallelism = 8,
)


class InvalidPassword(ValueError):
    '''Results from invalid hash parameters.'''

def ensure_bytes(s):
    if isinstance(s, str):
        s = bytes(s, 'utf-8')
    if not isinstance(s, bytes):
        raise TypeError(
            f'must be {bytes!r}, not {s!r}'
        )
    return s

def hash_password(password, *, salt, check_against=None):
    password = ensure_bytes(password)
    hash = argon2.low_level.hash_secret_raw(password, salt, **ARGON2_PARAMS)
    hash, key = hash[:HASH_LEN], hash[HASH_LEN:]
    if isinstance(check_against, bytes):
        if len(check_against) != len(hash):
            raise InvalidPassword('Lengths do not match')
        z = 0
        for a,b in zip(check_against, hash):
            z ^= a ^ b
        if z != 0:
            raise InvalidPassword('Inavlid hash parameters')
    return hash, key

def create_encrypted_container(data, password):
    data = gzip.compress(data)
    salt = os.urandom(SALT_LEN)
    hash, key = hash_password(password, salt=salt)
    key = base64.urlsafe_b64encode(key) 
    cipher = Fernet(key)
    data = cipher.encrypt(data)
    verification = hash + salt
    verification = base64.urlsafe_b64encode(verification)
    return verification + data
    
def unpack_encrypted_container(data, password): 
    v, data = data[:ENCODED_VERIFY_LEN], data[ENCODED_VERIFY_LEN:]
    if len(v) != ENCODED_VERIFY_LEN:
        raise ValueError(
            'Invalid container'
        )
    v = base64.urlsafe_b64decode(v)
    hash, salt = v[:HASH_LEN], v[HASH_LEN:]
    _, key = hash_password(password, salt=salt, check_against=hash)
    key = base64.urlsafe_b64encode(key) 
    cipher = Fernet(key)
    data = cipher.decrypt(data)
    return gzip.decompress(data)


def show_help():
    eprint('use -d to decrypt, or -e to encrypt')


def eprint(*args, **kwargs):
    if kwargs.get('file', None) is None:
        kwargs['file'] = sys.stderr
    print(*args, **kwargs)


def main():
    if len(sys.argv) < 2:
        show_help()
        return
    password = getpass.getpass('password> ')
    while password != getpass.getpass('confirm> '):
        eprint('\npassword do not match. try again:\n')
    data = sys.stdin.buffer.read()
    if sys.argv[1] == '-e':
        sys.stdout.buffer.write(create_encrypted_container(data, password))
    elif sys.argv[1] == '-d':
        sys.stdout.buffer.write(unpack_encrypted_container(data, password))
    else:
        show_help()

if __name__ == '__main__':
    main()