CLI of https://transfer.sh/ with encryption support.

transfersh.py

#!/usr/bin/env python3

import argparse
import hashlib
import os
import secrets
import requests
import typing

from Crypto.Cipher import AES
from Crypto.Cipher import ChaCha20_Poly1305

from Crypto.Cipher._mode_gcm import GcmMode
from Crypto.Cipher.ChaCha20_Poly1305 import ChaCha20Poly1305Cipher

# print(''.join(f'\\x{x:02X}' for x in secrets.token_bytes(32)))
KDF_SALT = b'\x44\x12\xBA\x44\xE9\xBB\x76\xB0\x26\x45\x5C\xF5\x8A\x02\xED\x98\xAD\xB0\x8E\xBD\xEB\xF4\xE3\xC9\xFF\xB1\xAC\xE0\x54\x93\x24\xD9'
BUFFER_SIZE = 65536

def uploadStream(cipher: typing.Union[GcmMode, ChaCha20Poly1305Cipher], path: str) -> typing.Generator[bytes, None, None]:
    yield cipher.nonce
    fileLen = os.path.getsize(path)
    fileLenBytes = bytes((fileLen >> (i << 3) & 0xFF) for i in range(8))
    yield cipher.encrypt(fileLenBytes)
    pathFilename = os.path.split(path)[1]
    pathLenBytes = bytes((len(pathFilename) >> (i << 3) & 0xFF) for i in range(2))
    yield cipher.encrypt(pathLenBytes)
    yield cipher.encrypt(pathFilename.encode('utf-8'))
    with open(path, 'rb') as f:
        for i in range(0, fileLen, BUFFER_SIZE):
            yield cipher.encrypt(f.read(min(fileLen - i, BUFFER_SIZE)))
    yield cipher.digest()

def downloadStream(cipher: typing.Union[GcmMode, ChaCha20Poly1305Cipher], request: requests.Response) -> None:
    fileLenBytes = cipher.decrypt(request.raw.read(8))
    fileLen = 0
    for i, j in enumerate(fileLenBytes):
        fileLen |= j << (i << 3)
    pathLenBytes = cipher.decrypt(request.raw.read(2))
    pathLen = 0
    for i, j in enumerate(pathLenBytes):
        pathLen |= j << (i << 3)
    path = cipher.decrypt(request.raw.read(pathLen)).decode('utf-8')
    with open(path, 'wb') as f:
        for i in range(0, fileLen, BUFFER_SIZE):
            f.write(cipher.decrypt(request.raw.read(min(fileLen - i, BUFFER_SIZE))))
    mac = request.raw.read(16)
    cipher.verify(mac)
    print('Downloaded', path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'command',
        type=str.lower,
        choices=('upload', 'download', 'delete', 'up', 'down', 'del'),
    )
    parser.add_argument(
        'source',
        type=str,
        help='The file to upload / URL of uploaded file / URL for deleting uploaded file.',
    )
    parser.add_argument(
        '-c',
        '--count',
        type=int,
        default=None,
        help='Limit the number of downloads while uploading.',
    )
    parser.add_argument(
        '-t',
        '--time',
        type=int,
        default=None,
        help='Set the number of days before deletion while uploading.',
    )
    parser.add_argument(
        '-m',
        '--method',
        type=str.lower,
        choices=('aes-gcm', 'chacha20-poly1305'),
        default='aes-gcm',
        help='Encryption method.',
    )
    parser.add_argument(
        '-n',
        '--iteration',
        type=int,
        default=262144,
        help='PBKDF2-SHA256 key iteration count.',
    )
    args = parser.parse_args()

    session = requests.Session()
    session.hooks['response'].append(lambda r, *args, **kwargs: r.raise_for_status())

    if args.command in {'up', 'upload'}:
        passphrase = input('Passphrase (Leave empty to generate one): ').strip()
        if not passphrase:
            passphrase = secrets.token_urlsafe(12)
            print('Generated passphrase:', passphrase)
        randomBytes = hashlib.pbkdf2_hmac('sha256', passphrase.encode('utf-8'), KDF_SALT, args.iteration, 64)
        if args.method == 'aes-gcm':
            key, aad = randomBytes[:16], randomBytes[16:]
            cipher = AES.new(key=key, mode=AES.MODE_GCM)
        elif args.method == 'chacha20-poly1305':
            key, aad = randomBytes[:32], randomBytes[32:]
            cipher = ChaCha20_Poly1305.new(key=key)
        cipher.update(aad)
        h = {}
        if args.count:
            h['Max-Downloads'] = str(args.count)
        if args.time:
            h['Max-Days'] = str(args.time)
        r = session.put(
            f'https://transfer.sh/{secrets.token_urlsafe(6)}',
            data=uploadStream(cipher, args.source),
            headers=h,
        )
        print('Download URL:', r.text)
        print('Delete URL:', r.headers['X-Url-Delete'])
    elif args.command in {'down', 'download'}:
        passphrase = None
        while not passphrase:
            passphrase = input('Passphrase: ').strip()
        randomBytes = hashlib.pbkdf2_hmac('sha256', passphrase.encode('utf-8'), KDF_SALT, args.iteration, 64)
        r = session.get(args.source, stream=True)
        if args.method == 'aes-gcm':
            key, aad = randomBytes[:16], randomBytes[16:]
            nonce = r.raw.read(16)
            cipher = AES.new(key=key, mode=AES.MODE_GCM, nonce=nonce)
        elif args.method == 'chacha20-poly1305':
            key, aad = randomBytes[:32], randomBytes[32:]
            nonce = r.raw.read(12)
            cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
        cipher.update(aad)
        downloadStream(cipher, r)
    elif args.command in {'del', 'delete'}:
        session.delete(args.source)
        print('File deleted.')