Fix download resumption and add stuff

requirements.txt

@@ -0,0 +1,4 @@
+aiohttp
+yarl
+telethon
+cryptg

setup.py

@@ -0,0 +1,45 @@
+import setuptools
+
+try:
+    long_desc = open("README.md").read()
+except IOError:
+    long_desc = "Failed to read README.md"
+
+setuptools.setup(
+    name="tgfilestream",
+    version="0.1.0",
+    url="https://mau.dev/tulir/tgfilestream",
+
+    author="Tulir Asokan",
+    author_email="[email protected]",
+
+    description="A Telegram bot that can stream Telegram files to users over HTTP.",
+    long_description=long_desc,
+    long_description_content_type="text/markdown",
+
+    packages=setuptools.find_packages(),
+
+    install_requires=[
+        "aiohttp>=3",
+        "telethon>=1.10",
+        "yarl>=1",
+    ],
+    extras_require={
+        "fast": ["cryptg>=0.2"],
+    },
+    python_requires="~=3.6",
+
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
+        "Framework :: AsyncIO",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+    ],
+    entry_points="""
+        [console_scripts]
+        tgfilestream=tgfilestream:main
+    """,
+)

tgfilestream.py

@@ -1,4 +1,5 @@
-from typing import Tuple, Union, cast
+from typing import Tuple, Union, AsyncIterable, cast
+import logging
 import asyncio
 import sys
 import os
@@ -7,6 +8,7 @@
 from telethon.tl.custom import Message
 from telethon.tl.types import TypeInputPeer, InputPeerChannel, InputPeerChat, InputPeerUser
 from aiohttp import web
+from yarl import URL
 
 try:
     api_id = int(os.environ["TG_API_ID"])
@@ -23,14 +25,20 @@
     print("Please make sure the PORT environment variable is an integer between 1 and 65535")
     sys.exit(1)
 
+debug = bool(os.environ.get("DEBUG"))
+trust_headers = bool(os.environ.get("TRUST_FORWARD_HEADERS"))
 host = os.environ.get("HOST", "localhost")
-public_url = os.environ.get("PUBLIC_URL", f"http://{host}:{port}")
+public_url = URL(os.environ.get("PUBLIC_URL", f"http://{host}:{port}"))
 session_name = os.environ.get("TG_SESSION_NAME", "tgfilestream")
+pack_bits = 32
+pack_bit_mask = (1 << pack_bits) - 1
+
 client = TelegramClient(session_name, api_id, api_hash)
 routes = web.RouteTableDef()
 
-pack_bits = 32
-pack_bit_mask = (1 << pack_bits) - 1
+log = logging.getLogger("tgfilestream")
+logging.basicConfig(level=logging.DEBUG if debug else logging.INFO)
+logging.getLogger("telethon").setLevel(logging.INFO if debug else logging.ERROR)
 
 group_bit = 0b01
 channel_bit = 0b10
@@ -70,32 +78,79 @@ def get_file_name(message: Union[Message, events.NewMessage.Event]) -> str:
     return f"{message.date.strftime('%Y-%m-%d_%H:%M:%S')}{ext}"
 
 
-@client.on(events.NewMessage)
-async def handle_message(evt: events.NewMessage.Event) -> None:
-    if not evt.is_private or not evt.file:
-        return
-    url = f"{public_url}/{pack_id(evt)}/{get_file_name(evt)}"
-    await evt.reply(f"Link to download file: [{url}]({url})")
+def get_requester_ip(req: web.Request) -> str:
+    if trust_headers:
+        try:
+            return req.headers["X-Forwarded-For"]
+        except KeyError:
+            pass
+    peername = req.transport.get_extra_info('peername')
+    if peername is not None:
+        return peername[0]
 
 
-@routes.get(r"/{id:\d+}/{name}")
-async def handle_request(req: web.Request) -> web.Response:
+async def cut_first_chunk(iterable: AsyncIterable[bytes], cut: int) -> AsyncIterable[bytes]:
+    first = True
+    async for chunk in iterable:
+        if first:
+            chunk = chunk[cut:]
+            first = False
+        yield chunk
+
+
+async def handle_request(req: web.Request, head: bool = False) -> web.Response:
     file_name = req.match_info["name"]
     file_id = int(req.match_info["id"])
     peer, msg_id = unpack_id(file_id)
     if not peer or not msg_id:
         return web.Response(status=404, text="404: Not Found")
+
     message = cast(Message, await client.get_messages(entity=peer, ids=msg_id))
     if not message or not message.file or get_file_name(message) != file_name:
         return web.Response(status=404, text="404: Not Found")
-    return web.Response(status=200, body=client.iter_download(message.media),
+
+    offset = req.http_range.start or 0
+    tg_offset = offset - offset % (2 ** 19)
+    size = message.file.size
+
+    if not head:
+        log.info(
+            f"Serving file in {message.id} (chat {message.chat_id}) to {get_requester_ip(req)}")
+        body = client.iter_download(message.media, file_size=message.file.size, offset=tg_offset)
+        body = cut_first_chunk(body, offset - tg_offset)
+    else:
+        body = None
+    return web.Response(status=206 if offset else 200,
+                        body=body,
                         headers={
                             "Content-Type": message.file.mime_type,
-                            "Content-Length": str(message.file.size),
-                            "Content-Disposition": f'attachment; filename="{file_name}"'
+                            "Content-Range": f"bytes {offset}-{size}/{size}",
+                            "Content-Length": str(size - offset),
+                            "Content-Disposition": f'attachment; filename="{file_name}"',
+                            "Accept-Ranges": "bytes",
                         })
 
 
+@routes.head(r"/{id:\d+}/{name}")
+async def handle_head_request(req: web.Request) -> web.Response:
+    return await handle_request(req, head=True)
+
+
+@routes.get(r"/{id:\d+}/{name}")
+async def handle_get_request(req: web.Request) -> web.Response:
+    return await handle_request(req, head=False)
+
+
+@client.on(events.NewMessage)
+async def handle_message(evt: events.NewMessage.Event) -> None:
+    if not evt.is_private or not evt.file:
+        return
+    url = public_url / str(pack_id(evt)) / get_file_name(evt)
+    await evt.reply(f"Link to download file: [{url}]({url})")
+    log.info(f"Replied with link for {evt.id} to {evt.from_id} in {evt.chat_id}")
+    log.debug(f"Link to {evt.id} in {evt.chat_id}: {url}")
+
+
 server = web.Application()
 server.add_routes(routes)
 runner = web.AppRunner(server)
@@ -116,15 +171,20 @@ async def stop() -> None:
     await client.disconnect()
 
 
-loop.run_until_complete(start())
-print("========== Initialization complete ==========")
-print(f"Listening at http://{host}:{port}")
-print(f"Public URL prefix is {public_url}")
-print("(Press CTRL+C to quit)")
+try:
+    loop.run_until_complete(start())
+except Exception:
+    log.fatal("Failed to initialize", exc_info=True)
+    sys.exit(2)
+
+log.info("Initialization complete")
+log.debug(f"Listening at http://{host}:{port}")
+log.debug(f"Public URL prefix is {public_url}")
+
 try:
     loop.run_forever()
 except KeyboardInterrupt:
     loop.run_until_complete(stop())
 except Exception:
-    print("Fatal error in event loop")
-    raise
+    log.fatal("Fatal error in event loop", exc_info=True)
+    sys.exit(3)