Skip to content

Commit 315ddee

Browse files
author
Your Name
committed
1. Fix download when multi client is disabled
2. File Unique Key, to prevent strangers to access your files 3. Performance fixes (multi client bots don't search for updates, import plugins via Pyrogram smart Plugins) 4. Black code formatting
1 parent 078aea4 commit 315ddee

15 files changed

Lines changed: 193 additions & 170 deletions

WebStreamer/__init__.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,10 @@
66
from .vars import Var
77
from WebStreamer.bot.clients import StreamBot
88

9-
print('\n')
10-
print('------------------- Initializing Telegram Bot -------------------')
9+
print("\n")
10+
print("------------------- Initializing Telegram Bot -------------------")
1111

1212
StreamBot.start()
1313
bot_info = StreamBot.get_me()
14-
__version__ = 1.06
14+
__version__ = 1.08
1515
StartTime = time.time()

WebStreamer/__main__.py

Lines changed: 21 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,9 @@
11
# This file is a part of TG-FileStreamBot
22
# Coding : Jyothis Jayanth [@EverythingSuckz]
33

4-
import sys
5-
import glob
64
import asyncio
75
import logging
8-
import importlib
9-
import importlib
106
from aiohttp import web
11-
from pathlib import Path
12-
from pathlib import Path
137
from pyrogram import idle
148
from WebStreamer import bot_info
159
from WebStreamer.vars import Var
@@ -19,59 +13,47 @@
1913

2014

2115
logging.basicConfig(
22-
level=logging.INFO,
23-
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
16+
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
2417
)
2518
logging.getLogger("pyrogram").setLevel(logging.ERROR)
2619
logging.getLogger("aiohttp").setLevel(logging.ERROR)
2720
logging.getLogger("aiohttp.web").setLevel(logging.ERROR)
2821

2922
loop = asyncio.get_event_loop()
3023

31-
_path = "WebStreamer/bot/plugins/*.py"
32-
files = glob.glob(_path)
3324

3425
async def start_services():
35-
print('----------------------------- DONE -----------------------------')
36-
print('\n')
37-
print('----------------------------- INITIALIZING CLIENTS -----------------------------')
26+
print("----------------------------- DONE -----------------------------")
27+
print("\n")
28+
print(
29+
"----------------------------- Initializing Clients -----------------------------"
30+
)
3831
await initialize_clients()
39-
print('----------------------------- DONE -----------------------------')
40-
print('--------------------------- Importing ---------------------------')
41-
for name in files:
42-
with open(name) as a:
43-
path_ = Path(a.name)
44-
plugin_name = path_.stem.replace(".py", "")
45-
plugins_dir = Path(f"WebStreamer/bot/plugins/{plugin_name}.py")
46-
import_path = ".plugins.{}".format(plugin_name)
47-
spec = importlib.util.spec_from_file_location(import_path, plugins_dir)
48-
load = importlib.util.module_from_spec(spec)
49-
spec.loader.exec_module(load)
50-
sys.modules["WebStreamer.bot.plugins." + plugin_name] = load
51-
print("Imported => " + plugin_name)
32+
print("----------------------------- DONE -----------------------------")
5233
if Var.ON_HEROKU:
53-
print('------------------ Starting Keep Alive Service ------------------')
54-
print('\n')
34+
print("------------------ Starting Keep Alive Service ------------------")
35+
print("\n")
5536
await asyncio.create_task(ping_server())
56-
print('-------------------- Initalizing Web Server --------------------')
37+
print("-------------------- Initalizing Web Server --------------------")
5738
app = web.AppRunner(await web_server())
5839
await app.setup()
5940
bind_address = "0.0.0.0" if Var.ON_HEROKU else Var.BIND_ADDRESS
6041
await web.TCPSite(app, bind_address, Var.PORT).start()
61-
print('----------------------------- DONE -----------------------------')
62-
print('\n')
63-
print('----------------------- Service Started -----------------------')
64-
print(' bot =>> {}'.format(bot_info.first_name))
42+
print("----------------------------- DONE -----------------------------")
43+
print("\n")
44+
print("----------------------- Service Started -----------------------")
45+
print(" bot =>> {}".format(bot_info.first_name))
6546
if bot_info.dc_id:
66-
print(' DC ID =>> {}'.format(str(bot_info.dc_id)))
67-
print(' server ip =>> {}'.format(bind_address, Var.PORT))
47+
print(" DC ID =>> {}".format(str(bot_info.dc_id)))
48+
print(" server ip =>> {}".format(bind_address, Var.PORT))
6849
if Var.ON_HEROKU:
69-
print(' app running on =>> {}'.format(Var.FQDN))
70-
print('---------------------------------------------------------------')
50+
print(" app running on =>> {}".format(Var.FQDN))
51+
print("---------------------------------------------------------------")
7152
await idle()
7253

73-
if __name__ == '__main__':
54+
55+
if __name__ == "__main__":
7456
try:
7557
loop.run_until_complete(start_services())
7658
except KeyboardInterrupt:
77-
logging.info('----------------------- Service Stopped -----------------------')
59+
logging.info("----------------------- Service Stopped -----------------------")

WebStreamer/bot/__init__.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,18 @@
44

55
from ..vars import Var
66
from pyrogram import Client
7+
from os import getcwd
8+
79
StreamBot = Client(
8-
session_name= 'WebStreamer',
10+
session_name="WebStreamer",
911
api_id=Var.API_ID,
1012
api_hash=Var.API_HASH,
13+
workdir="WebStreamer",
14+
plugins={"root": "WebStreamer/bot/plugins"},
1115
bot_token=Var.BOT_TOKEN,
1216
sleep_threshold=Var.SLEEP_THRESHOLD,
13-
workers=Var.WORKERS
17+
workers=Var.WORKERS,
1418
)
1519

1620
multi_clients = {}
17-
work_loads = {}
21+
work_loads = {}

WebStreamer/bot/clients.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,23 @@
88

99

1010
async def initialize_clients():
11+
multi_clients[0] = StreamBot
12+
work_loads[0] = 0
1113
if Var.MULTI_CLIENT:
1214
all_tokens = TokenParser().parse_from_env()
13-
multi_clients[0] = StreamBot
14-
work_loads[0] = 0
1515
for client_id, token in all_tokens.items():
16-
instace = Client(
17-
session_name= ':memory:',
16+
instance = Client(
17+
session_name=":memory:",
1818
api_id=Var.API_ID,
1919
api_hash=Var.API_HASH,
2020
bot_token=token,
2121
sleep_threshold=Var.SLEEP_THRESHOLD,
22-
workers=Var.WORKERS
22+
no_updates=True,
2323
)
24-
multi_clients[client_id] = await instace.start()
24+
try:
25+
multi_clients[client_id] = await instance.start()
26+
except Exception as e:
27+
print(f"Failed starting Client - {client_id}; Error: {e}")
28+
continue
2529
work_loads[client_id] = 0
26-
print(f"Started - Client {client_id}")
30+
print(f"Started - Client {client_id}")

WebStreamer/bot/plugins/start.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@
55
from pyrogram.types import Message
66
from WebStreamer.bot import StreamBot
77

8-
@StreamBot.on_message(filters.command(['start', 'help']))
8+
9+
@StreamBot.on_message(filters.command(["start", "help"]))
910
async def start(_, m: Message):
10-
await m.reply(f'Hi {m.from_user.mention(style="md")}, Send me a file to get an instant stream link.\nNote that this bot is still in beta and may not work as expected.\n')
11+
await m.reply(
12+
f'Hi {m.from_user.mention(style="md")}, Send me a file to get an instant stream link.\nNote that this bot is still in beta and may not work as expected.\n'
13+
)

WebStreamer/bot/plugins/stream.py

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
from urllib.parse import quote_plus
88
from WebStreamer.bot import StreamBot
99
from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton
10+
from WebStreamer.utils.file_id import get_unique_id
11+
1012

1113
def detect_type(m: Message):
1214
if m.document:
@@ -17,21 +19,25 @@ def detect_type(m: Message):
1719
return m.audio
1820
else:
1921
return
20-
2122

22-
@StreamBot.on_message(filters.private & (filters.document | filters.video | filters.audio | filters.animation), group=4)
23+
24+
@StreamBot.on_message(
25+
filters.private
26+
& (filters.document | filters.video | filters.audio | filters.animation),
27+
group=4,
28+
)
2329
async def media_receive_handler(_, m: Message):
2430
file = detect_type(m)
25-
file_name = ''
31+
file_name = ""
2632
if file:
2733
file_name = file.file_name
2834
log_msg = await m.forward(chat_id=Var.BIN_CHANNEL)
29-
stream_link = f"{Var.URL}{log_msg.message_id}"
30-
if file_name:
31-
stream_link += f'/{quote_plus(file_name)}'
35+
stream_link = f"{Var.URL}{log_msg.message_id}-{get_unique_id(log_msg)}"
3236
logging.info(f"Generated link: {stream_link} for {m.from_user.first_name}")
3337
await m.reply_text(
34-
text = "<code>{}</code>".format(stream_link),
35-
quote = True,
36-
reply_markup = InlineKeyboardMarkup([[InlineKeyboardButton('Open', url=stream_link)]])
37-
)
38+
text="<code>{}</code>".format(stream_link),
39+
quote=True,
40+
reply_markup=InlineKeyboardMarkup(
41+
[[InlineKeyboardButton("Open", url=stream_link)]]
42+
),
43+
)

WebStreamer/server/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,4 @@
1010
async def web_server():
1111
web_app = web.Application(client_max_size=30000000)
1212
web_app.add_routes(routes)
13-
return web_app
13+
return web_app

WebStreamer/server/stream_routes.py

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -13,38 +13,42 @@
1313
from WebStreamer import StartTime, __version__, bot_info
1414
from WebStreamer.bot.clients import multi_clients, work_loads
1515
from WebStreamer.utils.time_format import get_readable_time
16+
from WebStreamer.utils.file_id import get_unique_id
1617
from WebStreamer.utils.custom_dl import TGCustomYield, chunk_size, offset_fix
1718

1819

1920
routes = web.RouteTableDef()
2021

22+
2123
@routes.get("/", allow_head=True)
2224
async def root_route_handler(request):
23-
return web.json_response({"server_status": "running",
24-
"uptime": get_readable_time(time.time() - StartTime),
25-
"telegram_bot": '@'+ bot_info.username,
26-
"connected_bots": len(multi_clients),
27-
"loads": work_loads,
28-
"version": __version__})
25+
return web.json_response(
26+
{
27+
"server_status": "running",
28+
"uptime": get_readable_time(time.time() - StartTime),
29+
"telegram_bot": "@" + bot_info.username,
30+
"connected_bots": len(multi_clients),
31+
"loads": work_loads,
32+
"version": __version__,
33+
}
34+
)
2935

3036

31-
@routes.get(r"/{message_id:\S+}")
37+
@routes.get(r"/{message_id:\S+}-{secure_hash:\w+}")
3238
async def stream_handler(request: web.Request):
3339
try:
34-
message_id = request.match_info['message_id']
35-
message_id = int(re.search(r'(\d+)(?:\/\S+)?', message_id).group(1))
36-
return await media_streamer(request, message_id)
40+
message_id = request.match_info["message_id"]
41+
message_id = int(re.search(r"(\d+)(?:\/\S+)?", message_id).group(1))
42+
secure_hash = request.match_info["secure_hash"]
43+
return await media_streamer(request, message_id, secure_hash)
3744
except ValueError:
3845
raise web.HTTPNotFound
3946
except AttributeError:
4047
pass
4148

4249

43-
44-
45-
46-
async def media_streamer(request: web.Request, message_id: int):
47-
range_header = request.headers.get('Range', 0)
50+
async def media_streamer(request: web.Request, message_id: int, secure_hash: str):
51+
range_header = request.headers.get("Range", 0)
4852
_index = min(work_loads, key=work_loads.get)
4953
faster_client = multi_clients[_index]
5054
work_loads[_index] += 1
@@ -53,11 +57,14 @@ async def media_streamer(request: web.Request, message_id: int):
5357

5458
tg_connect = TGCustomYield(faster_client)
5559
media_msg = await faster_client.get_messages(Var.BIN_CHANNEL, message_id)
60+
if get_unique_id(media_msg) != secure_hash:
61+
work_loads[_index] -= 1
62+
raise web.HTTPForbidden
5663
file_properties = await tg_connect.generate_file_properties(media_msg)
5764
file_size = file_properties.file_size
5865

5966
if range_header:
60-
from_bytes, until_bytes = range_header.replace('bytes=', '').split('-')
67+
from_bytes, until_bytes = range_header.replace("bytes=", "").split("-")
6168
from_bytes = int(from_bytes)
6269
until_bytes = int(until_bytes) if until_bytes else file_size - 1
6370
else:
@@ -70,8 +77,9 @@ async def media_streamer(request: web.Request, message_id: int):
7077
first_part_cut = from_bytes - offset
7178
last_part_cut = (until_bytes % new_chunk_size) + 1
7279
part_count = math.ceil(req_length / new_chunk_size)
73-
body = tg_connect.yield_file(media_msg, offset, first_part_cut, last_part_cut, part_count,
74-
new_chunk_size)
80+
body = tg_connect.yield_file(
81+
media_msg, offset, first_part_cut, last_part_cut, part_count, new_chunk_size
82+
)
7583

7684
mime_type = file_properties.mime_type
7785
file_name = file_properties.file_name
@@ -87,7 +95,7 @@ async def media_streamer(request: web.Request, message_id: int):
8795
mime_type = mimetypes.guess_type(file_properties.file_name)
8896
else:
8997
mime_type = "application/octet-stream"
90-
file_name = f"{secrets.token_hex(2)}.unknown"
98+
file_name = f"{secrets.token_hex(2)}.unknown"
9199
if "video/" in mime_type or "audio/" in mime_type:
92100
dispo = "inline"
93101
return_resp = web.Response(
@@ -98,7 +106,7 @@ async def media_streamer(request: web.Request, message_id: int):
98106
"Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}",
99107
"Content-Disposition": f'{dispo}; filename="{file_name}"',
100108
"Accept-Ranges": "bytes",
101-
}
109+
},
102110
)
103111

