Fyaz Mohammed commited on
Commit
2fe1440
·
unverified ·
1 Parent(s): 0dce97e
Thunder/__init__.py CHANGED
@@ -3,4 +3,4 @@
3
  import time
4
 
5
  StartTime = time.time()
6
- __version__ = "2.0.0"
 
3
  import time
4
 
5
  StartTime = time.time()
6
+ __version__ = "2.1.0"
Thunder/__main__.py CHANGED
@@ -98,8 +98,8 @@ async def start_services():
98
  print_banner()
99
  print("╔════════════════ INITIALIZING BOT SERVICES ════════════════╗")
100
 
101
- print(" ▶ Starting Telegram Bot initialization...")
102
- try:
103
  try:
104
  await StreamBot.start()
105
  except FloodWait as e:
@@ -114,11 +114,14 @@ async def start_services():
114
  await asyncio.sleep(e.value)
115
  bot_info = await StreamBot.get_me()
116
 
117
- StreamBot.username = bot_info.username
118
- print(f" ✓ Bot initialized successfully as @{StreamBot.username}")
119
-
120
- await set_commands()
121
- print(" ✓ Bot commands set successfully.")
 
 
 
122
 
123
  restart_message_data = await db.get_restart_message()
124
  if restart_message_data:
 
98
  print_banner()
99
  print("╔════════════════ INITIALIZING BOT SERVICES ════════════════╗")
100
 
101
+ print(" ▶ Starting Telegram Bot initialization...")
102
+ try:
103
  try:
104
  await StreamBot.start()
105
  except FloodWait as e:
 
114
  await asyncio.sleep(e.value)
115
  bot_info = await StreamBot.get_me()
116
 
117
+ StreamBot.username = bot_info.username
118
+ print(f" ✓ Bot initialized successfully as @{StreamBot.username}")
119
+
120
+ await db.ensure_indexes()
121
+ print(" ✓ Database indexes ensured.")
122
+
123
+ await set_commands()
124
+ print(" ✓ Bot commands set successfully.")
125
 
126
  restart_message_data = await db.get_restart_message()
127
  if restart_message_data:
Thunder/bot/plugins/stream.py CHANGED
@@ -6,15 +6,16 @@ from typing import Any, Dict, Optional
6
 
7
  from pyrogram import Client, enums, filters
8
  from pyrogram.errors import FloodWait, MessageNotModified, MessageDeleteForbidden, MessageIdInvalid
9
- from pyrogram.types import (InlineKeyboardButton, InlineKeyboardMarkup,
10
- Message)
11
-
12
- from Thunder.bot import StreamBot
13
- from Thunder.utils.bot_utils import (gen_links, is_admin, log_newusr, notify_own,
14
- reply_user_err)
15
- from Thunder.utils.database import db
16
- from Thunder.utils.decorators import (check_banned, get_shortener_status,
17
- require_token)
 
18
  from Thunder.utils.force_channel import force_channel_check
19
  from Thunder.utils.logger import logger
