Add hmac

example-config.yaml

@@ -2,7 +2,12 @@ telegram:
     api_id: 0
     api_hash: https://my.telegram.org
     bot_token: https://t.me/BotFather
-# If authorized tokens does not exist, no authorization is required
+# If authorized tokens and hmac does not exist, no authorization is required
 authorized_tokens:
     - DIi4aXHn440PTPXJE1yVyIoU2L4bLGiyjC1Fd7usKAMZYWVcp5p0P792G4YlbNnIcWCLypXbUFJkVzKqhh0AkJYSWqJbsAy8TjA
     - RCywhzEkwqWmDjsqhBdNMwO1cQNi72SWsvAKZxdoeFylbioqKCZjEKyjgDXkYn6xeMQLj4dDq6QoonVAU1b1MDyOUX9CvT5W4MP
+hmac:
+    - key: QgTz0hoawBdIMUEULxMa0QXHOPjBujO5Vl3Liao87H0687JPDbd6ixB08vWWZOqXpOKv7kcDb1QGkkYeAirlrnvzkC7ZYwN63lR
+      digest: sha512
+    - key: ni1KHr52uvpO1tEBFXIYylha2CTfnTafY2b3qE0V7qNMj5eYtSMN0VGnuJHoFxDLSe2lUebyH3d2zA9q5IDfAs9qYWQ27mFV099
+      digest: sha256

streamtg.py

@@ -3,6 +3,7 @@ logging.basicConfig(level=logging.INFO)
 
 import os
 import yaml
+import hmac
 from aiohttp import web
 from telethon import TelegramClient
 from telethon.utils import _get_file_info
@@ -13,27 +14,52 @@ with open('config.yaml') as file:
 
 client = TelegramClient('streamtg', config['telegram']['api_id'], config['telegram']['api_hash'])
 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)
 
+def verify_token(token):
+    return token in authorized_tokens
+
+def verify_hmac(hexdigest, chat_id, message_ids):
+    text = f'{chat_id}|{"|".join(message_ids)}'.encode()
+    for i in hmacs:
+        i = i.copy()
+        i.update(text)
+        if hmac.compare_digest(hexdigest, i.hexdigest()):
+            return True
+    return False
+
 async def handler(request):
     query = request.query
-    if authorized_tokens:
-        if 'token' not in query:
-            return web.Response(status=401, text='Unauthorized')
-        if query['token'] not in authorized_tokens:
-            return web.Response(status=403, text='Forbidden')
+    token = query.get('token')
+    hexdigest = query.get('hmac')
+    if not token and not hexdigest and (authorized_tokens or hmacs):
+        return web.Response(status=401, text='Missing token or hmac')
     if 'chat_id' not in query:
         return web.Response(status=400, text='Missing chat_id')
     chat_id = query['chat_id']
     try:
         chat_id = int(chat_id)
     except ValueError:
-        pass
+        try:
+            chat_id = await client.get_peer_id(chat_id)
+        except BaseException:
+            if authorized_tokens or hmacs:
+                logging.exception('Exception occured while getting chat id of %s, returning 403 to hide known chats', chat_id)
+                return web.Response(status=403, text='Forbidden')
+            raise
     if 'message_id' not in query:
         return web.Response(status=400, text='Missing message_id')
     message_ids = query.getall('message_id')
     if any(True for i in message_ids if not i.isnumeric() or i == '0'):
         return web.Response(status=400, text='Invalid message_id')
+    if authorized_tokens or hmacs:
+        if not token or not verify_token(token):
+            if hexdigest:
+                if not verify_hmac(hexdigest, chat_id, message_ids):
+                    return web.Response(status=403, text='Forbidden')
+            else:
+                return web.Response(status=403, text='Forbidden')
     message_ids = list(map(int, message_ids))
     messages = await client.get_messages(chat_id, ids=message_ids)
     if any(True for i in messages if i is None):