#! /usr/bin/env python # -*- coding: utf-8 -*- # vim:fenc=utf-8 # # Copyright © 2018 jpk # # Distributed under terms of the MIT license. import os import click import requests import SuperBeam __VERSION__ = '0.5' @click.group() @click.version_option(version=__VERSION__, prog_name='SuperBeam CLI') 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 jsonlist(host): """Fetch jsonlist from HOST running SupberBeam. HOST is the IP of the SuperBeam server. """ with requests.get(f'http://{host}:8080/jsonlist', 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: fd.write(response.raw.read(int(filesize))) bar.update(len(filelist)) @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))) if __name__ == '__main__': try: cli() except requests.exceptions.ConnectionError as conerr: print(f'Failed to connect. Target alive?')