diff --git a/bale.py b/bale.py new file mode 100644 index 0000000..2361883 --- /dev/null +++ b/bale.py @@ -0,0 +1,300 @@ +import os +import re +import json +import time +import asyncio +from pathlib import Path +from typing import Optional + +from dotenv import load_dotenv +from aiobale import Client, Dispatcher +from aiobale.types import FileInput +from aiobale.enums import ChatType +import requests +import pyzipper +from urllib.parse import urlparse + +load_dotenv(override=True) + +BALE_TARGET_ID = int(os.getenv("BALE_TARGET_ID", "0")) + +BASE_DIR = Path(__file__).resolve().parent +DOWNLOAD_DIR = BASE_DIR / "downloads" +QUEUE_DIR = BASE_DIR / "queue" +QUEUE_FILE = QUEUE_DIR / "tasks.jsonl" +PROCESSING_FILE = QUEUE_DIR / "processing.json" +FAILED_FILE = QUEUE_DIR / "failed.jsonl" +STATUS_FILE = QUEUE_DIR / "status.jsonl" +URL_DIR = DOWNLOAD_DIR / "url" +CANCEL_FILE = QUEUE_DIR / "cancelled.jsonl" + +MAX_RETRIES = 8 +UPLOAD_TIMEOUT = 3600 + +DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) +QUEUE_DIR.mkdir(parents=True, exist_ok=True) +URL_DIR.mkdir(parents=True, exist_ok=True) + +dp = Dispatcher() +client = Client(dispatcher=dp, session_file=BASE_DIR / "session.bale") + +def safe_filename(name: Optional[str]) -> str: + name = (name or "file").strip() + name = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", name) + name = name.rstrip(". ") + return name[:200] or "file" + +def pretty_size(size) -> str: + size = float(size or 0) + units = ["B", "KB", "MB", "GB"] + index = 0 + while size >= 1024 and index < len(units) - 1: + size /= 1024 + index += 1 + return f"{size:.2f} {units[index]}" + +def eta_text(seconds) -> str: + if not seconds or seconds <= 0: + return "نامشخص" + seconds = int(seconds) + h = seconds // 3600 + m = (seconds % 3600) // 60 + s = seconds % 60 + if h: return f"{h}h {m}m" + if m: return f"{m}m {s}s" + return f"{s}s" + +def push_status(task: dict, text: str, status: str = "working", percent: float | None = None): + payload = { + "chat_id": task.get("chat_id"), + "message_id": task.get("status_message_id"), + "status": status, + "text": text, + "percent": percent, + "time": time.time(), + } + with open(STATUS_FILE, "a", encoding="utf-8") as file: + file.write(json.dumps(payload, ensure_ascii=False) + "\n") + +def is_cancelled(task: dict) -> bool: + job_id = str(task.get("job_id", "")) + if not job_id or not CANCEL_FILE.exists(): + return False + with open(CANCEL_FILE, "r", encoding="utf-8") as file: + for line in file: + if not line.strip(): continue + item = json.loads(line) + if str(item.get("job_id")) == job_id: + return True + return False + +def unique_path(path: Path) -> Path: + if not path.exists(): return path + stem, suffix = path.stem, path.suffix + index = 1 + while True: + candidate = path.with_name(f"{stem}_{index}{suffix}") + if not candidate.exists(): return candidate + index += 1 + +async def send_with_retry(file_path: str, caption: str = "", task: dict | None = None): + last_error = None + start_time = time.time() + + for attempt in range(1, MAX_RETRIES + 1): + if time.time() - start_time > UPLOAD_TIMEOUT: + raise RuntimeError("آپلود بیشتر از حد مجاز طول کشید و لغو شد.") + + if task and is_cancelled(task): + raise RuntimeError("ارسال لغو شد.") + + try: + if task: + push_status( + task, + f"🔼 در حال آپلود در بله (aiobale)...\n\n" + f"🔴 تلاش {attempt} از {MAX_RETRIES}\n\n" + f"برای لغو ارسال:\n" + f"`/del {task.get('job_id')}`", + "uploading" + ) + + await client.send_document( + file=FileInput(file_path), + chat_id=BALE_TARGET_ID, + chat_type=ChatType.PRIVATE, + caption=caption or "" + ) + return + + except Exception as e: + last_error = e + error_text = str(e).lower() + transient = any(k in error_text for k in ["timeout", "connection", "502", "503", "504", "500", "rate limit", "rate_limit", "flood"]) + if transient and attempt < MAX_RETRIES: + if "rate" in error_text or "flood" in error_text: + wait = min(30 * (2 ** (attempt - 1)), 300) + elif "500" in error_text: + wait = 30 + else: + wait = 5 + push_status(task, f"⏳ تلاش {attempt} ناموفق ({str(e)[:60]})\nتلاش مجدد در {wait} ثانیه...", "retrying") if task else None + await asyncio.sleep(wait) + continue + else: + raise e + + raise last_error if last_error else RuntimeError("Upload failed.") + +async def download_url(task: dict) -> Path: + url = task.get("url", "").strip() + if not url: raise RuntimeError("URL خالیه") + + push_status(task, "در حال دانلود ...", "downloading", 0) + + try: + resp = requests.get(url, stream=True, timeout=(10, 60), allow_redirects=True) + resp.raise_for_status() + except Exception as e: + raise RuntimeError(f"خطا در دریافت لینک: {str(e)}") + + cd = resp.headers.get("content-disposition", "") + match = re.findall(r'filename="(.+?)"', cd) + name = match[0] if match else Path(urlparse(url).path).name + name = safe_filename(name or f"file_{int(time.time())}") + if "." not in name: name += ".bin" + + target = unique_path(URL_DIR / name) + total = int(resp.headers.get("content-length") or 0) + downloaded, last_update, started = 0, 0, time.time() + + with open(target, "wb") as f: + for chunk in resp.iter_content(1024 * 1024): + if not chunk: continue + f.write(chunk) + downloaded += len(chunk) + + now = time.time() + if now - last_update < 3 and downloaded < total: continue + last_update = now + + speed = downloaded / max(now - started, 1) + eta = (total - downloaded) / speed if total and speed else None + percent = downloaded * 100 / total if total else None + + text = f"داره دانلود میکنه...\n\n{pretty_size(downloaded)}" + if total: text += f" از {pretty_size(total)}" + text += f"\nسرعت: {pretty_size(speed)}/s" + if eta: text += f"\nمونده: {eta_text(eta)}" + + push_status(task, text, "downloading", percent) + + if not target.exists() or target.stat().st_size == 0: + raise RuntimeError("فایل دانلود نشد") + + task["file_name"] = target.name + task["file_size"] = target.stat().st_size + return target + +def make_zip_with_password(file_path: Path, password: str) -> Path: + zip_path = unique_path(file_path.with_suffix(file_path.suffix + ".zip")) + with pyzipper.AESZipFile( + zip_path, "w", compression=pyzipper.ZIP_STORED, encryption=pyzipper.WZ_AES + ) as zip_file: + zip_file.setpassword(password.encode("utf-8")) + zip_file.write(file_path, arcname=file_path.name) + return zip_path + +def pop_first_task(): + if not QUEUE_FILE.exists(): return None + with open(QUEUE_FILE, "r", encoding="utf-8") as file: + lines = [line for line in file if line.strip()] + if not lines: return None + first_line = lines[0] + remaining = lines[1:] + with open(QUEUE_FILE, "w", encoding="utf-8") as file: + file.writelines(remaining) + return json.loads(first_line) + +def save_processing(task: dict) -> None: + with open(PROCESSING_FILE, "w", encoding="utf-8") as file: + json.dump(task, file, ensure_ascii=False, indent=2) + +def clear_processing() -> None: + if PROCESSING_FILE.exists(): PROCESSING_FILE.unlink() + +def append_failed(task: dict, error: str) -> None: + payload = {"task": task, "error": error} + with open(FAILED_FILE, "a", encoding="utf-8") as file: + file.write(json.dumps(payload, ensure_ascii=False) + "\n") + +async def process_task(task: dict): + task_type = task.get("type") + caption = task.get("caption", "") + safe_mode = task.get("safe_mode", False) + zip_password = task.get("zip_password", "") + + local_path: Path | None = None + + if task_type == "local_file": + local_path = Path(task.get("path", "")) + if not local_path.exists(): raise RuntimeError("Local file not found.") + elif task_type == "direct_url": + local_path = await download_url(task) + else: + raise RuntimeError("Unknown task type.") + + if safe_mode and zip_password: + push_status(task, "در حال تبدیل به فایل zip ...", "processing") + try: + zipped = make_zip_with_password(local_path, zip_password) + finally: + try: + if local_path.exists(): local_path.unlink() + except Exception: pass + send_path = zipped + else: + send_path = local_path + + try: + if is_cancelled(task): raise RuntimeError("ارسال لغو شد.") + await send_with_retry(str(send_path), caption, task) + push_status(task, "فایل با موفقیت در بله آپلود شد.", "done") + finally: + try: + if send_path and send_path.exists(): send_path.unlink() + except Exception: pass + +async def worker_loop(): + print("Bale aiobale worker starting...") + # NOTE: If this is the first run, it will ask for phone number in the terminal. + # User should run this once manually: python3 bale.py + await client.start(run_in_background=True) + me = await client.get_me() + print(f"Bale aiobale worker started. Logged in as: {me.id} ({me.name})") + if BALE_TARGET_ID != me.id: + print(f"WARNING: BALE_TARGET_ID={BALE_TARGET_ID} != your Bale user ID={me.id}. For Saved Messages, set BALE_TARGET_ID={me.id}") + + while True: + task = pop_first_task() + if not task: + await asyncio.sleep(1) + continue + + save_processing(task) + try: + await process_task(task) + except Exception as e: + append_failed(task, str(e)) + push_status(task, f"خطا: {str(e)}", "failed") + finally: + clear_processing() + +if __name__ == "__main__": + loop = asyncio.get_event_loop() + try: + loop.run_until_complete(worker_loop()) + except KeyboardInterrupt: + pass + finally: + loop.run_until_complete(client.stop())