140 lines
4.9 KiB
Python
Executable File
140 lines
4.9 KiB
Python
Executable File
#! /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.
|
|
|
|
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?')
|