share-via-ssh

#!/usr/bin/env python3

import argparse
import configparser
import datetime
import os
import pathlib
import random
import re
import shlex
import subprocess
import sys
import textwrap
import urllib.parse


try:
    import qrcode
except ImportError as e:
    qrcode = None


__version__ = '0.3.0'


# Filename for uploads from stdin
DEFAULT_FILENAME = "index.txt"


def main(arglist):

    # Parse command line arguments
    parser = argparse.ArgumentParser('share-via-ssh')
    parser.add_argument('path_strs', metavar='FILE', nargs='*',
        help='paths to files and/or directories to upload')
    parser.add_argument('-c', '--config', metavar='...',
        help='config file path')
    parser.add_argument('-g', '--group', metavar='...',
        help='existing directory to upload files to')
    parser.add_argument('-e', '--expire', metavar='...',
        help='expiration time (e.g. 15m, 1h, 1d, 2w, 2015-12-24 18:00, tomorrow)')
    parser.add_argument('--dry-run', action='store_true',
        help='do not actually upload anything')
    parser.add_argument('-i', '--stdin', action='store_true',
        help='upload text from stdin instead of files from the filesystem')
    parser.add_argument('--qr', metavar='TYPE', dest='qrcode_method',
        nargs='?', const='ascii',
        choices=['no', 'tty', 'ascii', 'invert'],
        help='enable QR code output and choose render method')
    args = parser.parse_args(arglist)

    # Check group name
    if args.group and ('/' in args.group or args.group in ('.', '..')):
        print('error: bad group name: %s' % args.group, file=sys.stderr)
        sys.exit(1)

    # Load config file(s)
    config = configparser.ConfigParser()
    if args.config:
        config_path_obj = pathlib.Path(args.config)
        if not config_path_obj.is_file():
            print('error: no such config file: %s' % config_path_obj, file=sys.stderr)
            sys.exit(1)
        config.read([args.config])
    else:
        config_path_objs = [pathlib.Path(s).expanduser() for s in [
            '~/.config/share-via-ssh.conf',
            '~/.share-via-ssh.conf'
        ]]
        if not config.read(map(str, config_path_objs)):
            print('error: found no config file in any of the following locations:', file=sys.stderr)
            print(file=sys.stderr)
            for config_path_obj in config_path_objs:
                print('- %s' % config_path_obj)
            sys.exit(1)

    # Check config file and extract settings
    if 'share-via-ssh' not in config:
        print('error: no [share-via-ssh] section in config file', file=sys.stderr)
        sys.exit(1)
    if 'host' not in config['share-via-ssh']:
        print('error: no "host" setting in config file', file=sys.stderr)
        sys.exit(1)
    host = config['share-via-ssh']['host']
    if 'base_dir' not in config['share-via-ssh']:
        print('error: no "base_dir" setting in config file', file=sys.stderr)
        sys.exit(1)
    base_dir = config['share-via-ssh']['base_dir']
    if 'base_url' not in config['share-via-ssh']:
        print('error: no "base_url" setting in config file', file=sys.stderr)
        sys.exit(1)
    base_url = config['share-via-ssh']['base_url']

    # Expiration date
    default_expiration = None  # never
    if 'expire' in config['share-via-ssh']:
        try:
            default_expiration = parse_expiration(config['share-via-ssh']['expire'])
        except ExpirationParseError as e:
            print('error in config file: %s' % e, file=sys.stderr)
            sys.exit(1)
    if args.expire is not None:
        try:
            expiration = parse_expiration(args.expire)
        except ExpirationParseError as e:
            print('error: %s' % e, file=sys.stderr)
            sys.exit(1)
    else:
        expiration = default_expiration

    # QR code method
    default_qrcode_method = 'tty' if qrcode else None
    if 'qr_code' in config['share-via-ssh']:
        choices = ['no', 'tty', 'ascii', 'invert']
        default_qrcode_method = config['share-via-ssh']['qr_code']
        if default_qrcode_method not in choices:
            print('error in config file: bad value for qr_code')
            print('allowed values are: %s' % ', '.join(choices))
            sys.exit(1)
        if default_qrcode_method == 'no':
            default_qrcode_method = None
    if args.qrcode_method == 'no':
        qrcode_method = None
    elif args.qrcode_method:
        qrcode_method = args.qrcode_method
    else:
        qrcode_method = default_qrcode_method
    if qrcode_method and not qrcode:
        print('warning: missing library. use "pip3 install qrcode" to fix.\n', file=sys.stderr)

    # Ensure --stdin and <files> aren't used together
    if args.stdin and args.path_strs:
        print('error: when --stdin is used, no files can be specified')
        sys.exit(1)
    elif not args.stdin and not args.path_strs:
        print('error: no files specified')
        sys.exit(1)

    # If taking data from stdin, read all before touching server so the user can safely ctrl-c
    if args.stdin:
        upload_text = sys.stdin.buffer.read()

    # Check paths
    path_objs = []
    for path_str in args.path_strs:
        path_obj = pathlib.Path(path_str).expanduser()
        if not path_obj.exists():
            print('error: file or directory doesn\'t exist: %s' % path_obj, file=sys.stderr)
            sys.exit(1)
        if path_obj.name in ('.htaccess', '.htpasswd'):
            print('error: refusing to upload .htaccess/.htpasswd')
            sys.exit(1)
        path_objs.append(path_obj)

    # Create group directory on server, if necessary
    group = args.group
    if not group:
        rnd = random.SystemRandom()
        group = "".join(rnd.choice("abcdefghijkmnpqrstuvwxyz23456789") for i in range(12))
        group_dir = os.path.join(base_dir, group) + "/"
        command = ['ssh', '-q', host, 'mkdir', group_dir]
        if args.dry_run:
            print(" ".join(map(shlex.quote, command)))
        else:
            subprocess.check_call(command)
    else:
        group_dir = os.path.join(base_dir, group) + "/"

    # Create .htaccess file with expiration date
    if expiration:
        htaccess_path = os.path.join(group_dir, ".htaccess")
        htaccess_text = textwrap.dedent("""\
            RewriteEngine on
            RewriteCond %%{TIME} >%s
            RewriteRule "(.*)" "-" [G]
        """).replace("\n", "\\n") % expiration.strftime("%Y%m%d%H%M%S")
        command = "echo %s > %s" % (shlex.quote(htaccess_text), htaccess_path)
        command = ['ssh', '-q', host, 'sh', '-c', shlex.quote(command)]
        if args.dry_run:
            print(" ".join(map(shlex.quote, command)))
        else:
            subprocess.check_call(command)

    # Copy text from stdin
    if args.stdin:
        command = ['ssh', '-o', "BatchMode yes", host, 'cat - > %s/%s' % (group_dir, DEFAULT_FILENAME)]
        proc = subprocess.Popen(command, stdin=subprocess.PIPE)
        proc.stdin.write(upload_text)
        proc.stdin.close()
        ret = proc.wait()
        if ret != 0:
            raise CalledProcessError(ret, cmd, None)

    # Copy files to server
    else:
        source_file_strs = []
        for path_obj in path_objs:
            source_file_strs.append(str(path_obj).replace(":", "\\:"))
        command = ['scp', '-q', '-r'] + source_file_strs + [host + ':' + group_dir]
        if args.dry_run:
            print(" ".join(map(shlex.quote, command)))
        else:
            subprocess.check_call(command)

    # Change permissions
    command = ['ssh', '-q', host, 'chmod', '-R', 'a+rX', group_dir]
    if args.dry_run:
        print(" ".join(map(shlex.quote, command)))
    else:
        subprocess.check_call(command)

    # Show report
    group_url = base_url + group + "/"
    if args.stdin:
        print("Link:       \x1b[1m%s \x1b[0m" % (group_url + urllib.parse.quote(DEFAULT_FILENAME)))
    else:
        for path_obj in path_objs:
            print(group_url + urllib.parse.quote(str(path_obj.name)))
        print()
        print("Directory:  \x1b[1m%s \x1b[0m" % group_url)
    print("Add files:  \x1b[1mshare-via-ssh --group=%s\x1b[0m" % group)
    if expiration:
        print("Expires:    %s" % expiration.strftime("%a, %Y-%m-%d %H:%M:%M"))

    # Show QR code
    if qrcode_method and qrcode:
        print()

        if len(path_objs) > 1:
            qr_url = group_url
        else:
            qr_url = group_url + urllib.parse.quote(str(path_objs[0].name))

        qr = qrcode.QRCode(image_factory=None)
        qr.add_data(qr_url)

        if qrcode_method == "no":
            pass
        elif qrcode_method == "ascii":
            qr.print_ascii(invert=True)
        elif qrcode_method == "invert":
            qr.print_ascii(invert=True)
        elif qrcode_method == "tty":
            qr.print_tty()


