From 10a15f33c847b76b106cf3f88bec1f3e715832dc Mon Sep 17 00:00:00 2001
From: blank X <theblankx@protonmail.com>
Date: Thu, 6 May 2021 09:42:10 +0700
Subject: [PATCH] Use built-in Range parser

---
 streamtg.py | 56 +++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 42 insertions(+), 14 deletions(-)

diff --git a/streamtg.py b/streamtg.py
index 4ce6a5b..04fec8f 100644
--- a/streamtg.py
+++ b/streamtg.py
@@ -2,6 +2,7 @@ import logging
 logging.basicConfig(level=logging.INFO)
 
 import os
+import re
 import yaml
 import hmac
 from aiohttp import web
@@ -16,6 +17,7 @@ client = TelegramClient('streamtg', config['telegram']['api_id'], config['telegr
 authorized_tokens = config.get('authorized_tokens')
 hmacs = [hmac.new(i['key'].encode(), digestmod=i['digest']) for i in config.get('hmac', ())]
 port = os.environ.get('PORT', 8080)
+bytes_regex = re.compile(r'^(\d*)-(\d*)$')
 
 def verify_token(token):
     return token in authorized_tokens
@@ -69,19 +71,45 @@ async def handler(request):
     max_size = 0
     for i in messages:
         max_size += _get_file_info(i).size
-    http_range = request.http_range
-    offset = http_range.start or 0
-    if offset < 0:
-        end = offset
-        offset = max_size - offset
-    else:
-        end = http_range.stop
-        if end is None:
-            end = max_size
-        elif end > max_size:
-            return web.Response(status=416, text='Range end size is bigger than file sizes', headers={'Content-Range': f'bytes */{max_size}'})
-        else:
-            end -= 1
+    status = 200
+    offset = 0
+    end = max_size - 1
+    if range := request.headers.get('Range'):
+        if range.startswith('bytes='):
+            match = bytes_regex.match(range[6:])
+            if not match:
+                return web.Response(status=416,
+                    text='Failed to match range with regex',
+                    headers={
+                        'Content-Range': f'bytes */{max_size}',
+                        'Accept-Ranges': 'bytes'
+                    }
+                )
+            tmp_offset = match.group(1)
+            tmp_end = match.group(2)
+            if tmp_offset:
+                offset = int(tmp_offset)
+                if tmp_end:
+                    end = int(tmp_end)
+            elif tmp_end:
+                offset = max_size - int(tmp_end)
+            else:
+                return web.Response(status=400,
+                    text='Both range sides missing',
+                    headers={
+                        'Content-Range': f'bytes */{max_size}',
+                        'Accept-Ranges': 'bytes'
+                    }
+                )
+            if offset < 0 or offset > max_size or end >= max_size or offset > end:
+                return web.Response(status=416,
+                    text='Range not satisfiable',
+                    headers={
+                        'Content-Range': f'bytes */{max_size}',
+                        'Accept-Ranges': 'bytes'
+                    }
+                )
+            status = 206
     length = end - offset + 1
 
     async def download():
@@ -101,7 +129,7 @@ async def handler(request):
                     break
             tmp_offset = 0
 
-    return web.Response(status=206 if (length != max_size) else 200,
+    return web.Response(status=status,
         body=download(),
         headers={
             'Content-Range': f'bytes {offset}-{end}/{max_size}',