dfd

Provider/telegram.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import re
+import shutil
 import sys
 import time
 from pathlib import Path
@@ -79,6 +80,215 @@ class Telegram(Provider):
 		telegram_conf = self.config.get("provider", {}).get("telegram", {}) if isinstance(self.config, dict) else {}
 		self._app_id = telegram_conf.get("app_id")
 		self._api_hash = telegram_conf.get("api_hash")
+		self._bot_token = telegram_conf.get("bot_token")
+		# Telethon downloads are chunked; larger parts mean fewer round-trips.
+		# Telethon typically expects 4..1024 KB and divisible by 4.
+		self._part_size_kb = telegram_conf.get("part_size_kb")
+		if self._part_size_kb is None:
+			self._part_size_kb = telegram_conf.get("chunk_kb")
+		if self._part_size_kb is None:
+			self._part_size_kb = telegram_conf.get("download_part_kb")
+
+	# Avoid repeatedly prompting during startup where validate() may be called multiple times.
+	_startup_auth_attempted: bool = False
+
+	def _legacy_session_base_path(self) -> Path:
+		# Older versions stored sessions under Log/medeia_macina.
+		root = Path(__file__).resolve().parents[1]
+		return root / "Log" / "medeia_macina" / "telegram"
+
+	def _migrate_legacy_session_if_needed(self) -> None:
+		"""If a legacy Telethon session exists, copy it to the new root location."""
+		try:
+			new_base = self._session_base_path()
+			new_session = Path(str(new_base) + ".session")
+			if new_session.is_file():
+				return
+
+			legacy_base = self._legacy_session_base_path()
+			legacy_session = Path(str(legacy_base) + ".session")
+			if not legacy_session.is_file():
+				return
+
+			for suffix in (".session", ".session-journal", ".session-wal", ".session-shm"):
+				src = Path(str(legacy_base) + suffix)
+				dst = Path(str(new_base) + suffix)
+				try:
+					if src.is_file() and not dst.exists():
+						shutil.copy2(str(src), str(dst))
+				except Exception:
+					continue
+		except Exception:
+			return
+
+	def _session_file_path(self) -> Path:
+		self._migrate_legacy_session_if_needed()
+		base = self._session_base_path()
+		return Path(str(base) + ".session")
+
+	def _has_session(self) -> bool:
+		self._migrate_legacy_session_if_needed()
+		try:
+			return self._session_file_path().is_file()
+		except Exception:
+			return False
+
+	def _ensure_session_interactive(self) -> bool:
+		"""Best-effort interactive auth to create a Telethon session file.
+
+		Returns True if a session exists after the attempt.
+		"""
+		if self._has_session():
+			return True
+
+		# Never prompt in non-interactive contexts.
+		try:
+			if not bool(getattr(sys.stdin, "isatty", lambda: False)()):
+				return False
+		except Exception:
+			return False
+
+		try:
+			from telethon.sync import TelegramClient
+		except Exception:
+			return False
+
+		try:
+			app_id, api_hash = self._credentials()
+		except Exception:
+			return False
+
+		self._ensure_event_loop()
+		loop = asyncio.get_event_loop()
+		if getattr(loop, "is_running", lambda: False)():
+			# Can't safely prompt/auth while another loop is running.
+			return False
+
+		def _resolve(value):
+			if asyncio.iscoroutine(value):
+				return loop.run_until_complete(value)
+			return value
+
+		try:
+			sys.stderr.write("[telegram] No session found; login required.\n")
+			sys.stderr.write("[telegram] Choose login method: 1) phone  2) bot token\n")
+			sys.stderr.write("[telegram] Enter 1 or 2: ")
+			sys.stderr.flush()
+			choice = ""
+			try:
+				choice = str(input()).strip().lower()
+			except EOFError:
+				choice = ""
+
+			use_bot = choice in {"2", "b", "bot", "token"}
+			bot_token = ""
+			if use_bot:
+				sys.stderr.write("[telegram] Bot token: ")
+				sys.stderr.flush()
+				try:
+					bot_token = str(input()).strip()
+				except EOFError:
+					bot_token = ""
+				if not bot_token:
+					return False
+				self._bot_token = bot_token
+			else:
+				sys.stderr.write("[telegram] Phone login selected (Telethon will prompt for phone + code).\n")
+				sys.stderr.flush()
+
+			session_base = self._session_base_path()
+			client = TelegramClient(str(session_base), app_id, api_hash)
+			try:
+				if use_bot:
+					_resolve(client.start(bot_token=bot_token))
+				else:
+					_resolve(client.start())
+			finally:
+				try:
+					_resolve(client.disconnect())
+				except Exception:
+					pass
+		finally:
+			try:
+				sys.stderr.write("\n")
+				sys.stderr.flush()
+			except Exception:
+				pass
+
+		return self._has_session()
+
+	def _ensure_session_with_bot_token(self, bot_token: str) -> bool:
+		"""Create a Telethon session using a bot token without prompting.
+
+		Returns True if a session exists after the attempt.
+		"""
+		if self._has_session():
+			return True
+		bot_token = str(bot_token or "").strip()
+		if not bot_token:
+			return False
+		try:
+			from telethon.sync import TelegramClient
+		except Exception:
+			return False
+		try:
+			app_id, api_hash = self._credentials()
+		except Exception:
+			return False
+
+		self._ensure_event_loop()
+		loop = asyncio.get_event_loop()
+		if getattr(loop, "is_running", lambda: False)():
+			return False
+
+		def _resolve(value):
+			if asyncio.iscoroutine(value):
+				return loop.run_until_complete(value)
+			return value
+
+		session_base = self._session_base_path()
+		client = TelegramClient(str(session_base), app_id, api_hash)
+		try:
+			_resolve(client.start(bot_token=bot_token))
+		finally:
+			try:
+				_resolve(client.disconnect())
+			except Exception:
+				pass
+
+		return self._has_session()
+
+	def _resolve_part_size_kb(self, file_size: Optional[int]) -> int:
+		# Default: bias to max throughput.
+		val = self._part_size_kb
+		try:
+			if val not in (None, ""):
+				ps = int(str(val).strip())
+			else:
+				ps = 1024
+		except Exception:
+			ps = 1024
+
+		# Clamp to Telethon-safe range.
+		if ps < 4:
+			ps = 4
+		if ps > 1024:
+			ps = 1024
+		# Must be divisible by 4.
+		ps = int(ps / 4) * 4
+		if ps <= 0:
+			ps = 64
+
+		# For very small files, reduce overhead a bit (still divisible by 4).
+		try:
+			if file_size is not None and int(file_size) > 0:
+				if int(file_size) < 2 * 1024 * 1024:
+					ps = min(ps, 256)
+				elif int(file_size) < 10 * 1024 * 1024:
+					ps = min(ps, 512)
+		except Exception:
+			pass
+		return ps
 
 	def validate(self) -> bool:
 		try:
@@ -91,16 +301,29 @@ class Telegram(Provider):
 		except Exception:
 			app_id = None
 		api_hash = str(self._api_hash).strip() if self._api_hash not in (None, "") else ""
-		return bool(app_id and api_hash)
+		if not bool(app_id and api_hash):
+			return False
+
+		# Consider Telegram "configured" only if a persisted session exists.
+		if self._has_session():
+			return True
+
+		# If a bot token is already configured, attempt a non-interactive login.
+		bot_token = str(self._bot_token or "").strip()
+		if bot_token:
+			return bool(self._ensure_session_with_bot_token(bot_token))
+
+		# Interactive startup prompt (only once per process).
+		if Telegram._startup_auth_attempted:
+			return False
+		Telegram._startup_auth_attempted = True
+		return bool(self._ensure_session_interactive())
 
 	def _session_base_path(self) -> Path:
+		# Store session alongside cookies.txt at repo root.
+		# Telethon uses this as base name and writes "<base>.session".
 		root = Path(__file__).resolve().parents[1]
-		session_dir = root / "Log" / "medeia_macina"
-		try:
-			session_dir.mkdir(parents=True, exist_ok=True)
-		except Exception:
-			pass
-		return session_dir / "telegram"
+		return root / "telegram"
 
 	def _credentials(self) -> Tuple[int, str]:
 		raw_app_id = self._app_id
@@ -266,7 +489,19 @@ class Telegram(Provider):
 				except Exception:
 					return
 
-			downloaded = _resolve(client.download_media(message, file=str(output_dir), progress_callback=_progress))
+			part_kb = self._resolve_part_size_kb(file_size)
+			try:
+				downloaded = _resolve(
+					client.download_media(
+						message,
+						file=str(output_dir),
+						progress_callback=_progress,
+						part_size_kb=part_kb,
+					)
+				)
+			except TypeError:
+				# Older/newer Telethon versions may not accept part_size_kb on download_media.
+				downloaded = _resolve(client.download_media(message, file=str(output_dir), progress_callback=_progress))
 			try:
 				sys.stderr.write("\n")
 				sys.stderr.flush()