class ExpirationParseError(Exception):
    pass


def parse_expiration(string, _now=None):
    now = _now or datetime.datetime.today()

    # Check for keywords
    patterns = {
        "never": None,
        "none": None,
        "": None,
        "now": now,
        "today": datetime.datetime.combine(now.date(), datetime.time(23, 59, 59)),
        "tomorrow": datetime.datetime.combine(
            now.date() + datetime.timedelta(days=1), datetime.time(23, 59, 59)),
    }
    for pattern, expiration in patterns.items():
        if string == pattern:
            return expiration

    # Try parsing it as an absolute date/time
    # (like "2015-12-24", "2015-12-24 18:00", "18:00")
    patterns = {
        "%Y%m%d%H%M%S": lambda dt: dt,
        "%Y-%m-%d %H:%M:%S": lambda dt: dt,
        "%Y-%m-%d %H:%M": lambda dt: dt,
        "%Y-%m-%d %H": lambda dt: dt,
        "%Y-%m-%d": lambda dt: dt,
        "%Y-%m": lambda dt: dt,
        "%Y": lambda dt: dt,
        "%H:%M:%S": lambda dt: datetime.datetime.combine(now.date(), dt.time()),
        "%H:%M": lambda dt: datetime.datetime.combine(now.date(), dt.time()),
        "%H": lambda dt: datetime.datetime.combine(now.date(), dt.time()),
    }
    for pattern, convert in patterns.items():
        try:
            match = datetime.datetime.strptime(string, pattern)
        except ValueError:
            continue
        else:
            return convert(match)

    # Try parsing it as a relative date
    # (like "3h", "3h15m", "3h 15m", "1week")
    expiration = now
    patterns = {
        r"(\d+) ?(?:s|secs?|seconds?)": datetime.timedelta(seconds=1),
        r"(\d+) ?(?:m|mins?|minutes?)": datetime.timedelta(minutes=1),
        r"(\d+) ?(?:h|hrs?|hours?)": datetime.timedelta(hours=1),
        r"(\d+) ?(?:d|days?)": datetime.timedelta(days=1),
        r"(\d+) ?(?:w|wks?|weeks?)": datetime.timedelta(weeks=1),
    }
    found_one = False
    for pattern, multiplier in patterns.items():
        for match in re.compile(pattern).findall(string):
            expiration += int(match) * multiplier
            found_one = True
    if found_one:
        return expiration

    # None of the above methods worked
    raise ExpirationParseError("couldn't parse expiration date: %s" % string)


if __name__ == '__main__':
    main(sys.argv[1:])