sbeam/src/sbeam.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 jpk <jpk+dev@goatpr0n.de>
#
# Distributed under terms of the MIT license.

"""
> GET /getstream HTTP/1.1
> Host: 192.168.178.50:8080
> User-Agent: sbeam
> Accept: */*
>
* Mark bundle as not supporting multiuse
< HTTP/1.1 200 OK
< ylfrettub: 1
< Content-Disposition: attachment; filename="RedHotChillyStream"
< Date: Fri, 09 Aug 2019 12:57:20 GMT
< Content-Length: 8137967
< Content-Type: application/octet-stream

Server QR Code = http://superbe.am/q?ATKyqMA=
Base64 = 01 32 b2 a8 c0 = 01 50 178 168 192
"""

import os
import click
import requests

import SuperBeam


@click.group()
def cli():
    pass


@cli.command()
@click.option('--recursive', default=True, is_flag=True,
              help='Recurse into sub directories')
@click.argument('filename', nargs=-1, type=click.Path(exists=True))
def serve(recursive, filename):
    """Start a SuperBeam serve to serve FILENAME [FILENAME ...].

    FILENAME can be one or more files or directories.

    :param recursive boolean: recursive directory processing
    :param filename str: filename or directory to serve
    """
    filelist = SuperBeam.build_filelist(filename, recursive)
    SuperBeam.serve_forever(filelist)


@cli.command()
@click.argument('host', default='127.0.0.1')
def superlist(host):
    """Fetch superlist from HOST running SupberBeam.

    HOST is the IP of the SuperBeam server.
    """
    with requests.get(f'http://{host}:8080/superlist',
                      headers={'User-Agent': 'sbeam'}) as response:
        assert response.status_code == 200
        return response.text.split('\n')


@cli.command()
@click.argument('host', default='127.0.0.1')
@click.argument('destination', type=click.Path(exists=False))
def receive(host, destination):
    """Download a file stream from HOST and save at DESTINATION.

    Receive a file stream from HOST and save all files listed in the superlist
    in DESTINATION.
    """
    def get_superlist(host):
        with requests.get(f'http://{host}:8080/superlist',
                          headers={'User-Agent': 'sbeam'}) as response:
            assert response.status_code == 200
            return response.text.split('\n')

    filelist = get_superlist(host=host)

    with requests.get(f'http://{host}:8080/getstream', stream=True,
                      headers={'User-Agent': 'sbeam'}) as response:
        assert response.status_code == 200

        with click.progressbar(filelist, label='Receiving') as bar:
            for rawinfo in bar:
                if not rawinfo:
                    break
                filename, filesize, filepath, timestamp = rawinfo.split('||')
                width, height = click.get_terminal_size()
                term_space = 70 + len(filename)
                printname = filename
                if term_space > width:
                    printname = filename[:-1*(term_space-width-4)]
                bar.bar_template = f'%(label)s [%(bar)s] %(info)s: {printname}'

                fullpath = os.path.join(destination, filepath[1:])
                os.makedirs(fullpath, exist_ok=True)
                with open(os.path.join(fullpath, filename), 'wb') as fd:
                    # TODO split filesize in chunks with remainder
                    fd.write(response.raw.read(int(filesize)))


@cli.command()
@click.argument('host', default='127.0.0.1')
@click.argument('filename', default='received.bin')
def raw_receive(host, filename):
    """Download a SuperBeam stream from HOST and save as FILENAME.

    HOST is the IP of the SuperBeam server. The SuperBeam stream will be saved
    as FILENAME. To split the files into individual files, the `split-stream`
    command with the Superlist from the `superlist` command has to be used.
    """
    with requests.get(f'http://{host}:8080/getstream', stream=True,
                      headers={'User-Agent': 'sbeam'}) as response:
        assert response.status_code == 200
        with open(filename, 'wb') as handle:
            for chunk in response.iter_content(chunk_size=8196):
                handle.write(chunk)


@cli.command()
@click.argument('filename', default='received.bin', type=click.File('r'))
@click.argument('superlist', default='superlist', type=click.File('r'))
@click.argument('directory', default='received', type=click.File('w'))
def split_stream(filename, superlist, directory):
    """Split FILENAME with SUPERLIST into single files.

    The SuperBeam stream FILENAME from the `receive` command is split into
    individual files as described by SUPERLIST. The SUPERLIST can be downloaded
    with the `superlist` command. The individual files are stored in DIRECTORY.
    """
    with open(filename, 'rb') as blob:
        for item in superlist:
            if not item:
                break
            filename, filesize, filepath, timestamp = item.split('||')
            fullpath = os.path.join(directory, filepath[1:])
            os.makedirs(fullpath, exist_ok=True)
            print(f'Saving {filename} ({filesize} bytes) to disk...')
            with open(os.path.join(fullpath, filename), 'wb') as handle:
                handle.write(blob.read(int(filesize)))


@cli.command()
def tui():
    SuperBeam.tui.run()


if __name__ == '__main__':
    cli()