20
  from Thunder.utils.messages import (
@@ -71,30 +72,55 @@ async def validate_request_common(client: Client, message: Message) -> Optional[
71
  return await get_shortener_status(client, message)
72
 
73
 
74
- async def send_channel_links(target_msg: Message, links: Dict[str, Any], source_info: str, source_id: int):
75
- try:
76
- await target_msg.reply_text(
77
- MSG_NEW_FILE_REQUEST.format(
78
- source_info=source_info,
79
- id_=source_id,
80
- online_link=links['online_link'],
81
- stream_link=links['stream_link']
82
- ),
83
- disable_web_page_preview=True,
84
- quote=True
85
- )
86
- except FloodWait as e:
87
- await asyncio.sleep(e.value)
88
- await target_msg.reply_text(
89
- MSG_NEW_FILE_REQUEST.format(
90
- source_info=source_info,
91
- id_=source_id,
92
- online_link=links['online_link'],
93
- stream_link=links['stream_link']
94
- ),
95
- disable_web_page_preview=True,
96
- quote=True
97
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
 
99
 
100
  async def safe_edit_message(message: Message, text: str, **kwargs):
@@ -313,17 +339,29 @@ async def channel_receive_handler(bot: Client, msg: Message):
313
  f"({message.chat.title or 'Unknown'}). Ignoring message.")
314
  return
315
 
316
- try:
317
- stored_msg = await fwd_media(message)
318
- if not stored_msg:
319
- logger.error(
320
- f"Failed to forward media from channel {message.chat.id}. Ignoring.")
321
- return
322
- shortener_val = await get_shortener_status(client, message)
323
- links = await gen_links(stored_msg, shortener=shortener_val)
324
- source_info = message.chat.title or "Unknown Channel"
325
-
326
- if notification_msg:
 
 
 
 
 
 
 
 
 
 
 
 
327
  try:
328
  try:
329
  await notification_msg.edit_text(
@@ -346,11 +384,23 @@ async def channel_receive_handler(bot: Client, msg: Message):
346
  ),
347
  disable_web_page_preview=True
348
  )
349
- except Exception as e:
350
- logger.error(f"Error editing notification message with links: {e}", exc_info=True)
351
- await send_channel_links(stored_msg, links, source_info, message.chat.id)
352
- else:
353
- await send_channel_links(stored_msg, links, source_info, message.chat.id)
 
 
 
 
 
 
 
 
 
 
 
 
354
 
355
  try:
356
  try:
@@ -389,16 +439,28 @@ async def process_single(
389
  shortener_val: bool,
390
  original_request_msg: Optional[Message] = None,
391
  notification_msg: Optional[Message] = None
392
- ):
393
- try:
394
- stored_msg = await fwd_media(file_msg)
395
- if not stored_msg:
396
- logger.error(f"Failed to forward media for message {file_msg.id}. Skipping.")
397
- return None
398
- links = await gen_links(stored_msg, shortener=shortener_val)
399
- if notification_msg:
400
- await safe_edit_message(
401
- notification_msg,
 
 
 
 
 
 
 
 
 
 
 
 
402
  MSG_LINKS.format(
403
  file_name=links['media_name'],
404
  file_size=links['media_size'],
@@ -421,35 +483,29 @@ async def process_single(
421
  if not source_info:
422
  source_info = f"@{source_msg.from_user.username}" if source_msg.from_user.username else "Unknown User"
423
  source_id = source_msg.from_user.id
424
- elif source_msg.chat.type == enums.ChatType.CHANNEL:
425
- source_info = source_msg.chat.title or "Unknown Channel"
426
- source_id = source_msg.chat.id
427
- if source_info and source_id:
428
- try:
429
- await stored_msg.reply_text(
430
- MSG_NEW_FILE_REQUEST.format(
431
- source_info=source_info,
432
- id_=source_id,
433
- online_link=links['online_link'],
434
- stream_link=links['stream_link']
435
- ),
436
- disable_web_page_preview=True,
437
- quote=True
438
- )
439
- except FloodWait as e:
440
- await asyncio.sleep(e.value)
441
- await stored_msg.reply_text(
442
- MSG_NEW_FILE_REQUEST.format(
443
- source_info=source_info,
444
- id_=source_id,
445
- online_link=links['online_link'],
446
- stream_link=links['stream_link']
447
- ),
448
- disable_web_page_preview=True,
449
- quote=True
450
- )
451
- if status_msg:
452
- await safe_delete_message(status_msg)
453
  return links
454
  except Exception as e:
455
  logger.error(f"Error processing single file for message {file_msg.id}: {e}", exc_info=True)
 
6
 
7
  from pyrogram import Client, enums, filters
8
  from pyrogram.errors import FloodWait, MessageNotModified, MessageDeleteForbidden, MessageIdInvalid
9
+ from pyrogram.types import (InlineKeyboardButton, InlineKeyboardMarkup,
10
+ Message)
11
+
12
+ from Thunder.bot import StreamBot
13
+ from Thunder.utils.bot_utils import (gen_canonical_links, gen_links, is_admin,
14
+ log_newusr, notify_own, reply_user_err)
15
+ from Thunder.utils.canonical_files import get_or_create_canonical_file
16
+ from Thunder.utils.database import db
17
+ from Thunder.utils.decorators import (check_banned, get_shortener_status,
18
+ require_token)
19
  from Thunder.utils.force_channel import force_channel_check
20
  from Thunder.utils.logger import logger
21
  from Thunder.utils.messages import (
 
72
  return await get_shortener_status(client, message)
73
 
74
 
75
+ async def send_channel_links(
76
+ links: Dict[str, Any],
77
+ source_info: str,
78
+ source_id: int,
79
+ *,
80
+ target_msg: Optional[Message] = None,
81
+ reply_to_message_id: Optional[int] = None
82
+ ):
83
+ try:
84
+ text = MSG_NEW_FILE_REQUEST.format(
85
+ source_info=source_info,
86
+ id_=source_id,
87
+ online_link=links['online_link'],
88
+ stream_link=links['stream_link']
89
+ )
90
+ if target_msg:
91
+ await target_msg.reply_text(
92
+ text,
93
+ disable_web_page_preview=True,
94
+ quote=True
95
+ )
96
+ else:
97
+ await StreamBot.send_message(
98
+ chat_id=Var.BIN_CHANNEL,
99
+ text=text,
100
+ disable_web_page_preview=True,
101
+ reply_to_message_id=reply_to_message_id
102
+ )
103
+ except FloodWait as e:
104
+ await asyncio.sleep(e.value)
105
+ text = MSG_NEW_FILE_REQUEST.format(
106
+ source_info=source_info,
107
+ id_=source_id,
108
+ online_link=links['online_link'],
109
+ stream_link=links['stream_link']
110
+ )
111
+ if target_msg:
112
+ await target_msg.reply_text(
113
+ text,
114
+ disable_web_page_preview=True,
115
+ quote=True
116
+ )
117
+ else:
118
+ await StreamBot.send_message(
119
+ chat_id=Var.BIN_CHANNEL,
120
+ text=text,
121
+ disable_web_page_preview=True,
122
+ reply_to_message_id=reply_to_message_id
123
+ )
124
 
125
 
126
  async def safe_edit_message(message: Message, text: str, **kwargs):
 
339
  f"({message.chat.title or 'Unknown'}). Ignoring message.")
340
  return
341
 
342
+ try:
343
+ shortener_val = await get_shortener_status(client, message)
344
+ canonical_record, stored_msg, _ = await get_or_create_canonical_file(message, fwd_media)
345
+ if canonical_record:
346
+ links = await gen_canonical_links(
347
+ file_name=canonical_record["file_name"],
348
+ file_size=int(canonical_record.get("file_size", 0) or 0),
349
+ public_hash=canonical_record["public_hash"],
350
+ shortener=shortener_val
351
+ )
352
+ reply_to_message_id = int(canonical_record["canonical_message_id"])
353
+ else:
354
+ if not stored_msg:
355
+ stored_msg = await fwd_media(message)
356
+ if not stored_msg:
357
+ logger.error(
358
+ f"Failed to forward media from channel {message.chat.id}. Ignoring.")
359
+ return
360
+ links = await gen_links(stored_msg, shortener=shortener_val)
361
+ reply_to_message_id = stored_msg.id
362
+ source_info = message.chat.title or "Unknown Channel"
363
+
364
+ if notification_msg:
365
  try:
366
  try:
367
  await notification_msg.edit_text(
 
384
  ),
385
  disable_web_page_preview=True
386
  )
387
+ except Exception as e:
388
+ logger.error(f"Error editing notification message with links: {e}", exc_info=True)
389
+ await send_channel_links(
390
+ links,
391
+ source_info,
392
+ message.chat.id,
393
+ target_msg=stored_msg,
394
+ reply_to_message_id=reply_to_message_id
395
+ )
396
+ else:
397
+ await send_channel_links(
398
+ links,
399
+ source_info,
400
+ message.chat.id,
401
+ target_msg=stored_msg,
402
+ reply_to_message_id=reply_to_message_id
403
+ )
404
 
405
  try:
406
  try:
 
439
  shortener_val: bool,
440
  original_request_msg: Optional[Message] = None,
441
  notification_msg: Optional[Message] = None
442
+ ):
443
+ try:
444
+ canonical_record, stored_msg, _ = await get_or_create_canonical_file(file_msg, fwd_media)
445
+ if canonical_record:
446
+ links = await gen_canonical_links(
447
+ file_name=canonical_record["file_name"],
448
+ file_size=int(canonical_record.get("file_size", 0) or 0),
449
+ public_hash=canonical_record["public_hash"],
450
+ shortener=shortener_val
451
+ )
452
+ canonical_reply_id = int(canonical_record["canonical_message_id"])
453
+ else:
454
+ if not stored_msg:
455
+ stored_msg = await fwd_media(file_msg)
456
+ if not stored_msg:
457
+ logger.error(f"Failed to forward media for message {file_msg.id}. Skipping.")
458
+ return None
459
+ links = await gen_links(stored_msg, shortener=shortener_val)
460
+ canonical_reply_id = stored_msg.id
461
+ if notification_msg:
462
+ await safe_edit_message(
463
+ notification_msg,
464
  MSG_LINKS.format(
465
  file_name=links['media_name'],
466
  file_size=links['media_size'],
 
483
  if not source_info:
484
  source_info = f"@{source_msg.from_user.username}" if source_msg.from_user.username else "Unknown User"
485
  source_id = source_msg.from_user.id
486
+ elif source_msg.chat.type == enums.ChatType.CHANNEL:
487
+ source_info = source_msg.chat.title or "Unknown Channel"
488
+ source_id = source_msg.chat.id
489
+ if source_info and source_id:
490
+ try:
491
+ await send_channel_links(
492
+ links,
493
+ source_info,
494
+ source_id,
495
+ target_msg=stored_msg,
496
+ reply_to_message_id=canonical_reply_id
497
+ )
498
+ except FloodWait as e:
499
+ await asyncio.sleep(e.value)
500
+ await send_channel_links(
501
+ links,
502
+ source_info,
503
+ source_id,
504
+ target_msg=stored_msg,
505
+ reply_to_message_id=canonical_reply_id
506
+ )
507
+ if status_msg:
508
+ await safe_delete_message(status_msg)
 
 
 
 
 
 
509
  return links
510
  except Exception as e:
511
  logger.error(f"Error processing single file for message {file_msg.id}: {e}", exc_info=True)
Thunder/server/stream_routes.py CHANGED
@@ -1,19 +1,27 @@
1
  # Thunder/server/stream_routes.py
2
 
3
- import re
4
- import secrets
5
- import time
6
- from urllib.parse import quote, unquote
7
-
8
- from aiohttp import web
9
-
10
- from Thunder import __version__, StartTime
11
- from Thunder.bot import StreamBot, multi_clients, work_loads
12
- from Thunder.server.exceptions import FileNotFound, InvalidHash
13
- from Thunder.utils.custom_dl import ByteStreamer
14
- from Thunder.utils.logger import logger
15
- from Thunder.utils.render_template import render_page
16
- from Thunder.utils.time_format import get_readable_time
 
 
 
 
 
 
 
 
17
 
18
  routes = web.RouteTableDef()
19
 
@@ -24,8 +32,9 @@ RANGE_REGEX = re.compile(r"bytes=(?P<start>\d*)-(?P<end>\d*)")
24
  PATTERN_HASH_FIRST = re.compile(
25
  rf"^([a-zA-Z0-9_-]{{{SECURE_HASH_LENGTH}}})(\d+)(?:/.*)?$")
26
  PATTERN_ID_FIRST = re.compile(r"^(\d+)(?:/.*)?$")
27
- VALID_HASH_REGEX = re.compile(r'^[a-zA-Z0-9_-]+$')
28
- VALID_DISPOSITIONS = {"inline", "attachment"}
 
29
 
30
  CORS_HEADERS = {
31
  "Access-Control-Allow-Origin": "*",
@@ -43,7 +52,7 @@ def get_streamer(client_id: int) -> ByteStreamer:
43
  return streamers[client_id]
44
 
45
 
46
- def parse_media_request(path: str, query: dict) -> tuple[int, str]:
47
  clean_path = unquote(path).strip('/')
48
 
49
  match = PATTERN_HASH_FIRST.match(clean_path)
@@ -70,7 +79,14 @@ def parse_media_request(path: str, query: dict) -> tuple[int, str]:
70
  except ValueError as e:
71
  raise InvalidHash(f"Invalid message ID format in path: {e}") from e
72
 
73
- raise InvalidHash("Invalid URL structure or missing hash")
 
 
 
 
 
 
 
74
 
75
 
76
  def select_optimal_client() -> tuple[int, ByteStreamer]:
@@ -96,7 +112,7 @@ def get_content_disposition(request: web.Request) -> str:
96
  return disposition if disposition in VALID_DISPOSITIONS else "attachment"
97
 
98
 
99
- def parse_range_header(range_header: str, file_size: int) -> tuple[int, int]:
100
  if not range_header:
101
  return 0, file_size - 1
102
 
@@ -124,7 +140,114 @@ def parse_range_header(range_header: str, file_size: int) -> tuple[int, int]:
124
  headers={"Content-Range": f"bytes */{file_size}"}
125
  )
126
 
127
- return start, end
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
128
 
129
 
130
  @routes.get("/", allow_head=True)
@@ -167,16 +290,49 @@ async def status_options(request: web.Request):
167
  })
168
 
169
 
170
- @routes.options(r"/{path:.+}")
171
- async def media_options(request: web.Request):
172
  return web.Response(headers={
173
  **CORS_HEADERS,
174
  "Access-Control-Max-Age": "86400"
175
- })
176
-
177
-
178
- @routes.get(r"/watch/{path:.+}", allow_head=True)
179
- async def media_preview(request: web.Request):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
180
  try:
181
  path = request.match_info["path"]
182
  message_id, secure_hash = parse_media_request(path, request.query)
@@ -204,113 +360,90 @@ async def media_preview(request: web.Request):
204
  except Exception as e:
205
  error_id = secrets.token_hex(6)
206
  logger.error(f"Preview error {error_id}: {e}", exc_info=True)
207
- raise web.HTTPInternalServerError(
208
- text=f"Server error occurred: {error_id}") from e
209
-
210
-
211
- @routes.get(r"/{path:.+}", allow_head=True)
212
- async def media_delivery(request: web.Request):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
213
  try:
214
  path = request.match_info["path"]
215
  message_id, secure_hash = parse_media_request(path, request.query)
216
 
217
  client_id, streamer = select_optimal_client()
218
 
219
- work_loads[client_id] += 1
220
-
221
- try:
222
- file_info = await streamer.get_file_info(message_id)
223
- if not file_info.get('unique_id'):
224
- raise FileNotFound("File unique ID not found in info.")
225
-
226
- if (file_info['unique_id'][:SECURE_HASH_LENGTH] !=
227
- secure_hash):
228
- raise InvalidHash(
229
- "Provided hash does not match file's unique ID.")
230
-
231
- file_size = file_info.get('file_size', 0)
232
- if file_size == 0:
233
- raise FileNotFound(
234
- "File size is reported as zero or unavailable.")
235
-
236
- range_header = request.headers.get("Range", "")
237
- start, end = parse_range_header(range_header, file_size)
238
- content_length = end - start + 1
239
-
240
- if start == 0 and end == file_size - 1:
241
- range_header = ""
242
-
243
- mime_type = (
244
- file_info.get('mime_type') or 'application/octet-stream')
245
-
246
- filename = file_info.get('file_name')
247
- if not filename:
248
- ext = mime_type.split('/')[-1] if '/' in mime_type else 'bin'
249
- ext_map = {'jpeg': 'jpg', 'mpeg': 'mp3', 'octet-stream': 'bin'}
250
- ext = ext_map.get(ext, ext)
251
- filename = f"file_{secrets.token_hex(4)}.{ext}"
252
-
253
- disposition = get_content_disposition(request)
254
-
255
- headers = {
256
- "Content-Type": mime_type,
257
- "Content-Length": str(content_length),
258
- "Content-Disposition": (
259
- f"{disposition}; filename*=UTF-8''{quote(filename)}"),
260
- "Accept-Ranges": "bytes",
261
- "Cache-Control": "public, max-age=31536000",
262
- "Connection": "keep-alive",
263
- "Access-Control-Allow-Origin": "*",
264
- "Access-Control-Allow-Headers": "Range, Content-Type, *",
265
- "Access-Control-Expose-Headers": (
266
- "Content-Length, Content-Range, Content-Disposition"),
267
- "X-Content-Type-Options": "nosniff"
268
- }
269
-
270
- if range_header:
271
- headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"
272
-
273
- if request.method == 'HEAD':
274
- work_loads[client_id] -= 1
275
- return web.Response(
276
- status=206 if range_header else 200,
277
- headers=headers
278
- )
279
-
280
- async def stream_generator():
281
- try:
282
- bytes_sent = 0
283
- bytes_to_skip = start % CHUNK_SIZE
284
-
285
- async for chunk in streamer.stream_file(
286
- message_id, offset=start, limit=content_length):
287
- if bytes_to_skip > 0:
288
- if len(chunk) <= bytes_to_skip:
289
- bytes_to_skip -= len(chunk)
290
- continue
291
- chunk = chunk[bytes_to_skip:]
292
- bytes_to_skip = 0
293
-
294
- remaining = content_length - bytes_sent
295
- if len(chunk) > remaining:
296
- chunk = chunk[:remaining]
297
-
298
- if chunk:
299
- yield chunk
300
- bytes_sent += len(chunk)
301
-
302
- if bytes_sent >= content_length:
303
- break
304
- finally:
305
- work_loads[client_id] -= 1
306
-
307
- return web.Response(
308
- status=206 if range_header else 200,
309
- body=stream_generator(),
310
- headers=headers
311
- )
312
-
313
- except (FileNotFound, InvalidHash):
314
  work_loads[client_id] -= 1
315
  raise
316
  except Exception as e:
 
1
  # Thunder/server/stream_routes.py
2
 
3
+ import re
4
+ import secrets
5
+ import time
6
+ from urllib.parse import quote, unquote
7
+
8
+ from aiohttp import web
9
+
10
+ from Thunder import __version__, StartTime
11
+ from Thunder.bot import StreamBot, multi_clients, work_loads
12
+ from Thunder.server.exceptions import FileNotFound, InvalidHash
13
+ from Thunder.utils.bot_utils import quote_media_name
14
+ from Thunder.utils.canonical_files import (
15
+ PUBLIC_HASH_LENGTH,
16
+ get_file_by_hash,
17
+ update_cached_file_id,
18
+ )
19
+ from Thunder.utils.custom_dl import ByteStreamer
20
+ from Thunder.utils.file_properties import get_media
21
+ from Thunder.utils.logger import logger
22
+ from Thunder.utils.render_template import render_media_page, render_page
23
+ from Thunder.utils.time_format import get_readable_time
24
+ from Thunder.vars import Var
25
 
26
  routes = web.RouteTableDef()
27
 
 
32
  PATTERN_HASH_FIRST = re.compile(
33
  rf"^([a-zA-Z0-9_-]{{{SECURE_HASH_LENGTH}}})(\d+)(?:/.*)?$")
34
  PATTERN_ID_FIRST = re.compile(r"^(\d+)(?:/.*)?$")
35
+ VALID_HASH_REGEX = re.compile(r'^[a-zA-Z0-9_-]+$')
36
+ VALID_PUBLIC_HASH_REGEX = re.compile(rf'^[0-9a-f]{{{PUBLIC_HASH_LENGTH}}}$')
37
+ VALID_DISPOSITIONS = {"inline", "attachment"}
38
 
39
  CORS_HEADERS = {
40
  "Access-Control-Allow-Origin": "*",
 
52
  return streamers[client_id]
53
 
54
 
55
+ def parse_media_request(path: str, query: dict) -> tuple[int, str]:
56
  clean_path = unquote(path).strip('/')
57
 
58
  match = PATTERN_HASH_FIRST.match(clean_path)
 
79
  except ValueError as e:
80
  raise InvalidHash(f"Invalid message ID format in path: {e}") from e
81
 
82
+ raise InvalidHash("Invalid URL structure or missing hash")
83
+
84
+
85
+ def validate_public_hash(public_hash: str) -> str:
86
+ secure_hash = public_hash.strip().lower()
87
+ if len(secure_hash) != PUBLIC_HASH_LENGTH or not VALID_PUBLIC_HASH_REGEX.match(secure_hash):
88
+ raise InvalidHash("Invalid canonical file hash")
89
+ return secure_hash
90
 
91
 
92
  def select_optimal_client() -> tuple[int, ByteStreamer]:
 
112
  return disposition if disposition in VALID_DISPOSITIONS else "attachment"
113
 
114
 
115
+ def parse_range_header(range_header: str, file_size: int) -> tuple[int, int]:
116
  if not range_header:
117
  return 0, file_size - 1
118
 
 
140
  headers={"Content-Range": f"bytes */{file_size}"}
141
  )
142
 
143
+ return start, end
144
+
145
+
146
+ def _resolve_unique_id(file_info: dict) -> str:
147
+ unique_id = file_info.get("unique_id") or file_info.get("file_unique_id")
148
+ if not unique_id:
149
+ raise FileNotFound("File unique ID not found in info.")
150
+ return unique_id
151
+
152
+
153
+ def _resolve_filename(file_info: dict, mime_type: str) -> str:
154
+ filename = file_info.get("file_name")
155
+ if filename:
156
+ return filename
157
+
158
+ ext = mime_type.split('/')[-1] if '/' in mime_type else 'bin'
159
+ ext_map = {'jpeg': 'jpg', 'mpeg': 'mp3', 'octet-stream': 'bin'}
160
+ ext = ext_map.get(ext, ext)
161
+ return f"file_{secrets.token_hex(4)}.{ext}"
162
+
163
+
164
+ async def _serve_media_response(
165
+ request: web.Request,
166
+ *,
167
+ file_info: dict,
168
+ streamer: ByteStreamer,
169
+ client_id: int,
170
+ media_ref: int | str,
171
+ fallback_message_id: int | None = None,
172
+ on_fallback_message=None
173
+ ):
174
+ file_size = int(file_info.get('file_size', 0) or 0)
175
+ if file_size == 0:
176
+ raise FileNotFound("File size is reported as zero or unavailable.")
177
+
178
+ range_header = request.headers.get("Range", "")
179
+ start, end = parse_range_header(range_header, file_size)
180
+ content_length = end - start + 1
181
+
182
+ if start == 0 and end == file_size - 1:
183
+ range_header = ""
184
+
185
+ mime_type = file_info.get('mime_type') or 'application/octet-stream'
186
+ filename = _resolve_filename(file_info, mime_type)
187
+ disposition = get_content_disposition(request)
188
+
189
+ headers = {
190
+ "Content-Type": mime_type,
191
+ "Content-Length": str(content_length),
192
+ "Content-Disposition": (
193
+ f"{disposition}; filename*=UTF-8''{quote(filename, safe='')}"),
194
+ "Accept-Ranges": "bytes",
195
+ "Cache-Control": "public, max-age=31536000",
196
+ "Connection": "keep-alive",
197
+ "Access-Control-Allow-Origin": "*",
198
+ "Access-Control-Allow-Headers": "Range, Content-Type, *",
199
+ "Access-Control-Expose-Headers": (
200
+ "Content-Length, Content-Range, Content-Disposition"),
201
+ "X-Content-Type-Options": "nosniff"
202
+ }
203
+
204
+ if range_header:
205
+ headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"
206
+
207
+ if request.method == 'HEAD':
208
+ work_loads[client_id] -= 1
209
+ return web.Response(
210
+ status=206 if range_header else 200,
211
+ headers=headers
212
+ )
213
+
214
+ async def stream_generator():
215
+ try:
216
+ bytes_sent = 0
217
+ bytes_to_skip = start % CHUNK_SIZE
218
+
219
+ async for chunk in streamer.stream_file(
220
+ media_ref,
221
+ offset=start,
222
+ limit=content_length,
223
+ fallback_message_id=fallback_message_id,
224
+ on_fallback_message=on_fallback_message
225
+ ):
226
+ if bytes_to_skip > 0:
227
+ if len(chunk) <= bytes_to_skip:
228
+ bytes_to_skip -= len(chunk)
229
+ continue
230
+ chunk = chunk[bytes_to_skip:]
231
+ bytes_to_skip = 0
232
+
233
+ remaining = content_length - bytes_sent
234
+ if len(chunk) > remaining:
235
+ chunk = chunk[:remaining]
236
+
237
+ if chunk:
238
+ yield chunk
239
+ bytes_sent += len(chunk)
240
+
241
+ if bytes_sent >= content_length:
242
+ break
243
+ finally:
244
+ work_loads[client_id] -= 1
245
+
246
+ return web.Response(
247
+ status=206 if range_header else 200,
248
+ body=stream_generator(),
249
+ headers=headers
250
+ )
251
 
252
 
253
  @routes.get("/", allow_head=True)
 
290
  })
