wba/backend/keephelper/__main__.py

import aiohttp
from aiohttp import web
from argparse import ArgumentParser
import ssl
import json


def options_fetch_blocks(request):
    return web.Response(headers={ 'Access-Control-Allow-Origin': '*' })


async def post_fetch_blocks(request):
    body = await request.read()
    body = json.loads(body)
    proxy_host = body['keepProxyHost']
    arv_token = body['arvToken']
    segments = body['segments']
    use_ssl = body['useSsl'] \
        if 'useSsl' in body \
        else True
    protocol = 'https://' \
        if use_ssl \
        else 'http://'
    name = body['name'] \
        if 'name' in body \
        else None
    content_type = body['contentType'] \
        if 'contentType' in body \
        else 'application/octet-stream'
    res = web.StreamResponse()
    res.content_type = content_type
    if name:
        res.headers['Content-Disposition'] = 'attachment; filename*=UTF-8\'\'"' + name + '"'
    else:
        res.headers['Content-Disposition'] = 'inline'
    res.headers['Access-Control-Allow-Origin'] = '*'
    await res.prepare(request)
    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(verify_ssl=False)) as session:
        for seg in segments:
            url = protocol + proxy_host + '/' + seg[0]
            async with session.get(url, headers={ 'Authorization': 'OAuth2 ' + arv_token }) as response:
                block = await response.read()
                block = block[seg[1]:seg[2]]
                await res.write(block)
    return res


def get_index(request):
    return web.Response(text='Use /fetch-blocks to stream files from Keep')


def create_parser():
    parser = ArgumentParser()
    parser.add_argument('--port', type=int, default=50080)
    parser.add_argument('--ssl-cert', type=str, default=None)
    return parser


def main():
    parser = create_parser()
    args = parser.parse_args()
    app = web.Application()
    app.add_routes([
        web.get('/', get_index),
        web.post('/fetch-blocks', post_fetch_blocks),
        web.options('/fetch-blocks', options_fetch_blocks)
    ])
    if args.ssl_cert:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
        ssl_context.load_cert_chain(args.ssl_cert)
    else:
        ssl_context = None
    web.run_app(app, port=args.port, ssl_context=ssl_context)


if __name__ == '__main__':
    main()