first commit

app/download.py

@@ -0,0 +1,156 @@
+"""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)