crypt

#!/usr/bin/env python

import os, sys, hashlib as hl, contextlib as cl


class CryptConf:
	bs = 1 * 2**20
	scrypt_salt = b'crypt.rFnuzEwQBUXITkIs'
	scrypt_nx, scrypt_r, scrypt_p = 18, 8, 1
	scrypt_maxmem = 300 * 2**20


def nacl_init(cache=list()):
	if not cache:
		import nacl.bindings.crypto_secretstream as csm # https://github.com/pyca/pynacl/
		class CS:
			def __getattr__(self, k):
				return getattr(csm, f'crypto_secretstream_xchacha20poly1305_{k}')
		cache.append(CS())
	return cache[0]

def encrypt_stream(src, dst, sym_key, bs=256*2**10):
	csx = nacl_init()
	chunk = src.read(bs)
	if not chunk: return
	if len(sym_key) != csx.KEYBYTES:
		raise ValueError(f'Key must be exactly {csx.KEYBYTES}B long')
	state = csx.state()
	dst.write(csx.init_push(state, sym_key))
	dst.write(csx.push( state,
		bs.to_bytes(8, 'big'), tag=csx.TAG_MESSAGE ))
	while chunk:
		chunk_next = src.read(bs)
		dst.write(csx.push( state, chunk,
			tag=csx.TAG_FINAL if not chunk_next else 0 ))
		chunk = chunk_next

def decrypt_stream(src, dst, sym_key):
	csx = nacl_init()
	chunk = src.read(csx.HEADERBYTES)
	if not chunk: return
	state = csx.state()
	csx.init_pull(state, chunk, sym_key)
	chunk, tag = csx.pull(state, src.read(8 + csx.ABYTES))
	bs = int.from_bytes(chunk, 'big') + csx.ABYTES
	while tag != csx.TAG_FINAL:
		frame = src.read(bs)
		if not frame: raise ValueError('Corrupted stream, missing chunks')
		chunk, tag = csx.pull(state, frame)
		dst.write(chunk)


def main(args=None, conf=None):
	if not conf: conf = CryptConf()

	import argparse, re, textwrap
	dd = lambda text: re.sub( r' \t+', ' ',
		textwrap.dedent(text).strip('\n') + '\n' ).replace('\t', '  ')
	parser = argparse.ArgumentParser(
		formatter_class=argparse.RawTextHelpFormatter,
		usage='%(prog)s [options] path',
		description=dd('''
			Simple streaming encryption tool.
			Uses PyNaCl's crypto_secretstream_xchacha20poly1305 API
				in an all-or-nothing fashion (no support for partial encryption/decryption).
			Example use: crypt -ek mykey secret.tar secret.tar.enc'''))

	parser.add_argument('src', nargs='?',
		help='Source stream file. Use "-" for stdin (default).')
	parser.add_argument('dst', nargs='?',
		help='Destination stream file. Use "-" for stdout (default).')

	parser.add_argument('-k', '--key', metavar='string/@file',
		help=dd(f'''
			Key seed to use for encryption/decryption.
			If prefixed by @ then first stripped line from a specified file will be used as key seed.
			Use @@ prefix for @ as the first character of seed instead.
			Always stretched via scrypt (n=2^{conf.scrypt_nx}, \
				r={conf.scrypt_r}, p={conf.scrypt_p}) before use.'''))
	parser.add_argument('-e', '--encrypt', action='store_true',
		help='Encrypt stream with a specified -k/--key.')
	parser.add_argument('-d', '--decrypt', action='store_true',
		help='Decrypt stream using specified -k/--key.')

	opts = parser.parse_args(sys.argv[1:] if args is None else args)

	if not (opts.encrypt or opts.decrypt):
		parser.error('One of -e/--encrypt or -d/--decrypt must be specified.')
	if opts.encrypt and opts.decrypt:
		parser.error('Both -e/--encrypt and -d/--decrypt must not be specified at once.')
	if not opts.key: parser.error('-k/--key option is required.')

	csx, enc = nacl_init(), bool(opts.encrypt)

	key_raw = opts.key
	if key_raw.startswith('@@'): key_raw = ('@' + key_raw[2:]).encode()
	elif key_raw[0] == '@':
		with open(os.path.expanduser(key_raw[1:]), 'rb') as key_file:
			key_raw = key_file.readline().strip()
	else: key_raw = key_raw.encode()
	enc_key = hl.scrypt( key_raw, salt=conf.scrypt_salt,
		n=2**conf.scrypt_nx, r=conf.scrypt_r, p=conf.scrypt_p,
		maxmem=conf.scrypt_maxmem, dklen=32 )

	with cl.ExitStack() as ctx:
		src = ctx.enter_context(open(
			sys.stdin.fileno() if (opts.src or '-') == '-' else opts.src, 'rb', conf.bs ))
		dst = ctx.enter_context(open(
			sys.stdout.fileno() if (opts.dst or '-') == '-' else opts.dst, 'wb', conf.bs ))
		if enc: encrypt_stream(src, dst, enc_key)
		else: decrypt_stream(src, dst, enc_key)

if __name__ == '__main__': sys.exit(main())