x360-web-queue/app/download.py

"""aria2 JSON-RPC. aria2 is started alongside the API server, not exposed publicly."""
import asyncio
import base64
import json
import os
import urllib.parse
from pathlib import Path

import httpx

from . import archive_auth, store

RPC = 'http://127.0.0.1:6800/jsonrpc'


async def rpc(method, *params):
    payload = {
        'jsonrpc': '2.0', 'id': 'x360', 'method': 'aria2.' + method,
        'params': ['token:' + os.environ['ARIA2_SECRET'], *params],
    }
    # On a container restart the API may start slightly before aria2 RPC.
    # Retry the connection instead of permanently failing restored jobs.
    async with httpx.AsyncClient(timeout=15) as client:
        for attempt in range(15):
            try:
                r = await client.post(RPC, json=payload)
                r.raise_for_status()
                answer = r.json()
                if 'error' in answer:
                    raise RuntimeError(answer['error'].get('message', 'Bledna odpowiedz aria2'))
                return answer['result']
            except (httpx.ConnectError, httpx.ConnectTimeout):
                if attempt == 14:
                    raise
                await asyncio.sleep(1)


def source_urls(job):
    if job['source_type'] == 'archive_many':
        return json.loads(job['source'])
    return [job['source']]


def target_filename(uri):
    if uri.lower().startswith('magnet:'):
        return None
    name = urllib.parse.unquote(Path(urllib.parse.urlsplit(uri).path).name)
    if not name or name in ('.', '..'):
        return 'pobrany-plik.bin'
    # Do not allow URI-provided paths to escape the download directory.
    return ''.join(c if c not in '\\/:*?"<>|\x00\n\r' else '_' for c in name)[:200]


async def begin(job):
    base = store.DATA / 'raw' / job['id']
    base.mkdir(parents=True, exist_ok=True)
    urls = source_urls(job)
    ix = job['part_index']
    if ix >= len(urls):
        store.update(job['id'], state='preparing')
        return
    uri = urls[ix]
    opts = {
        'dir': str(base), 'continue': 'true', 'max-connection-per-server': '2',
        'split': '2', 'seed-time': '0', 'auto-file-renaming': 'false',
        'allow-overwrite': 'false', 'max-tries': '8', 'retry-wait': '30',
    }
    if job['source_type'] != 'torrent_selected':
        name = target_filename(uri)
        if name and '.' in name:
            opts['out'] = name
        elif name:
            # Let the server's Content-Disposition supply the real filename.
            opts['content-disposition'] = 'true'
    if job['source_type'] == 'torrent_selected':
        torrent = store.DATA / 'torrents' / (job['id'] + '.torrent')
        encoded = base64.b64encode(torrent.read_bytes()).decode('ascii')
        select = json.loads(job['source'])['selection']
        opts['select-file'] = select
        opts['bt-remove-unselected-file'] = 'true'
        gid = await rpc('addTorrent', encoded, [], opts)
        note = 'Pobieranie wybranych plików torrenta'
    else:
        if archive_auth.is_archive_download(uri):
            # aria2 downloads from the local proxy. It never receives IA cookies.
            opts['header'] = [archive_auth.signed_header(job['id'], ix)]
            uri = archive_auth.internal_url(job['id'], ix)
            opts['split'] = '1'
            opts['max-connection-per-server'] = '1'
        gid = await rpc('addUri', [uri], opts)
        note = f'Pobieranie pliku {ix+1}/{len(urls)}'
    store.update(job['id'], gid=gid, state='downloading', note=note)


async def poll(job):
    urls = source_urls(job)
    gid = job['gid']
    if not gid:
        await begin(job)
        return
    try:
        response = await rpc('tellStatus', gid, ['gid','status','totalLength','completedLength',
                                                'downloadSpeed','followedBy','errorMessage'])
    except RuntimeError as ex:
        # On restart aria2 may lose an old GID. Re-add to resume available .aria2 files.
        if 'not found' in str(ex).lower():
            store.update(job['id'], gid='', note='Odzyskiwanie transferu po restarcie')
            return
        raise

    followed = response.get('followedBy', [])
    if followed:
        store.update(job['id'], gid=followed[0], note='Magnet: metadane gotowe')
        return
    state = response.get('status')
    length = int(response.get('totalLength', 0) or 0)
    completed = int(response.get('completedLength', 0) or 0)
    speed = int(response.get('downloadSpeed', 0) or 0)
    part_progress = (completed / length) if length else 0
    ix = job['part_index']
    pct = (ix + part_progress) / len(urls) * 100
    store.update(job['id'], progress=round(pct, 2), speed=speed, total=length,
                 completed=completed, note=f'Pobieranie {ix+1}/{len(urls)}: {state}')
    if state == 'complete':
        if ix + 1 < len(urls):
            store.update(job['id'], gid='', part_index=ix+1)
        else:
            store.update(job['id'], state='preparing', progress=100, speed=0,
                         note='Pobieranie ukonczone, przygotowuje pliki')
    elif state in ('error','removed'):
        message = response.get('errorMessage') or state
        if job['source_type'] == 'archive_many' or (job['source_type'] == 'url' and archive_auth.is_archive_download(urls[ix])):
            if any(x in message.lower() for x in ('authorization', '403', '401', 'forbidden')):
                message += ' — Internet Archive: sprawdź logowanie i dostęp do pliku.'
        raise RuntimeError('aria2: ' + message)
    elif state == 'paused':
        store.update(job['id'], state='paused', speed=0)


async def cancel_download(gid):
    if not gid:
        return
    try:
        await rpc('forceRemove', gid)
    except (RuntimeError, httpx.HTTPError):
        pass


async def pause_download(gid):
    if gid:
        await rpc('forcePause', gid)


async def resume_download(gid):
    if gid:
        await rpc('unpause', gid)