291
 
292
 
293
+ @routes.options(r"/{path:.+}")
294
+ async def media_options(request: web.Request):
295
  return web.Response(headers={
296
  **CORS_HEADERS,
297
  "Access-Control-Max-Age": "86400"
298
+ })
299
+
300
+
301
+ @routes.get(r"/watch/f/{secure_hash}/{name:.+}", allow_head=True)
302
+ async def canonical_media_preview(request: web.Request):
303
+ try:
304
+ secure_hash = validate_public_hash(request.match_info["secure_hash"])
305
+ file_record = await get_file_by_hash(secure_hash)
306
+ if not file_record:
307
+ raise FileNotFound("Canonical file not found")
308
+
309
+ file_name = file_record.get("file_name") or f"file_{secure_hash}"
310
+ src = f"{Var.URL.rstrip('/')}/f/{secure_hash}/{quote_media_name(file_name)}"
311
+ rendered_page = await render_media_page(file_name, src, requested_action='stream')
312
+
313
+ response = web.Response(
314
+ text=rendered_page,
315
+ content_type='text/html',
316
+ headers={
317
+ "Access-Control-Allow-Origin": "*",
318
+ "Access-Control-Allow-Headers": "Range, Content-Type, *",
319
+ "X-Content-Type-Options": "nosniff",
320
+ }
321
+ )
322
+ response.enable_compression()
323
+ return response
324
+ except (InvalidHash, FileNotFound) as e:
325
+ logger.debug(f"Canonical preview error: {type(e).__name__} - {e}", exc_info=True)
326
+ raise web.HTTPNotFound(text="Resource not found") from e
327
+ except Exception as e:
328
+ error_id = secrets.token_hex(6)
329
+ logger.error(f"Canonical preview error {error_id}: {e}", exc_info=True)
330
+ raise web.HTTPInternalServerError(
331
+ text=f"Server error occurred: {error_id}") from e
332
+
333
+
334
+ @routes.get(r"/watch/{path:.+}", allow_head=True)
335
+ async def media_preview(request: web.Request):
336
  try:
