Test / bot.py
Reooiq's picture
Update bot.py
d837e7a verified
Raw History Blame Contribute Delete
24.8 kB
import os
import asyncio
import logging
import time
import threading
import shutil
from typing import Optional, List, Dict, Any
import aiosqlite
import gradio as gr
from hydrogram import Client, filters
from hydrogram.errors import FloodWait, RPCError
from hydrogram.types import (
InputMediaPhoto, InputMediaVideo, InputMediaDocument,
InputMediaAudio, Message
)
from fastapi import FastAPI
import uvicorn
try:
from hydrogram.errors import FloodPremiumWait
except ImportError:
class FloodPremiumWait(FloodWait):
pass
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s | %(levelname)-7s | %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("UniversalArchiver")
API_ID = int(os.environ.get("API_ID", 0))
API_HASH = os.environ.get("API_HASH", "")
SESSION_STRING = os.environ.get("SESSION_STRING", "")
TARGET_CHAT_RAW = os.environ.get("TARGET_CHAT", "me")
AUTO_SCAN_ALL = os.environ.get("AUTO_SCAN_ALL", "1") == "1"
DATA_DIR = "/tmp/data"
DOWNLOAD_DIR = os.path.join(DATA_DIR, "downloads")
DB_PATH = os.path.join(DATA_DIR, "archiver_v7.db")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)
app = Client(
"universal_archiver",
api_id=API_ID,
api_hash=API_HASH,
session_string=SESSION_STRING,
in_memory=True,
workers=8,
max_concurrent_transmissions=4,
device_model="Archiver Pro",
system_version="Linux 64bit",
app_version="7.0.0",
lang_code="ar"
)
TARGET_CHAT = "me"
TARGET_CHAT_ID: Optional[int] = None
MAIN_LOOP: Optional[asyncio.AbstractEventLoop] = None
KNOWN_CHATS: Dict[int, Dict[str, Any]] = {}
stats = {
"total_queued": 0,
"total_transferred": 0,
"total_failed": 0,
"discovered_chats": 0,
"active_scans": 0,
"status": "Starting..."
}
class SmartFloodManager:
def __init__(self, normal_delay: float = 1.5):
self.normal_delay = normal_delay
self.last_call = 0.0
self.lock = asyncio.Lock()
self.flood_until = 0.0
async def wait(self):
async with self.lock:
now = time.time()
if self.flood_until > now:
sleep_time = self.flood_until - now
logger.warning(f"FloodWait active, sleeping {sleep_time:.1f}s...")
await asyncio.sleep(sleep_time)
now = time.time()
elapsed = now - self.last_call
if elapsed < self.normal_delay:
await asyncio.sleep(self.normal_delay - elapsed)
self.last_call = time.time()
def handle_flood(self, seconds: int):
wait_seconds = max(seconds, 5) + 2
self.flood_until = time.time() + wait_seconds
logger.warning(f"FloodWait received: {seconds}s. Cooling down for {wait_seconds}s.")
flood_manager = SmartFloodManager(normal_delay=1.5)
class SafeDBManager:
def __init__(self, db_path: str):
self.db_path = db_path
self.lock = asyncio.Lock()
self.conn: Optional[aiosqlite.Connection] = None
async def connect(self):
async with self.lock:
if self.conn:
try:
await self.conn.close()
except Exception:
pass
self.conn = None
if os.path.exists(self.db_path):
try:
test_conn = await aiosqlite.connect(self.db_path)
cur = await test_conn.execute("PRAGMA integrity_check;")
res = await cur.fetchone()
await test_conn.close()
if not res or res[0] != "ok":
raise Exception("Corrupt DB check failed")
except Exception as e:
logger.error(f"Resetting corrupt DB: {e}")
for ext in ["", "-wal", "-shm", "-journal"]:
f = f"{self.db_path}{ext}"
if os.path.exists(f):
try:
os.remove(f)
except Exception:
pass
self.conn = await aiosqlite.connect(self.db_path)
self.conn.row_factory = aiosqlite.Row
await self.conn.execute("PRAGMA journal_mode=WAL;")
await self.conn.execute("PRAGMA synchronous=NORMAL;")
await self.conn.execute("PRAGMA busy_timeout=60000;")
await self.conn.execute("""
CREATE TABLE IF NOT EXISTS jobs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id INTEGER NOT NULL,
message_id INTEGER NOT NULL,
media_group_id TEXT,
media_type TEXT,
caption TEXT,
source_name TEXT,
source_username TEXT,
album_total INTEGER DEFAULT 1,
album_index INTEGER DEFAULT 1,
status TEXT DEFAULT 'pending',
attempts INTEGER DEFAULT 0,
error TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(chat_id, message_id)
)
""")
await self.conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)")
await self.conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_media_group ON jobs(media_group_id)")
await self.conn.execute("""
CREATE TABLE IF NOT EXISTS scrape_checkpoints (
chat_id INTEGER PRIMARY KEY,
chat_title TEXT,
chat_type TEXT,
last_scraped_id INTEGER DEFAULT 0,
total_queued INTEGER DEFAULT 0,
status TEXT DEFAULT 'idle',
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
await self.conn.commit()
logger.info("Database initialized with single connection lock.")
async def execute(self, query: str, params: tuple = ()) -> aiosqlite.Cursor:
async with self.lock:
try:
cur = await self.conn.execute(query, params)
await self.conn.commit()
return cur
except aiosqlite.DatabaseError as e:
logger.error(f"Database error during query, auto-recovering: {e}")
await self.connect()
cur = await self.conn.execute(query, params)
await self.conn.commit()
return cur
async def fetchone(self, query: str, params: tuple = ()):
async with self.lock:
try:
cur = await self.conn.execute(query, params)
row = await cur.fetchone()
return dict(row) if row else None
except aiosqlite.DatabaseError as e:
logger.error(f"Database error during fetchone, auto-recovering: {e}")
await self.connect()
cur = await self.conn.execute(query, params)
row = await cur.fetchone()
return dict(row) if row else None
async def fetchall(self, query: str, params: tuple = ()):
async with self.lock:
try:
cur = await self.conn.execute(query, params)
rows = await cur.fetchall()
return [dict(r) for r in rows]
except aiosqlite.DatabaseError as e:
logger.error(f"Database error during fetchall, auto-recovering: {e}")
await self.connect()
cur = await self.conn.execute(query, params)
rows = await cur.fetchall()
return [dict(r) for r in rows]
db = SafeDBManager(DB_PATH)
async def add_job_to_queue(message: Message, album_info: Optional[Dict[str, int]] = None) -> bool:
if not message.media:
return False
chat_id = int(message.chat.id)
message_id = int(message.id)
media_group_id = str(message.media_group_id) if message.media_group_id else None
media_type = "unknown"
if message.photo: media_type = "photo"
elif message.video: media_type = "video"
elif message.document: media_type = "document"
elif message.audio: media_type = "audio"
elif message.voice: media_type = "voice"
elif message.video_note: media_type = "video_note"
elif message.animation: media_type = "animation"
chat = message.chat
source_name = chat.title or chat.first_name or chat.username or f"Chat_{chat_id}"
source_username = chat.username or ""
raw_caption = message.caption or ""
album_total = 1
album_index = 1
if album_info:
album_total = album_info.get("total", 1)
album_index = album_info.get("index", 1)
try:
cur = await db.execute("""
INSERT OR IGNORE INTO jobs
(chat_id, message_id, media_group_id, media_type, caption, source_name, source_username, album_total, album_index, status)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')
""", (chat_id, message_id, media_group_id, media_type, raw_caption, source_name, source_username, album_total, album_index))
if cur.rowcount > 0:
stats["total_queued"] += 1
return True
return False
except Exception as e:
logger.error(f"Error queuing message {message_id}: {e}")
return False
def format_smart_caption(raw_cap: str, source_name: str, source_username: str, msg_id: int, album_index: int, album_total: int) -> str:
lines = []
if raw_cap and raw_cap.strip():
lines.append(raw_cap.strip())
lines.append("")
if source_username:
source_tag = f"[{source_name}](https://t.me/{source_username}/{msg_id})"
else:
source_tag = f"**{source_name}**"
lines.append(f"๐Ÿ“ข ุงู„ู…ุตุฏุฑ: {source_tag}")
if album_total > 1:
lines.append(f"๐Ÿ”ข ุงู„ู…ู‚ุทุน: {album_index} ู…ู† ุฅุฌู…ุงู„ูŠ {album_total}")
return "\n".join(lines)
async def process_job(job: Dict[str, Any]) -> str:
chat_id = job["chat_id"]
msg_id = job["message_id"]
raw_caption = job["caption"]
source_name = job["source_name"]
source_username = job.get("source_username", "")
album_index = job["album_index"]
album_total = job["album_total"]
mtype = job["media_type"]
caption = format_smart_caption(raw_caption, source_name, source_username, msg_id, album_index, album_total)
await flood_manager.wait()
try:
await app.copy_message(
chat_id=TARGET_CHAT,
from_chat_id=chat_id,
message_id=msg_id,
caption=caption
)
logger.info(f"โšก Fast copied: {msg_id} from [{source_name}]")
return "SUCCESS"
except (FloodWait, FloodPremiumWait) as e:
flood_manager.handle_flood(getattr(e, 'value', 30))
raise e
except Exception:
pass
downloaded_file = None
try:
msg = await app.get_messages(chat_id, msg_id)
if not msg or not msg.media:
return "DELETED"
ext_map = {
"photo": ".jpg",
"video": ".mp4",
"animation": ".mp4",
"video_note": ".mp4",
"audio": ".mp3",
"voice": ".ogg",
"document": ""
}
ext = ext_map.get(mtype, "")
target_path = os.path.join(DOWNLOAD_DIR, f"media_{chat_id}_{msg_id}{ext}")
downloaded = await app.download_media(msg, file_name=target_path)
if not downloaded or not os.path.exists(downloaded):
raise Exception("Download failed")
downloaded_file = downloaded
if mtype == "photo":
try:
await app.send_photo(TARGET_CHAT, photo=downloaded_file, caption=caption)
except Exception:
await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)
elif mtype == "video":
await app.send_video(TARGET_CHAT, video=downloaded_file, caption=caption, supports_streaming=True)
elif mtype == "document":
await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)
elif mtype == "audio":
await app.send_audio(TARGET_CHAT, audio=downloaded_file, caption=caption)
elif mtype == "voice":
await app.send_voice(TARGET_CHAT, voice=downloaded_file, caption=caption)
elif mtype == "video_note":
await app.send_video_note(TARGET_CHAT, video_note=downloaded_file)
elif mtype == "animation":
await app.send_animation(TARGET_CHAT, animation=downloaded_file, caption=caption)
else:
await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)
logger.info(f"๐Ÿ“ฅ Downloaded & uploaded: {msg_id} from [{source_name}]")
return "SUCCESS"
except (FloodWait, FloodPremiumWait) as e:
flood_manager.handle_flood(getattr(e, 'value', 30))
raise e
except Exception as e:
logger.error(f"โŒ Failed processing {msg_id}: {e}")
return "RETRY"
finally:
if downloaded_file and os.path.exists(downloaded_file):
try:
if os.path.isdir(downloaded_file):
shutil.rmtree(downloaded_file, ignore_errors=True)
else:
os.remove(downloaded_file)
except Exception:
pass
async def queue_worker_loop():
logger.info("โš™๏ธ Queue worker loop started")
while True:
try:
if not app.is_connected:
await asyncio.sleep(5)
continue
job = await db.fetchone("""
SELECT * FROM jobs
WHERE status = 'pending' AND attempts < 10
ORDER BY attempts ASC, id ASC LIMIT 1
""")
if not job:
await asyncio.sleep(2)
continue
await db.execute("UPDATE jobs SET status='processing' WHERE id=?", (job["id"],))
result = await process_job(job)
if result == "SUCCESS":
await db.execute("UPDATE jobs SET status='completed', error=NULL WHERE id=?", (job["id"],))
stats["total_transferred"] += 1
elif result == "DELETED":
await db.execute("UPDATE jobs SET status='failed', error='Deleted' WHERE id=?", (job["id"],))
stats["total_failed"] += 1
else:
await db.execute("""
UPDATE jobs SET status='pending', attempts=attempts+1, error='Retry' WHERE id=?
""", (job["id"],))
except (FloodWait, FloodPremiumWait) as e:
sec = getattr(e, 'value', 30)
logger.warning(f"FloodWait sleeping {sec}s")
await asyncio.sleep(sec + 1)
except Exception as e:
logger.error(f"Worker exception: {e}")
await asyncio.sleep(4)
async def warm_up_dialogs_cache() -> List[Dict[str, Any]]:
logger.info("๐Ÿ”„ Warming up dialogs cache...")
discovered = []
try:
async for dialog in app.get_dialogs():
chat = dialog.chat
if not chat:
continue
if TARGET_CHAT_ID and chat.id == TARGET_CHAT_ID:
continue
if str(chat.id) == str(TARGET_CHAT):
continue
type_str = str(chat.type).lower()
if any(k in type_str for k in ["channel", "supergroup", "group"]):
title = chat.title or chat.first_name or f"Chat_{chat.id}"
chat_info = {
"id": chat.id,
"title": title,
"username": chat.username or "",
"type": type_str
}
KNOWN_CHATS[chat.id] = chat_info
discovered.append(chat_info)
logger.info(f"๐Ÿ“Œ Found: [{title}] ({chat.id}) - {type_str}")
stats["discovered_chats"] = len(discovered)
logger.info(f"โœ… Discovered {len(discovered)} chats/channels")
return discovered
except Exception as e:
logger.error(f"Dialogs error: {e}")
return []
async def scrape_single_chat(chat_id: int, chat_title: str, limit: int = 0) -> int:
try:
last_row = await db.fetchone("SELECT last_scraped_id FROM scrape_checkpoints WHERE chat_id=?", (chat_id,))
last_id = last_row["last_scraped_id"] if last_row else 0
logger.info(f"๐Ÿ” Scraping: {chat_title} (from msg {last_id})...")
count = 0
media_groups: Dict[str, List[Message]] = {}
highest_msg_id = last_id
async for message in app.get_chat_history(chat_id, limit=limit if limit > 0 else 0):
if message.id <= last_id and last_id > 0:
break
if message.id > highest_msg_id:
highest_msg_id = message.id
if message.media:
if message.media_group_id:
mg_id = str(message.media_group_id)
media_groups.setdefault(mg_id, []).append(message)
else:
if await add_job_to_queue(message):
count += 1
if count > 0 and count % 100 == 0:
await db.execute("""
INSERT INTO scrape_checkpoints (chat_id, chat_title, last_scraped_id, total_queued, status)
VALUES (?, ?, ?, ?, 'in_progress')
ON CONFLICT(chat_id) DO UPDATE SET
last_scraped_id=excluded.last_scraped_id,
total_queued=total_queued+excluded.total_queued,
updated_at=CURRENT_TIMESTAMP
""", (chat_id, chat_title, highest_msg_id, 100))
await asyncio.sleep(0.1)
for mg_id, msgs in media_groups.items():
msgs.sort(key=lambda m: m.id)
total = len(msgs)
for idx, msg in enumerate(msgs, start=1):
if await add_job_to_queue(msg, album_info={"total": total, "index": idx}):
count += 1
await db.execute("""
INSERT INTO scrape_checkpoints (chat_id, chat_title, last_scraped_id, total_queued, status)
VALUES (?, ?, ?, ?, 'completed')
ON CONFLICT(chat_id) DO UPDATE SET
last_scraped_id=excluded.last_scraped_id,
total_queued=total_queued+excluded.total_queued,
status='completed',
updated_at=CURRENT_TIMESTAMP
""", (chat_id, chat_title, highest_msg_id, count))
logger.info(f"โœ… Finished {chat_title}: queued {count} files.")
return count
except Exception as e:
logger.warning(f"Skipping {chat_title} ({chat_id}): {e}")
return 0
async def scrape_all_joined_chats(limit_per_chat: int = 0) -> str:
stats["active_scans"] += 1
try:
chats = await warm_up_dialogs_cache()
if not chats:
return "No chats found."
logger.info(f"๐Ÿš€ Scraping all {len(chats)} chats...")
total_all = 0
for idx, chat in enumerate(chats, 1):
c_id = chat["id"]
c_title = chat["title"]
logger.info(f"[{idx}/{len(chats)}] Scraping: {c_title}...")
queued = await scrape_single_chat(c_id, c_title, limit=limit_per_chat)
total_all += queued
await asyncio.sleep(1.0)
msg = f"Completed all chats ({len(chats)}). Queued {total_all} files."
logger.info(msg)
return msg
except Exception as e:
err = f"Scrape all error: {e}"
logger.error(err)
return err
finally:
stats["active_scans"] = max(0, stats["active_scans"] - 1)
@app.on_message(filters.incoming & ~filters.me & filters.media)
async def realtime_all_chats_handler(client, message: Message):
try:
if TARGET_CHAT_ID and message.chat.id == TARGET_CHAT_ID:
return
if str(message.chat.id) == str(TARGET_CHAT):
return
type_str = str(message.chat.type).lower()
if message.chat.id not in KNOWN_CHATS and any(k in type_str for k in ["channel", "supergroup", "group"]):
KNOWN_CHATS[message.chat.id] = {
"id": message.chat.id,
"title": message.chat.title or f"Chat_{message.chat.id}",
"username": message.chat.username or "",
"type": type_str
}
if message.media_group_id:
await asyncio.sleep(2.0)
try:
album_msgs = await client.get_media_group(message.chat.id, message.id)
album_msgs.sort(key=lambda m: m.id)
total = len(album_msgs)
for idx, msg in enumerate(album_msgs, start=1):
await add_job_to_queue(msg, album_info={"total": total, "index": idx})
except Exception:
await add_job_to_queue(message)
else:
await add_job_to_queue(message)
logger.info(f"โœจ Realtime media: [{message.chat.title or message.chat.id}]")
except Exception as e:
logger.error(f"Realtime error: {e}")
async def resolve_target_chat():
global TARGET_CHAT, TARGET_CHAT_ID
raw = TARGET_CHAT_RAW.strip()
if raw.lower() in ["me", "saved", ""]:
TARGET_CHAT = "me"
TARGET_CHAT_ID = None
logger.info("Target: Saved Messages (me)")
return
try:
chat_id_int = int(raw)
chat = await app.get_chat(chat_id_int)
TARGET_CHAT = chat_id_int
TARGET_CHAT_ID = chat.id
logger.info(f"Target: {chat.title} ({TARGET_CHAT})")
except ValueError:
try:
chat = await app.get_chat(raw)
TARGET_CHAT = chat.id
TARGET_CHAT_ID = chat.id
logger.info(f"Target: {chat.title} ({TARGET_CHAT})")
except Exception as e:
logger.error(f"Target error ({raw}): {e}")
TARGET_CHAT = "me"
def create_dashboard():
with gr.Blocks(title="Telegram Archiver") as demo:
gr.Markdown("# ๐ŸŒ Universal Telegram Archiver Pro")
with gr.Row():
gr.Textbox(label="Status", value="Active โœ…", interactive=False)
gr.Number(label="Discovered Chats", value=0, interactive=False)
gr.Number(label="Queue Size", value=0, interactive=False)
gr.Number(label="Transferred", value=0, interactive=False)
with gr.Row():
scan_all_btn = gr.Button("๐Ÿ”„ Scan & Scrape All Joined Chats", variant="primary", scale=2)
limit_input = gr.Number(label="Limit per chat (0 = All)", value=0, scale=1)
scrape_status = gr.Textbox(label="Operation Result", interactive=False)
def trigger_scan_all(lim):
if MAIN_LOOP and MAIN_LOOP.is_running():
asyncio.run_coroutine_threadsafe(
scrape_all_joined_chats(int(lim)),
MAIN_LOOP
)
return "Started full scan in background!"
return "Bot is not connected."
scan_all_btn.click(fn=trigger_scan_all, inputs=[limit_input], outputs=[scrape_status])
return demo
def run_web_server():
fastapi_app = FastAPI()
@fastapi_app.get("/health")
async def health():
return {"status": "ok", "stats": stats}
@fastapi_app.get("/ping")
async def ping():
return {"status": "alive", "stats": stats}
demo = create_dashboard()
gr.mount_gradio_app(fastapi_app, demo, path="/")
uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
async def main():
global MAIN_LOOP
MAIN_LOOP = asyncio.get_running_loop()
logger.info("Connecting to Telegram...")
await app.start()
me = await app.get_me()
logger.info(f"Logged in: {me.first_name} (@{me.username or me.id})")
await resolve_target_chat()
await db.connect()
asyncio.create_task(queue_worker_loop())
discovered = await warm_up_dialogs_cache()
logger.info(f"Ready: {len(discovered)} chats.")
if AUTO_SCAN_ALL:
logger.info("Auto-scanning all chats in background...")
asyncio.create_task(scrape_all_joined_chats(limit_per_chat=0))
stats["status"] = "Active"
stop_event = asyncio.Event()
await stop_event.wait()
if __name__ == "__main__":
web_thread = threading.Thread(target=run_web_server, daemon=True)
web_thread.start()
loop = asyncio.get_event_loop()
try:
loop.run_until_complete(main())
except KeyboardInterrupt:
pass
finally:
if app.is_connected:
loop.run_until_complete(app.stop())