Files
telegram-backup-downloader/downloadv4.py

633 lines
21 KiB
Python

import asyncio
import json
import logging
import os
import re
import signal
import sys
import shutil
from datetime import datetime, timezone, timedelta
from logging.handlers import RotatingFileHandler
from pathlib import Path
import aiofiles
import aiosqlite
from telethon import TelegramClient, events, utils
from telethon.errors import (
FloodWaitError,
FileReferenceExpiredError,
RPCError,
)
# =========================
# LOAD CONFIG & CONSTANTS
# =========================
with open("config.json", "r", encoding="utf-8") as f:
CONFIG = json.load(f)
API_ID = CONFIG["api_id"]
API_HASH = CONFIG["api_hash"]
PHONE = CONFIG["phone_number"]
SESSION_NAME = CONFIG["session_name"]
CHANNELS = CONFIG["channels"]
PARALLEL_DOWNLOADS = CONFIG.get("parallel_downloads", 3)
DOWNLOAD_TIMEOUT = CONFIG.get("download_timeout_seconds", 600)
DOWNLOAD_RETRY = CONFIG.get("download_retry_count", 3)
QUEUE_MAX_SIZE = CONFIG.get("queue_max_size", 5000)
MIN_DISK_SPACE = 6 * 1024 * 1024 * 1024 # 6GB in bytes
BASE_DIR = Path("channels")
LOG_DIR = Path("logs")
BASE_DIR.mkdir(exist_ok=True)
LOG_DIR.mkdir(exist_ok=True)
# =========================
# LOGGING
# =========================
logger = logging.getLogger("telegram_archiver")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
"%(asctime)s [%(levelname)s] %(message)s"
)
file_handler = RotatingFileHandler(
LOG_DIR / "app.log",
maxBytes=10 * 1024 * 1024,
backupCount=10,
encoding="utf-8"
)
file_handler.setFormatter(formatter)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.addHandler(console_handler)
# =========================
# SQLITE
# =========================
DB_FILE = "telegram_state.db"
CREATE_TABLES_SQL = """
CREATE TABLE IF NOT EXISTS channel_state (
channel_id TEXT PRIMARY KEY,
last_message_id INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS downloaded_media (
channel_id TEXT,
message_id INTEGER,
file_name TEXT,
file_size INTEGER,
PRIMARY KEY(channel_id, message_id)
);
CREATE TABLE IF NOT EXISTS logged_messages (
channel_id TEXT,
message_id INTEGER,
PRIMARY KEY(channel_id, message_id)
);
CREATE TABLE IF NOT EXISTS failed_downloads (
channel_id TEXT,
message_id INTEGER,
last_error TEXT,
retry_count INTEGER,
PRIMARY KEY(channel_id,message_id)
);
"""
db = None
# =========================
# GLOBALS & SAFEGUARDS
# =========================
INVALID_CHARS = r'[<>:"/\\\\|?*]'
client = TelegramClient(
SESSION_NAME,
API_ID,
API_HASH,
auto_reconnect=True,
connection_retries=None,
retry_delay=5,
)
download_queue = asyncio.Queue(maxsize=QUEUE_MAX_SIZE)
priority_download_queue = asyncio.Queue()
shutdown_event = asyncio.Event()
shutdown_lock = asyncio.Lock()
CONFIG_LOCK = asyncio.Lock()
# Will be initialized inside main() to tie it to the correct event loop
download_allowed = None
CONFIG_FILE = "config.json"
MANUAL_IN_PROGRESS = set()
active_tasks = set()
worker_tasks = []
config_task = None
sync_task = None
disk_task = None
MONITORED_CHANNEL_IDS = set()
CHANNEL_ENTITY_MAP = {}
CHANNEL_OBJECT_MAP = {}
MY_TZ = timezone(timedelta(hours=8))
# =========================
# HELPERS
# =========================
def sanitize_filename(name):
name = re.sub(INVALID_CHARS, "_", name)
name = name.strip().rstrip(".")
return name[:200]
def parse_channel_id(raw_id):
if isinstance(raw_id, str) and raw_id.lstrip('-').isdigit():
return int(raw_id)
return raw_id
# =========================
# DATABASE OPERATIONS
# =========================
async def init_db():
global db
db = await aiosqlite.connect(DB_FILE)
await db.execute("PRAGMA journal_mode=WAL;")
await db.execute("PRAGMA synchronous=NORMAL;")
await db.executescript(CREATE_TABLES_SQL)
await db.commit()
async def add_logged_message(channel_id, message_id):
await db.execute("INSERT OR IGNORE INTO logged_messages VALUES (?,?)", (channel_id, message_id))
await db.commit()
async def resolve_channels():
MONITORED_CHANNEL_IDS.clear()
CHANNEL_ENTITY_MAP.clear()
CHANNEL_OBJECT_MAP.clear()
for channel in CHANNELS:
try:
parsed_channel = parse_channel_id(channel)
entity = await client.get_entity(parsed_channel)
entity_id = str(utils.get_peer_id(entity))
entity_title = getattr(entity, "title", entity_id)
MONITORED_CHANNEL_IDS.add(entity_id)
CHANNEL_ENTITY_MAP[entity_id] = entity_title
CHANNEL_OBJECT_MAP[entity_id] = entity
except Exception as e:
logger.error(f"[RESOLVE ERROR] Failed to resolve channel '{channel}': {e}")
async def get_last_message_id(channel_id):
async with db.execute("SELECT last_message_id FROM channel_state WHERE channel_id=?", (channel_id,)) as cursor:
row = await cursor.fetchone()
return row[0] if row else 0
async def update_last_message_id(channel_id, message_id):
await db.execute(
"""
INSERT INTO channel_state(channel_id, last_message_id)
VALUES (?, ?) ON CONFLICT(channel_id)
DO UPDATE SET last_message_id = MAX(last_message_id, excluded.last_message_id)
""", (channel_id, message_id)
)
await db.commit()
async def media_exists(channel_id, message_id):
async with db.execute("SELECT 1 FROM downloaded_media WHERE channel_id=? AND message_id=?", (channel_id, message_id)) as cursor:
return await cursor.fetchone() is not None
async def add_media_record(channel_id, message_id, file_name, file_size):
await db.execute(
"INSERT OR IGNORE INTO downloaded_media (channel_id, message_id, file_name, file_size) VALUES (?, ?, ?, ?)",
(channel_id, message_id, file_name, file_size)
)
await db.commit()
async def add_failed_download(channel_id, message_id, error):
await db.execute(
"""
INSERT INTO failed_downloads (channel_id, message_id, last_error, retry_count)
VALUES (?, ?, ?, 1) ON CONFLICT(channel_id,message_id)
DO UPDATE SET retry_count = retry_count + 1, last_error = excluded.last_error
""", (channel_id, message_id, error[:1000])
)
await db.commit()
async def text_logged(channel_id, message_id):
async with db.execute("SELECT 1 FROM logged_messages WHERE channel_id=? AND message_id=?", (channel_id, message_id)) as cursor:
return await cursor.fetchone() is not None
async def save_config():
async with CONFIG_LOCK:
async with aiofiles.open(CONFIG_FILE, "w", encoding="utf-8") as f:
await f.write(json.dumps(CONFIG, indent=2, ensure_ascii=False))
def get_channel_paths(channel_name):
safe_name = sanitize_filename(channel_name)
root = BASE_DIR / safe_name
folders = {
"root": root,
"photo": root / "photo",
"video": root / "video",
"document": root / "document",
"messages": root / "messages.txt",
"metadata": root / "metadata.jsonl",
}
for p in folders.values():
if isinstance(p, Path) and p.suffix == "":
p.mkdir(parents=True, exist_ok=True)
return folders
async def append_message_log(file_path, text):
async with aiofiles.open(file_path, "a", encoding="utf-8") as f:
await f.write(text)
async def append_metadata(file_path, data):
async with aiofiles.open(file_path, "a", encoding="utf-8") as f:
await f.write(json.dumps(data, ensure_ascii=False) + "\n")
def detect_media_type(message):
if message.photo: return "photo"
if message.video: return "video"
if message.document: return "document"
return None
# =========================
# CORE PROCESSING
# =========================
async def download_media(message, paths, channel_id):
if message.out: return
media_type = detect_media_type(message)
if not media_type: return
if await media_exists(channel_id, message.id): return
# === DISK SPACE SAFEGUARD ===
if not download_allowed.is_set():
logger.info(f"⏳ [PAUSED] Waiting for disk space to process media for message ID: {message.id}")
await download_allowed.wait()
if shutdown_event.is_set(): return
# ============================
entity = CHANNEL_OBJECT_MAP.get(channel_id)
if not entity: return
file_name = message.file.name if message.file else None
if not file_name:
ext = message.file.ext if message.file else ""
file_name = f"{message.id}{ext}"
file_name = sanitize_filename(file_name)
final_name = f"{message.id}_{file_name}"
target_dir = paths[media_type]
temp_path = target_dir / f"{final_name}.part"
target_path = target_dir / final_name
if target_path.exists():
size = target_path.stat().st_size
await add_media_record(channel_id, message.id, final_name, size)
return
logger.info(f"[DOWNLOAD] {final_name}")
success = False
last_error = "Unknown"
for attempt in range(1, DOWNLOAD_RETRY + 1):
if shutdown_event.is_set() or not client.is_connected():
break
try:
temp_path.unlink(missing_ok=True)
fresh_message = await client.get_messages(entity=entity, ids=message.id)
if not fresh_message:
raise Exception(f"Message {message.id} not found")
await asyncio.wait_for(fresh_message.download_media(file=temp_path), timeout=DOWNLOAD_TIMEOUT)
if not temp_path.exists():
raise Exception("Temp file missing")
size = temp_path.stat().st_size
if size <= 0:
raise Exception("Downloaded file is 0 bytes")
os.replace(temp_path, target_path)
await add_media_record(channel_id, message.id, final_name, size)
logger.info(f"[DOWNLOAD OK] {final_name} ({size} bytes)")
success = True
break
except Exception as e:
last_error = str(e)
logger.exception(f"[DOWNLOAD ERROR] {final_name}: {e}")
finally:
if temp_path.exists():
try:
if temp_path.stat().st_size == 0:
temp_path.unlink(missing_ok=True)
except Exception:
pass
await asyncio.sleep(min(5 * attempt, 5))
if not success and not shutdown_event.is_set():
await add_failed_download(channel_id, message.id, last_error)
logger.error(f"[FAILED] {final_name} after {DOWNLOAD_RETRY} attempts")
elif success:
await db.execute("DELETE FROM failed_downloads WHERE channel_id=? AND message_id=?", (channel_id, message.id))
await db.commit()
async def process_message(message, entity_title, channel_id):
paths = get_channel_paths(entity_title)
sender = await message.get_sender()
sender_name = "Unknown"
username = ""
if sender:
first = getattr(sender, "first_name", "") or ""
last = getattr(sender, "last_name", "") or ""
sender_name = f"{first} {last}".strip() or "Unknown"
username = getattr(sender, "username", "") or ""
timestamp = message.date.astimezone(MY_TZ).strftime("%Y-%m-%d %I:%M:%S %p")
if message.message:
if not await text_logged(channel_id, message.id):
media_info = ""
media_name = None
media_size = None
if message.media and message.file:
media_name = sanitize_filename(message.file.name or f"{message.id}{message.file.ext or ''}")
media_size = getattr(message.file, "size", 0) or 0
media_size_mb = media_size / (1024 * 1024)
media_info = f"\n\n[Media]\nFilename: {media_name}\nSize: {media_size_mb:.3f} MB"
text_log = (
f"[{timestamp}]\nSender Name: {sender_name}\nUsername: {username}\n"
f"Message ID: {message.id}\n{media_info}\n\n{message.message}\n{'-'*50}\n"
)
await append_message_log(paths["messages"], text_log)
await append_metadata(paths["metadata"], {
"timestamp": timestamp, "message_id": message.id, "sender_name": sender_name,
"username": username, "text": message.message, "media_filename": media_name, "media_size_bytes": media_size
})
await add_logged_message(channel_id, message.id)
if message.media:
await download_media(message, paths, channel_id)
await update_last_message_id(channel_id, message.id)
# =========================
# BACKGROUND MONITORS
# =========================
async def disk_monitor():
"""Continuously monitors disk space and pauses downloads if below 6GB."""
loop = asyncio.get_running_loop()
while not shutdown_event.is_set():
try:
free_space = shutil.disk_usage(BASE_DIR).free
if free_space < MIN_DISK_SPACE:
download_allowed.clear()
logger.warning(f"🚨 DISK SPACE CRITICAL: Only {free_space / 1024**3:.2f} GB left! (< 6GB)")
print("\n" + "="*60)
print("🛑 DOWNLOADS PAUSED DUE TO LOW DISK SPACE.")
print(f"Current Free Space: {free_space / 1024**3:.2f} GB")
print("Please free up storage space on the drive.")
print("="*60 + "\n")
# Asynchronously wait for terminal input without blocking the main event loop
await loop.run_in_executor(None, input, "👉 Press [ENTER] here once space is cleared to re-check...\n")
logger.info("Re-checking disk space...")
else:
if not download_allowed.is_set():
logger.info("✅ Disk space sufficient. Resuming operations.")
download_allowed.set()
await asyncio.sleep(5)
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"[DISK MONITOR] Error: {e}")
await asyncio.sleep(5)
# =========================
# WORKERS & QUEUES
# =========================
async def worker(worker_id):
logger.info(f"[WORKER {worker_id}] Started")
while True:
if shutdown_event.is_set() and download_queue.empty() and priority_download_queue.empty():
break
task = None
queue_used = None
try:
try:
item = priority_download_queue.get_nowait()
queue_used = priority_download_queue
except asyncio.QueueEmpty:
item = await asyncio.wait_for(download_queue.get(), timeout=1)
queue_used = download_queue
message, entity_title, channel_id, is_manual, raw_channel_id = item
task = asyncio.create_task(process_message(message, entity_title, channel_id))
active_tasks.add(task)
await task
if is_manual:
await remove_manual_download(raw_channel_id, message.id)
MANUAL_IN_PROGRESS.discard(f"{channel_id}:{message.id}")
except asyncio.TimeoutError:
continue
except asyncio.CancelledError:
break
except Exception:
logger.exception(f"[WORKER {worker_id}]")
finally:
if task:
active_tasks.discard(task)
if queue_used:
queue_used.task_done()
logger.info(f"[WORKER {worker_id}] Exited")
# (Omitted config watchers and specific round-robin code below for brevity, use your existing functions)
# ... [Insert reload_config, config_watcher, apply_channel_overrides, remove_manual_download, process_manual_downloads, initial_sync, new_message_handler as written previously] ...
# NOTE: Make sure to include all of them when running. I am keeping the response focused on the modifications and the main block below!
# =========================
# INITIAL SYNC (ROUND-ROBIN)
# =========================
async def initial_sync():
try:
active_syncs = []
for entity_id, entity in CHANNEL_OBJECT_MAP.items():
if shutdown_event.is_set() or not client.is_connected():
return
try:
entity_title = CHANNEL_ENTITY_MAP[entity_id]
last_message_id = await get_last_message_id(entity_id)
min_id = max(0, last_message_id - 5)
logger.info(f"[SYNC INIT] {entity_title} (last_id={last_message_id}, sync min_id={min_id})")
iterator = client.iter_messages(entity, min_id=min_id, reverse=True).__aiter__()
active_syncs.append({"entity_id": entity_id, "entity_title": entity_title, "iterator": iterator, "count": 0})
except Exception:
logger.exception(f"[SYNC INIT ERROR] Failed to initialize iterator for {entity_id}")
chunk_size = 15
while active_syncs and not shutdown_event.is_set():
if not client.is_connected():
break
for sync in list(active_syncs):
if shutdown_event.is_set() or not client.is_connected(): break
exhausted = False
for _ in range(chunk_size):
try:
message = await sync["iterator"].__anext__()
await download_queue.put((message, sync["entity_title"], sync["entity_id"], False, None))
sync["count"] += 1
except StopAsyncIteration:
exhausted = True
break
except Exception as e:
logger.error(f"[SYNC ITER ERROR] Channel {sync['entity_title']}: {e}")
exhausted = True
break
if exhausted:
logger.info(f"[SYNC DONE] {sync['entity_title']} complete. Queued: {sync['count']}")
active_syncs.remove(sync)
await asyncio.sleep(0.01)
except asyncio.CancelledError:
pass
except Exception:
logger.exception("[SYNC CRITICAL ERROR]")
# =========================
# SHUTDOWN
# =========================
async def shutdown():
async with shutdown_lock:
if shutdown_event.is_set(): return
shutdown_event.set()
logger.info("[SHUTDOWN] Graceful shutdown started")
global config_task, sync_task, disk_task
if config_task and not config_task.done(): config_task.cancel()
if sync_task and not sync_task.done(): sync_task.cancel()
if disk_task and not disk_task.done(): disk_task.cancel()
try:
if client.is_connected(): await asyncio.wait_for(client.disconnect(), timeout=10)
except Exception:
logger.exception("[DISCONNECT ERROR]")
while not download_queue.empty():
try: download_queue.get_nowait(); download_queue.task_done()
except asyncio.QueueEmpty: break
while not priority_download_queue.empty():
try: priority_download_queue.get_nowait(); priority_download_queue.task_done()
except asyncio.QueueEmpty: break
if active_tasks: await asyncio.gather(*active_tasks, return_exceptions=True)
for task in worker_tasks: task.cancel()
await asyncio.gather(*worker_tasks, return_exceptions=True)
if db:
await db.commit()
await db.close()
logger.info("[SHUTDOWN COMPLETE]")
# =========================
# MAIN
# =========================
async def main():
global config_task, sync_task, disk_task, download_allowed
# Initialize the global event tracker inside the running loop
download_allowed = asyncio.Event()
download_allowed.set()
await init_db()
await client.start(phone=PHONE)
me = await client.get_me()
logger.info(f"[CONNECTED] {me.first_name}")
await resolve_channels()
for i in range(PARALLEL_DOWNLOADS):
worker_tasks.append(asyncio.create_task(worker(i + 1)))
logger.info("[LISTENER ACTIVE]")
# Track background tasks
disk_task = asyncio.create_task(disk_monitor())
sync_task = asyncio.create_task(initial_sync())
listener_task = asyncio.create_task(client.run_until_disconnected())
await listener_task
while not shutdown_event.is_set():
await asyncio.sleep(5)
# =========================
# SIGNAL & ENTRY
# =========================
def handle_exit(*args):
loop.call_soon_threadsafe(lambda: asyncio.create_task(shutdown()))
if __name__ == "__main__":
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
signal.signal(signal.SIGINT, handle_exit)
if hasattr(signal, "SIGTERM"):
signal.signal(signal.SIGTERM, handle_exit)
try:
loop.run_until_complete(main())
finally:
loop.run_until_complete(shutdown())