337
  path = request.match_info["path"]
338
  message_id, secure_hash = parse_media_request(path, request.query)
 
360
  except Exception as e:
361
  error_id = secrets.token_hex(6)
362
  logger.error(f"Preview error {error_id}: {e}", exc_info=True)
363
+ raise web.HTTPInternalServerError(
364
+ text=f"Server error occurred: {error_id}") from e
365
+
366
+
367
+ @routes.get(r"/f/{secure_hash}/{name:.+}", allow_head=True)
368
+ async def canonical_media_delivery(request: web.Request):
369
+ try:
370
+ secure_hash = validate_public_hash(request.match_info["secure_hash"])
371
+ file_record = await get_file_by_hash(secure_hash)
372
+ if not file_record:
373
+ raise FileNotFound("Canonical file not found")
374
+
375
+ client_id, streamer = select_optimal_client()
376
+ work_loads[client_id] += 1
377
+
378
+ try:
379
+ _resolve_unique_id(file_record)
380
+ media_ref = int(file_record["canonical_message_id"])
381
+ if client_id == 0 and file_record.get("file_id"):
382
+ media_ref = file_record["file_id"]
383
+ fallback_message_id = int(file_record["canonical_message_id"])
384
+
385
+ async def persist_refreshed_file_id(message):
386
+ if client_id != 0:
387
+ return
388
+ media = get_media(message)
389
+ new_file_id = getattr(media, "file_id", None) if media else None
390
+ if new_file_id and new_file_id != file_record.get("file_id"):
391
+ await update_cached_file_id(file_record, new_file_id)
392
+
393
+ return await _serve_media_response(
394
+ request,
395
+ file_info=file_record,
396
+ streamer=streamer,
397
+ client_id=client_id,
398
+ media_ref=media_ref,
399
+ fallback_message_id=fallback_message_id,
400
+ on_fallback_message=persist_refreshed_file_id
401
+ )
402
+ except (FileNotFound, InvalidHash):
403
+ work_loads[client_id] -= 1
404
+ raise
405
+ except Exception as e:
406
+ work_loads[client_id] -= 1
407
+ error_id = secrets.token_hex(6)
408
+ logger.error(f"Canonical stream error {error_id}: {e}", exc_info=True)
409
+ raise web.HTTPInternalServerError(
410
+ text=f"Server error during streaming: {error_id}") from e
411
+ except (InvalidHash, FileNotFound) as e:
412
+ logger.debug(f"Canonical client error: {type(e).__name__} - {e}", exc_info=True)
413
+ raise web.HTTPNotFound(text="Resource not found") from e
414
+ except Exception as e:
415
+ error_id = secrets.token_hex(6)
416
+ logger.error(f"Canonical server error {error_id}: {e}", exc_info=True)
417
+ raise web.HTTPInternalServerError(
418
+ text=f"An unexpected server error occurred: {error_id}") from e
419
+
420
+
421
+ @routes.get(r"/{path:.+}", allow_head=True)
422
+ async def media_delivery(request: web.Request):
423
  try:
424
  path = request.match_info["path"]
425
  message_id, secure_hash = parse_media_request(path, request.query)
426
 
427
  client_id, streamer = select_optimal_client()
428
 
429
+ work_loads[client_id] += 1
430
+
431
+ try:
432
+ file_info = await streamer.get_file_info(message_id)
433
+ unique_id = _resolve_unique_id(file_info)
434
+
435
+ if unique_id[:SECURE_HASH_LENGTH] != secure_hash:
436
+ raise InvalidHash(
437
+ "Provided hash does not match file's unique ID.")
438
+ return await _serve_media_response(
439
+ request,
440
+ file_info=file_info,
441
+ streamer=streamer,
442
+ client_id=client_id,
443
+ media_ref=message_id
444
+ )
445
+
446
+ except (FileNotFound, InvalidHash):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
447
  work_loads[client_id] -= 1
