streamtg/streamtg.py

124 lines
4.6 KiB
Python

import logging
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
from telethon.client.downloads import _GenericDownloadIter
with open('config.yaml') as file:
config = yaml.safe_load(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
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:
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):
return web.Response(status=400, text='At least one of the messages does not exist')
if any(True for i in messages if not i.media):
return web.Response(status=400, text='At least one of the messages do not contain media')
http_range = request.http_range
offset = http_range.start or 0
end = http_range.stop
max_size = 0
for i in messages:
max_size += _get_file_info(i).size
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}'})
length = end - offset
async def download():
tmp_offset = offset
tmp_length = length
for i in messages:
if tmp_length < 1:
break
size = _get_file_info(i).size
if tmp_offset > size:
tmp_offset -= size
continue
stream = client._iter_download(i, offset=tmp_offset, request_size=131072, msg_data=(chat_id, i.id))
if isinstance(stream, _GenericDownloadIter):
__load_next_chunk = stream._load_next_chunk
def _load_next_chunk(mask=131071):
return __load_next_chunk(mask=mask)
stream._load_next_chunk = _load_next_chunk
async for chunk in stream:
yield chunk[:tmp_length]
tmp_length -= len(chunk)
if tmp_length < 1:
break
return web.Response(status=206 if (end - offset != max_size) else 200,
body=download(),
headers={
'Content-Range': f'bytes {offset}-{end}/{max_size}',
'Content-Length': str(end - offset),
'Accept-Ranges': 'bytes'
}
)
app = web.Application()
app.add_routes([web.get('/', handler)])
async def main():
await client.start(bot_token=config['telegram'].get('bot_token'))
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '0.0.0.0', port)
await site.start()
await client.run_until_disconnected()
client.loop.run_until_complete(main())