104112
if return_resp.status == 200:

WebStreamer/utils/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
# This file is a part of TG-FileStreamBot
2-
# Coding : Jyothis Jayanth [@EverythingSuckz]
2+
# Coding : Jyothis Jayanth [@EverythingSuckz]

WebStreamer/utils/config_parser.py

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,19 @@
11
import os
2-
import re
32
from typing import Dict, Optional
43

4+
55
class TokenParser:
66
def __init__(self, config_file: Optional[str] = None):
77
self.tokens = {}
88
self.config_file = config_file
99

10-
# def parse_from_dotenv(self):
11-
# with open(self.config_file, 'r') as f:
12-
# for line in f.readlines():
13-
# if line.startswith('#'):
14-
# continue
15-
# elif match := re.search(r'^MULTI_TOKEN(\d+)=(\S+)$', line):
16-
# count = int(match.group(1))
17-
# token = match.group(2)
18-
# if not count and not token:
19-
# continue
20-
# self.tokens[count] = token
21-
# return self.config
22-
2310
def parse_from_env(self) -> Dict[int, str]:
2411
self.tokens = dict(
25-
(c+1, t) for c, (_, t) in enumerate(filter(
26-
lambda n: n[0].startswith('MULTI_TOKEN'),
27-
sorted(os.environ.items())
28-
))
12+
(c + 1, t)
13+
for c, (_, t) in enumerate(
14+
filter(
15+
lambda n: n[0].startswith("MULTI_TOKEN"), sorted(os.environ.items())
16+
)
17+
)
2918
)
30-
return self.tokens
19+
return self.tokens

0 commit comments

Comments
 (0)