448
  raise
449
  except Exception as e:
Thunder/utils/bot_utils.py CHANGED
@@ -1,8 +1,8 @@
1
  # Thunder/utils/bot_utils.py
2
 
3
- import asyncio
4
- from typing import Any, Dict, Optional
5
- from urllib.parse import quote
6
 
7
  from pyrogram import Client
8
  from pyrogram.enums import ChatMemberStatus
@@ -16,12 +16,64 @@ from Thunder.utils.human_readable import humanbytes
16
  from Thunder.utils.logger import logger
17
  from Thunder.utils.messages import (MSG_BUTTON_GET_HELP, MSG_DC_UNKNOWN,
18
  MSG_DC_USER_INFO, MSG_NEW_USER)
19
- from Thunder.utils.shortener import shorten
20
- from Thunder.vars import Var
21
-
22
-
23
-
24
- async def notify_ch(cli: Client, txt: str):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  if not (hasattr(Var, 'BIN_CHANNEL') and isinstance(Var.BIN_CHANNEL, int) and Var.BIN_CHANNEL != 0):
26
  return
27
  try:
@@ -78,32 +130,20 @@ async def log_newusr(cli: Client, uid: int, fname: str):
78
  logger.error(f"Database error in log_newusr for user {uid}: {e}")
79
 
80
 
81
- async def gen_links(fwd_msg: Message, shortener: bool = True) -> Dict[str, str]:
82
- base_url = Var.URL.rstrip("/")
83
- fid = fwd_msg.id
84
- m_name_raw = get_fname(fwd_msg)
85
- m_name = m_name_raw.decode('utf-8', errors='replace') if isinstance(m_name_raw, bytes) else str(m_name_raw)
86
- m_size_hr = humanbytes(get_fsize(fwd_msg))
87
- enc_fname = quote(m_name)
88
- f_hash = get_hash(fwd_msg)
89
- slink = f"{base_url}/watch/{f_hash}{fid}/{enc_fname}"
90
- olink = f"{base_url}/{f_hash}{fid}/{enc_fname}"
91
-
92
- if shortener and getattr(Var, "SHORTEN_MEDIA_LINKS", False):
93
- try:
94
- s_results = await asyncio.gather(shorten(slink), shorten(olink), return_exceptions=True)
95
- if not isinstance(s_results[0], Exception):
96
- slink = s_results[0]
97
- else:
98
- logger.warning(f"Failed to shorten stream_link: {s_results[0]}")
99
- if not isinstance(s_results[1], Exception):
100
- olink = s_results[1]
101
- else:
102
- logger.warning(f"Failed to shorten online_link: {s_results[1]}")
103
- except Exception as e:
104
- logger.error(f"Error during link shortening: {e}")
105
-
106
- return {"stream_link": slink, "online_link": olink, "media_name": m_name, "media_size": m_size_hr}
107
 
108
 
109
  async def gen_dc_txt(usr: User) -> str:
 
1
  # Thunder/utils/bot_utils.py
2
 
3
+ import asyncio
4
+ from typing import Any, Dict, Optional
5
+ from urllib.parse import quote
6
 
7
  from pyrogram import Client
8
  from pyrogram.enums import ChatMemberStatus
 
16
  from Thunder.utils.logger import logger
17
  from Thunder.utils.messages import (MSG_BUTTON_GET_HELP, MSG_DC_UNKNOWN,
18
  MSG_DC_USER_INFO, MSG_NEW_USER)
19
+ from Thunder.utils.shortener import shorten
20
+ from Thunder.vars import Var
21
+
22
+
23
+ def quote_media_name(file_name: str) -> str:
24
+ return quote(str(file_name).replace("/", "_"), safe="")
25
+
26
+
27
+ async def _build_links(
28
+ *,
29
+ download_path: str,
30
+ stream_path: str,
31
+ media_name: str,
32
+ media_size: str,
33
+ shortener: bool = True
34
+ ) -> Dict[str, str]:
35
+ base_url = Var.URL.rstrip("/")
36
+ slink = f"{base_url}{stream_path}"
37
+ olink = f"{base_url}{download_path}"
38
+
39
+ if shortener and getattr(Var, "SHORTEN_MEDIA_LINKS", False):
40
+ try:
41
+ s_results = await asyncio.gather(shorten(slink), shorten(olink), return_exceptions=True)
42
+ if not isinstance(s_results[0], Exception):
43
+ slink = s_results[0]
44
+ else:
45
+ logger.warning(f"Failed to shorten stream_link: {s_results[0]}")
46
+ if not isinstance(s_results[1], Exception):
47
+ olink = s_results[1]
48
+ else:
49
+ logger.warning(f"Failed to shorten online_link: {s_results[1]}")
50
+ except Exception as e:
51
+ logger.error(f"Error during link shortening: {e}")
52
+
53
+ return {"stream_link": slink, "online_link": olink, "media_name": media_name, "media_size": media_size}
54
+
55
+
56
+ async def gen_canonical_links(
57
+ *,
58
+ file_name: str,
59
+ file_size: int,
60
+ public_hash: str,
61
+ shortener: bool = True
62
+ ) -> Dict[str, str]:
63
+ media_name = str(file_name)
64
+ media_size = humanbytes(file_size)
65
+ encoded_name = quote_media_name(media_name)
66
+ return await _build_links(
67
+ download_path=f"/f/{public_hash}/{encoded_name}",
68
+ stream_path=f"/watch/f/{public_hash}/{encoded_name}",
69
+ media_name=media_name,
70
+ media_size=media_size,
71
+ shortener=shortener
72
+ )
73
+
74
+
75
+
76
+ async def notify_ch(cli: Client, txt: str):
77
  if not (hasattr(Var, 'BIN_CHANNEL') and isinstance(Var.BIN_CHANNEL, int) and Var.BIN_CHANNEL != 0):
78
  return
79
  try:
 
130
  logger.error(f"Database error in log_newusr for user {uid}: {e}")
131
 
132
 
133
+ async def gen_links(fwd_msg: Message, shortener: bool = True) -> Dict[str, str]:
134
+ fid = fwd_msg.id
135
+ m_name_raw = get_fname(fwd_msg)
136
+ m_name = m_name_raw.decode('utf-8', errors='replace') if isinstance(m_name_raw, bytes) else str(m_name_raw)
137
+ m_size_hr = humanbytes(get_fsize(fwd_msg))
138
+ enc_fname = quote_media_name(m_name)
139
+ f_hash = get_hash(fwd_msg)
140
+ return await _build_links(
141
+ download_path=f"/{f_hash}{fid}/{enc_fname}",
142
+ stream_path=f"/watch/{f_hash}{fid}/{enc_fname}",
143
+ media_name=m_name,
144
+ media_size=m_size_hr,
145
+ shortener=shortener
146
+ )
 
 
 
 
 
 
 
 
 
 
 
 
147
 
148
 
149
  async def gen_dc_txt(usr: User) -> str:
