arsam
43729eeaa6
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
301 lines
10 KiB
Python
301 lines
10 KiB
Python
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())
|