Thunder/utils/canonical_files.py ADDED
@@ -0,0 +1,333 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import datetime
3
+ import hashlib
4
+ from collections import OrderedDict
5
+ from contextlib import asynccontextmanager
6
+ from typing import Any, Awaitable, Callable, Dict, Optional, Tuple
7
+
8
+ from pyrogram.errors import FloodWait
9
+ from pyrogram.types import Message
10
+ from pymongo.errors import DuplicateKeyError
11
+
12
+ from Thunder.bot import StreamBot
13
+ from Thunder.utils.database import db
14
+ from Thunder.utils.file_properties import get_fname, get_media, get_uniqid
15
+ from Thunder.utils.logger import logger
16
+ from Thunder.vars import Var
17
+
18
+ PUBLIC_HASH_LENGTH = 20
19
+ _CACHE_TTL_SECONDS = 600
20
+ _CACHE_MAX_ITEMS = 4096
21
+
22
+ _cache_by_unique_id: "OrderedDict[str, Tuple[float, Dict[str, Any]]]" = OrderedDict()
23
+ _cache_by_hash: "OrderedDict[str, Tuple[float, Dict[str, Any]]]" = OrderedDict()
24
+ _cache_by_message_id: "OrderedDict[int, Tuple[float, Dict[str, Any]]]" = OrderedDict()
25
+
26
+ _upload_locks: dict[str, asyncio.Lock] = {}
27
+ _upload_lock_counts: dict[str, int] = {}
28
+ _upload_locks_guard = asyncio.Lock()
29
+
30
+
31
+ def build_public_hash(file_unique_id: str) -> str:
32
+ return hashlib.sha256(file_unique_id.encode("utf-8")).hexdigest()[:PUBLIC_HASH_LENGTH]
33
+
34
+
35
+ def _infer_mime_type(media: Any) -> str:
36
+ mime_type = getattr(media, "mime_type", None)
37
+ if mime_type:
38
+ return mime_type
39
+
40
+ mime_map = {
41
+ "photo": "image/jpeg",
42
+ "voice": "audio/ogg",
43
+ "videonote": "video/mp4",
44
+ }
45
+ return mime_map.get(type(media).__name__.lower(), "application/octet-stream")
46
+
47
+
48
+ def build_file_record(
49
+ stored_message: Message,
50
+ *,
51
+ source_chat_id: Optional[int] = None,
52
+ source_message_id: Optional[int] = None
53
+ ) -> Optional[Dict[str, Any]]:
54
+ media = get_media(stored_message)
55
+ file_unique_id = get_uniqid(stored_message)
56
+ if not media or not file_unique_id:
57
+ return None
58
+
59
+ now = datetime.datetime.utcnow()
60
+ return {
61
+ "file_unique_id": file_unique_id,
62
+ "public_hash": build_public_hash(file_unique_id),
63
+ "canonical_message_id": stored_message.id,
64
+ "file_id": getattr(media, "file_id", None),
65
+ "file_name": get_fname(stored_message),
66
+ "mime_type": _infer_mime_type(media),
67
+ "file_size": getattr(media, "file_size", 0) or 0,
68
+ "media_type": type(media).__name__.lower(),
69
+ "first_source_chat_id": source_chat_id,
70
+ "first_source_message_id": source_message_id,
71
+ "created_at": now,
72
+ "last_seen_at": now,
73
+ "seen_count": 1,
74
+ "reuse_count": 0
75
+ }
76
+
77
+
78
+ def _prune_cache(cache: "OrderedDict[Any, Tuple[float, Dict[str, Any]]]") -> None:
79
+ now = asyncio.get_running_loop().time()
80
+ expired_keys = [key for key, (ts, _) in cache.items() if now - ts > _CACHE_TTL_SECONDS]
81
+ for key in expired_keys:
82
+ cache.pop(key, None)
83
+ while len(cache) > _CACHE_MAX_ITEMS:
84
+ cache.popitem(last=False)
85
+
86
+
87
+ def _cache_get(
88
+ cache: "OrderedDict[Any, Tuple[float, Dict[str, Any]]]",
89
+ key: Any
90
+ ) -> Optional[Dict[str, Any]]:
91
+ if key not in cache:
92
+ return None
93
+ ts, value = cache[key]
94
+ now = asyncio.get_running_loop().time()
95
+ if now - ts > _CACHE_TTL_SECONDS:
96
+ cache.pop(key, None)
97
+ return None
98
+ cache.move_to_end(key)
99
+ return value
100
+
101
+
102
+ def _remember(record: Dict[str, Any]) -> Dict[str, Any]:
103
+ now = asyncio.get_running_loop().time()
104
+ file_unique_id = record.get("file_unique_id")
105
+ public_hash = record.get("public_hash")
106
+ canonical_message_id = record.get("canonical_message_id")
107
+
108
+ if file_unique_id:
109
+ _cache_by_unique_id[file_unique_id] = (now, record)
110
+ _cache_by_unique_id.move_to_end(file_unique_id)
111
+ _prune_cache(_cache_by_unique_id)
112
+ if public_hash:
113
+ _cache_by_hash[public_hash] = (now, record)
114
+ _cache_by_hash.move_to_end(public_hash)
115
+ _prune_cache(_cache_by_hash)
116
+ if canonical_message_id is not None:
117
+ _cache_by_message_id[canonical_message_id] = (now, record)
118
+ _cache_by_message_id.move_to_end(canonical_message_id)
119
+ _prune_cache(_cache_by_message_id)
120
+ return record
121
+
122
+
123
+ def _forget(record: Dict[str, Any]) -> None:
124
+ file_unique_id = record.get("file_unique_id")
125
+ public_hash = record.get("public_hash")
126
+ canonical_message_id = record.get("canonical_message_id")
127
+
128
+ if file_unique_id:
129
+ _cache_by_unique_id.pop(file_unique_id, None)
130
+ if public_hash:
131
+ _cache_by_hash.pop(public_hash, None)
132
+ if canonical_message_id is not None:
133
+ _cache_by_message_id.pop(canonical_message_id, None)
134
+
135
+
136
+ async def get_file_by_unique_id(file_unique_id: str) -> Optional[Dict[str, Any]]:
137
+ cached = _cache_get(_cache_by_unique_id, file_unique_id)
138
+ if cached:
139
+ return cached
140
+ record = await db.get_file_by_unique_id(file_unique_id)
141
+ return _remember(record) if record else None
142
+
143
+
144
+ async def get_file_by_hash(public_hash: str) -> Optional[Dict[str, Any]]:
145
+ cached = _cache_get(_cache_by_hash, public_hash)
146
+ if cached:
147
+ return cached
148
+ record = await db.get_file_by_hash(public_hash)
149
+ return _remember(record) if record else None
150
+
151
+
152
+ async def get_file_by_message_id(canonical_message_id: int) -> Optional[Dict[str, Any]]:
153
+ cached = _cache_get(_cache_by_message_id, canonical_message_id)
154
+ if cached:
155
+ return cached
156
+ record = await db.get_file_by_message_id(canonical_message_id)
157
+ return _remember(record) if record else None
158
+
159
+
160
+ async def touch_file_record(record: Dict[str, Any], *, reused: bool = False) -> None:
161
+ if not record.get("public_hash"):
162
+ return
163
+ record["last_seen_at"] = datetime.datetime.utcnow()
164
+ record["seen_count"] = int(record.get("seen_count", 0)) + 1
165
+ if reused:
166
+ record["reuse_count"] = int(record.get("reuse_count", 0)) + 1
167
+ _remember(record)
168
+ await db.touch_file_record(record["public_hash"], reused=reused)
169
+
170
+
171
+ def schedule_touch_file_record(record: Dict[str, Any], *, reused: bool = False) -> None:
172
+ if not record.get("public_hash"):
173
+ return
174
+
175
+ record["last_seen_at"] = datetime.datetime.utcnow()
176
+ record["seen_count"] = int(record.get("seen_count", 0)) + 1
177
+ if reused:
178
+ record["reuse_count"] = int(record.get("reuse_count", 0)) + 1
179
+ _remember(record)
180
+
181
+ task = asyncio.create_task(
182
+ db.touch_file_record(record["public_hash"], reused=reused),
183
+ name=f"touch_file_record:{record['public_hash']}"
184
+ )
185
+
186
+ def _log_touch_failure(done_task: asyncio.Task) -> None:
187
+ try:
188
+ done_task.result()
189
+ except Exception as e:
190
+ logger.error(
191
+ f"Background touch failed for canonical file {record['public_hash']}: {e}",
192
+ exc_info=True
193
+ )
194
+
195
+ task.add_done_callback(_log_touch_failure)
196
+
197
+
198
+ async def update_cached_file_id(record: Dict[str, Any], file_id: str) -> None:
199
+ if not record.get("public_hash") or not file_id:
200
+ return
201
+ record["file_id"] = file_id
202
+ _remember(record)
203
+ await db.update_file_id(record["public_hash"], file_id)
204
+
205
+
206
+ async def _fetch_canonical_message(record: Dict[str, Any]) -> Optional[Message]:
207
+ canonical_message_id = record.get("canonical_message_id")
208
+ if canonical_message_id is None:
209
+ return None
210
+
211
+ try:
212
+ try:
213
+ message = await StreamBot.get_messages(
214
+ chat_id=int(Var.BIN_CHANNEL),
215
+ message_ids=int(canonical_message_id)
216
+ )
217
+ except FloodWait as e:
218
+ await asyncio.sleep(e.value)
219
+ message = await StreamBot.get_messages(
220
+ chat_id=int(Var.BIN_CHANNEL),
221
+ message_ids=int(canonical_message_id)
222
+ )
223
+ except Exception as e:
224
+ logger.warning(
225
+ f"Transient error fetching canonical message {canonical_message_id}: {e}",
226
+ exc_info=True
227
+ )
228
+ raise
229
+
230
+ if not message or not message.media:
231
+ return None
232
+ return message
233
+
234
+
235
+ async def _is_canonical_record_valid(record: Dict[str, Any], file_unique_id: str) -> bool:
236
+ message = await _fetch_canonical_message(record)
237
+ return bool(message and get_uniqid(message) == file_unique_id)
238
+
239
+
240
+ def _merge_replacement_record(
241
+ existing: Dict[str, Any],
242
+ refreshed: Dict[str, Any]
243
+ ) -> Dict[str, Any]:
244
+ refreshed["created_at"] = existing.get("created_at", refreshed["created_at"])
245
+ refreshed["seen_count"] = int(existing.get("seen_count", 0)) + 1
246
+ refreshed["reuse_count"] = int(existing.get("reuse_count", 0))
247
+ refreshed["first_source_chat_id"] = existing.get(
248
+ "first_source_chat_id",
249
+ refreshed.get("first_source_chat_id")
250
+ )
251
+ refreshed["first_source_message_id"] = existing.get(
252
+ "first_source_message_id",
253
+ refreshed.get("first_source_message_id")
254
+ )
255
+ return refreshed
256
+
257
+
258
+ @asynccontextmanager
259
+ async def file_ingest_lock(file_unique_id: str):
260
+ async with _upload_locks_guard:
261
+ lock = _upload_locks.get(file_unique_id)
262
+ if lock is None:
263
+ lock = asyncio.Lock()
264
+ _upload_locks[file_unique_id] = lock
265
+ _upload_lock_counts[file_unique_id] = 0
266
+ _upload_lock_counts[file_unique_id] += 1
267
+
268
+ await lock.acquire()
269
+ try:
270
+ yield
271
+ finally:
272
+ lock.release()
273
+ async with _upload_locks_guard:
274
+ remaining = _upload_lock_counts.get(file_unique_id, 1) - 1
275
+ if remaining <= 0:
276
+ _upload_lock_counts.pop(file_unique_id, None)
277
+ _upload_locks.pop(file_unique_id, None)
278
+ else:
279
+ _upload_lock_counts[file_unique_id] = remaining
280
+
281
+
282
+ async def get_or_create_canonical_file(
283
+ source_message: Message,
284
+ copy_media: Callable[[Message], Awaitable[Optional[Message]]]
285
+ ) -> Tuple[Optional[Dict[str, Any]], Optional[Message], bool]:
286
+ file_unique_id = get_uniqid(source_message)
287
+ if not file_unique_id:
288
+ return None, None, False
289
+
290
+ async with file_ingest_lock(file_unique_id):
291
+ existing = await get_file_by_unique_id(file_unique_id)
292
+ if existing:
293
+ try:
294
+ is_valid = await _is_canonical_record_valid(existing, file_unique_id)
295
+ except Exception:
296
+ schedule_touch_file_record(existing, reused=True)
297
+ return existing, None, True
298
+ if is_valid:
299
+ schedule_touch_file_record(existing, reused=True)
300
+ return existing, None, True
301
+ _forget(existing)
302
+
303
+ stored_message = await copy_media(source_message)
304
+ if not stored_message:
305
+ return None, None, False
306
+
307
+ record = build_file_record(
308
+ stored_message,
309
+ source_chat_id=source_message.chat.id if source_message.chat else None,
310
+ source_message_id=source_message.id
311
+ )
312
+ if not record:
313
+ return None, stored_message, False
314
+
315
+ try:
316
+ if existing:
317
+ record = _merge_replacement_record(existing, record)
318
+ await db.replace_file_record(record)
319
+ else:
320
+ await db.create_file_record(record)
321
+ _remember(record)
322
+ return record, stored_message, False
323
+ except DuplicateKeyError:
324
+ existing = await get_file_by_unique_id(file_unique_id)
325
+ if existing:
326
+ schedule_touch_file_record(existing, reused=True)
327
+ return existing, None, True
328
+ raise
329
+ except FloodWait:
330
+ raise
331
+ except Exception as e:
332
+ logger.error(f"Error creating canonical file for {file_unique_id}: {e}", exc_info=True)
333
+ return None, stored_message, False
Thunder/utils/custom_dl.py CHANGED
@@ -1,7 +1,7 @@
1
  # Thunder/utils/custom_dl.py
2
 
3
  import asyncio
4
- from typing import Any, AsyncGenerator, Dict
5
 
6
  from pyrogram import Client
7
  from pyrogram.errors import FloodWait
@@ -37,26 +37,52 @@ class ByteStreamer:
37
  return message
38
 
39
  async def stream_file(
40
- self, message_id: int, offset: int = 0, limit: int = 0
 
 
 
 
 
41
  ) -> AsyncGenerator[bytes, None]:
42
- message = await self.get_message(message_id)
43
-
44
  chunk_offset = offset // (1024 * 1024)
45
-
46
  chunk_limit = 0
47
  if limit > 0:
48
  chunk_limit = ((limit + (1024 * 1024) - 1) // (1024 * 1024)) + 1
49
 
50
- while True:
51
- try:
52
- async for chunk in self.client.stream_media(
53
- message, offset=chunk_offset, limit=chunk_limit
54
- ):
55
- yield chunk
56
- break
57
- except FloodWait as e:
58
- logger.debug(f"FloodWait: stream_file, sleep {e.value}s")
59
- await asyncio.sleep(e.value)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
60
 
61
  def get_file_info_sync(self, message: Message) -> Dict[str, Any]:
62
  media = get_media(message)
 
1
  # Thunder/utils/custom_dl.py
2
 
3
  import asyncio
4
+ from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Optional
5
 
6
  from pyrogram import Client
7
  from pyrogram.errors import FloodWait
 
37
  return message
38
 
39
  async def stream_file(
40
+ self,
41
+ media_ref: int | str | Message,
42
+ offset: int = 0,
43
+ limit: int = 0,
44
+ fallback_message_id: int | None = None,
45
+ on_fallback_message: Optional[Callable[[Message], Awaitable[None]]] = None
46
  ) -> AsyncGenerator[bytes, None]:
 
 
47
  chunk_offset = offset // (1024 * 1024)
 
48
  chunk_limit = 0
49
  if limit > 0:
50
  chunk_limit = ((limit + (1024 * 1024) - 1) // (1024 * 1024)) + 1
51
 
52
+ refs: list[int | str | Message] = [media_ref]
53
+ if fallback_message_id is not None and fallback_message_id != media_ref:
54
+ refs.append(fallback_message_id)
55
+
56
+ last_error: Exception | None = None
57
+ for ref in refs:
58
+ started_stream = False
59
+ while True:
60
+ try:
61
+ target = await self.get_message(ref) if isinstance(ref, int) else ref
62
+ if (
63
+ on_fallback_message is not None and
64
+ fallback_message_id is not None and
65
+ ref == fallback_message_id and
66
+ isinstance(target, Message)
67
+ ):
68
+ await on_fallback_message(target)
69
+ async for chunk in self.client.stream_media(
70
+ target, offset=chunk_offset, limit=chunk_limit
71
+ ):
72
+ started_stream = True
73
+ yield chunk
74
+ return
75
+ except FloodWait as e:
76
+ logger.debug(f"FloodWait: stream_file, sleep {e.value}s")
77
+ await asyncio.sleep(e.value)
78
+ except Exception as e:
79
+ last_error = e
80
+ logger.debug(f"Error streaming media ref {ref}: {e}", exc_info=True)
81
+ if started_stream:
82
+ raise
83
+ break
84
+
85
+ raise FileNotFound(f"Unable to stream file: {last_error}")
86
 
87
  def get_file_info_sync(self, message: Message) -> Dict[str, Any]:
88
  media = get_media(message)
Thunder/utils/database.py CHANGED
@@ -1,7 +1,7 @@
1
  # Thunder/utils/database.py
2
 
3
  import datetime
4
- from typing import Optional, Dict, Any
5
  from pymongo import AsyncMongoClient
6
  from pymongo.asynchronous.collection import AsyncCollection
7
  from Thunder.vars import Var
@@ -17,6 +17,7 @@ class Database:
17
  self.token_col: AsyncCollection = self.db.tokens
18
  self.authorized_users_col: AsyncCollection = self.db.authorized_users
19
  self.restart_message_col: AsyncCollection = self.db.restart_message
 
20
 
21
  async def ensure_indexes(self):
22
  try:
@@ -29,6 +30,11 @@ class Database:
29
  await self.token_col.create_index("activated")
30
  await self.restart_message_col.create_index("message_id", unique=True)
31
  await self.restart_message_col.create_index("timestamp", expireAfterSeconds=3600)
 
 
 
 
 
32
 
33
  logger.debug("Database indexes ensured.")
34
  except Exception as e:
@@ -245,6 +251,80 @@ class Database:
245
  logger.error(f"Error in is_user_authorized for user {user_id}: {e}", exc_info=True)
246
  return False
247
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
248
  async def close(self):
249
  if self._client:
250
  await self._client.close()
 
1
  # Thunder/utils/database.py
2
 
3
  import datetime
4
+ from typing import Any, Dict, Optional
5
  from pymongo import AsyncMongoClient
6
  from pymongo.asynchronous.collection import AsyncCollection
7
  from Thunder.vars import Var
 
17
  self.token_col: AsyncCollection = self.db.tokens
18
  self.authorized_users_col: AsyncCollection = self.db.authorized_users
19
  self.restart_message_col: AsyncCollection = self.db.restart_message
20
+ self.files_col: AsyncCollection = self.db.files
21
 
22
  async def ensure_indexes(self):
23
  try:
 
30
  await self.token_col.create_index("activated")
31
  await self.restart_message_col.create_index("message_id", unique=True)
32
  await self.restart_message_col.create_index("timestamp", expireAfterSeconds=3600)
33
+ await self.files_col.create_index("file_unique_id", unique=True)
34
+ await self.files_col.create_index("public_hash", unique=True)
35
+ await self.files_col.create_index("canonical_message_id", unique=True)
36
+ await self.files_col.create_index("created_at")
37
+ await self.files_col.create_index("last_seen_at")
38
 
39
  logger.debug("Database indexes ensured.")
40
  except Exception as e:
 
251
  logger.error(f"Error in is_user_authorized for user {user_id}: {e}", exc_info=True)
252
  return False
253
 
254
+ async def get_file_by_unique_id(self, file_unique_id: str) -> Optional[Dict[str, Any]]:
255
+ try:
256
+ return await self.files_col.find_one({"file_unique_id": file_unique_id})
257
+ except Exception as e:
258
+ logger.error(f"Error getting file by unique_id {file_unique_id}: {e}", exc_info=True)
259
+ return None
260
+
261
+ async def get_file_by_hash(self, public_hash: str) -> Optional[Dict[str, Any]]:
262
+ try:
263
+ return await self.files_col.find_one({"public_hash": public_hash})
264
+ except Exception as e:
265
+ logger.error(f"Error getting file by hash {public_hash}: {e}", exc_info=True)
266
+ raise
267
+
268
+ async def get_file_by_message_id(self, canonical_message_id: int) -> Optional[Dict[str, Any]]:
269
+ try:
270
+ return await self.files_col.find_one({"canonical_message_id": canonical_message_id})
271
+ except Exception as e:
272
+ logger.error(
273
+ f"Error getting file by message_id {canonical_message_id}: {e}",
274
+ exc_info=True
275
+ )
276
+ return None
277
+
278
+ async def create_file_record(self, file_record: Dict[str, Any]) -> None:
279
+ try:
280
+ await self.files_col.insert_one(file_record)
281
+ except Exception as e:
282
+ logger.error(
283
+ f"Error creating canonical file record for {file_record.get('file_unique_id')}: {e}",
284
+ exc_info=True
285
+ )
286
+ raise
287
+
288
+ async def replace_file_record(self, file_record: Dict[str, Any]) -> None:
289
+ try:
290
+ await self.files_col.replace_one(
291
+ {"file_unique_id": file_record["file_unique_id"]},
292
+ file_record,
293
+ upsert=True
294
+ )
295
+ except Exception as e:
296
+ logger.error(
297
+ f"Error replacing canonical file record for {file_record.get('file_unique_id')}: {e}",
298
+ exc_info=True
299
+ )
300
+ raise
301
+
302
+ async def touch_file_record(self, public_hash: str, *, reused: bool = False) -> None:
303
+ try:
304
+ update_doc: Dict[str, Any] = {
305
+ "$set": {"last_seen_at": datetime.datetime.utcnow()},
306
+ "$inc": {"seen_count": 1}
307
+ }
308
+ if reused:
309
+ update_doc["$inc"]["reuse_count"] = 1
310
+ await self.files_col.update_one({"public_hash": public_hash}, update_doc)
311
+ except Exception as e:
312
+ logger.error(f"Error touching canonical file {public_hash}: {e}", exc_info=True)
313
+
314
+ async def update_file_id(self, public_hash: str, file_id: str) -> None:
315
+ try:
316
+ await self.files_col.update_one(
317
+ {"public_hash": public_hash},
318
+ {
319
+ "$set": {
320
+ "file_id": file_id,
321
+ "last_seen_at": datetime.datetime.utcnow()
322
+ }
323
+ }
324
+ )
325
+ except Exception as e:
326
+ logger.error(f"Error updating file_id for {public_hash}: {e}", exc_info=True)
327
+
328
  async def close(self):
329
  if self._client:
330
  await self._client.close()
Thunder/utils/render_template.py CHANGED
@@ -1,8 +1,8 @@
1
  # Thunder/utils/render_template.py
2
 
3
- import asyncio
4
- import html as html_module
5
- import urllib.parse
6
 
7
  from jinja2 import Environment, FileSystemLoader
8
  from pyrogram.errors import FloodWait
@@ -13,18 +13,36 @@ from Thunder.utils.file_properties import get_fname, get_uniqid
13
  from Thunder.utils.logger import logger
14
  from Thunder.vars import Var
15
 
16
- template_env = Environment(
17
  loader=FileSystemLoader('Thunder/template'),
18
  enable_async=True,
19
  cache_size=200,
20
  auto_reload=False,
21
- optimized=True
22
  )
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  async def render_page(id: int, secure_hash: str, requested_action: str | None = None) -> str:
25
- try:
26
- try:
27
- message = await StreamBot.get_messages(chat_id=int(Var.BIN_CHANNEL), message_ids=id)
28
  except FloodWait as e:
29
  await asyncio.sleep(e.value)
30
  message = await StreamBot.get_messages(chat_id=int(Var.BIN_CHANNEL), message_ids=id)
@@ -34,27 +52,13 @@ async def render_page(id: int, secure_hash: str, requested_action: str | None =
34
 
35
  file_unique_id = get_uniqid(message)
36
  file_name = get_fname(message)
37
-
38
  if not file_unique_id or file_unique_id[:6] != secure_hash:
39
  raise InvalidHash("File unique ID or secure hash mismatch during rendering.")
40
 
41
- quoted_filename = urllib.parse.quote(file_name.replace('/', '_'))
42
  src = urllib.parse.urljoin(Var.URL, f'{secure_hash}{id}/{quoted_filename}')
43
- safe_filename = html_module.escape(file_name)
44
- if requested_action == 'stream':
45
- template = template_env.get_template('req.html')
46
- context = {
47
- 'heading': f"View {safe_filename}",
48
- 'file_name': safe_filename,
49
- 'src': f"{src}?disposition=inline"
50
- }
51
- else:
52
- template = template_env.get_template('dl.html')
53
- context = {
54
- 'file_name': safe_filename,
55
- 'src': src
56
- }
57
- return await template.render_async(**context)
58
- except Exception as e:
59
- logger.error(f"Error in render_page for ID {id} and hash {secure_hash}: {e}", exc_info=True)
60
- raise
 
1
  # Thunder/utils/render_template.py
2
 
3
+ import asyncio
4
+ import html as html_module
5
+ import urllib.parse
6
 
7
  from jinja2 import Environment, FileSystemLoader
8
  from pyrogram.errors import FloodWait
 
13
  from Thunder.utils.logger import logger
14
  from Thunder.vars import Var
15
 
16
+ template_env = Environment(
17
  loader=FileSystemLoader('Thunder/template'),
18
  enable_async=True,
19
  cache_size=200,
20
  auto_reload=False,
21
+ optimized=True
22
  )
23
 
24
+ async def render_media_page(file_name: str, src: str, requested_action: str | None = None) -> str:
25
+ safe_filename = html_module.escape(file_name)
26
+ if requested_action == 'stream':
27
+ template = template_env.get_template('req.html')
28
+ context = {
29
+ 'heading': f"View {safe_filename}",
30
+ 'file_name': safe_filename,
31
+ 'src': f"{src}?disposition=inline"
32
+ }
33
+ else:
34
+ template = template_env.get_template('dl.html')
35
+ context = {
36
+ 'file_name': safe_filename,
37
+ 'src': src
38
+ }
39
+ return await template.render_async(**context)
40
+
41
+
42
  async def render_page(id: int, secure_hash: str, requested_action: str | None = None) -> str:
43
+ try:
44
+ try:
45
+ message = await StreamBot.get_messages(chat_id=int(Var.BIN_CHANNEL), message_ids=id)
46
  except FloodWait as e:
47
  await asyncio.sleep(e.value)
48
  message = await StreamBot.get_messages(chat_id=int(Var.BIN_CHANNEL), message_ids=id)
 
52
 
53
  file_unique_id = get_uniqid(message)
54
  file_name = get_fname(message)
55
+
56
  if not file_unique_id or file_unique_id[:6] != secure_hash:
57
  raise InvalidHash("File unique ID or secure hash mismatch during rendering.")
58
 
59
+ quoted_filename = urllib.parse.quote(file_name.replace('/', '_'), safe="")
60
  src = urllib.parse.urljoin(Var.URL, f'{secure_hash}{id}/{quoted_filename}')
61
+ return await render_media_page(file_name, src, requested_action)
62
+ except Exception as e:
63
+ logger.error(f"Error in render_page for ID {id} and hash {secure_hash}: {e}", exc_info=True)
64
+ raise