Import source: rub

Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.

rub.py

@@ -0,0 +1,444 @@
+import os
+import re
+import json
+import time
+from pathlib import Path
+from typing import Optional
+
+from dotenv import load_dotenv
+from rubpy import Client as RubikaClient
+import requests
+import pyzipper
+from urllib.parse import urlparse
+import threading
+
+load_dotenv()
+
+SESSION = os.getenv("RUBIKA_SESSION", "rubika_session").strip()
+
+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 = 5
+UPLOAD_TIMEOUT = 1800
+TARGET = "me"
+
+DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
+QUEUE_DIR.mkdir(parents=True, exist_ok=True)
+URL_DIR.mkdir(parents=True, exist_ok=True)
+
+
+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 get_per_attempt_timeout(file_path: str) -> int:
+    size_mb = Path(file_path).stat().st_size / (1024 * 1024)
+
+    if size_mb < 100:
+        return 180
+    elif size_mb < 500:
+        return 420
+    elif size_mb < 1000:
+        return 720
+    else:
+        return 1200
+    
+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 = path.stem
+    suffix = path.suffix
+    index = 1
+
+    while True:
+        candidate = path.with_name(f"{stem}_{index}{suffix}")
+        if not candidate.exists():
+            return candidate
+        index += 1
+
+
+def has_session(session_name: str) -> bool:
+    candidates = [
+        Path(session_name),
+        Path(f"{session_name}.session"),
+        Path(f"{session_name}.sqlite"),
+    ]
+    return any(path.exists() for path in candidates)
+
+
+def ensure_session():
+    if has_session(SESSION):
+        return
+
+    client = RubikaClient(name=SESSION)
+
+    try:
+        client.start()
+        print("Login successful.")
+    finally:
+        try:
+            client.disconnect()
+        except Exception:
+            pass
+
+
+def send_document(file_path: str, caption: str = ""):
+    client = RubikaClient(name=SESSION)
+
+    try:
+        client.start()
+        return client.send_document(
+            TARGET,
+            file_path,
+            caption=caption or ""
+        )
+    finally:
+        try:
+            client.disconnect()
+        except Exception:
+            pass
+
+def send_with_timeout(file_path, caption, timeout):
+    result = {}
+    error = {}
+
+    def target():
+        try:
+            result["data"] = send_document(file_path, caption)
+        except Exception as e:
+            error["err"] = e
+
+    t = threading.Thread(target=target)
+    t.start()
+    t.join(timeout)
+
+    if t.is_alive():
+        raise RuntimeError("آپلود بیشتر از حد مجاز طول کشید و لغو شد.")
+
+    if "err" in error:
+        raise error["err"]
+
+    return result.get("data")
+
+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"🔼 در حال آپلود در روبیکا...\n\n"
+                    f"🔴 تلاش {attempt} از {MAX_RETRIES}\n\n"
+                    f"برای لغو ارسال:\n"
+                    f"`/del {task.get('job_id')}`",
+                    "uploading"
+                )
+
+            elapsed = time.time() - start_time
+            remaining = UPLOAD_TIMEOUT - elapsed
+
+            if remaining <= 0:
+                raise RuntimeError("آپلود بیشتر از ۳۰ دقیقه طول کشید و لغو شد.")
+
+            per_attempt = min(get_per_attempt_timeout(file_path), remaining)
+
+            return send_with_timeout(file_path, caption, per_attempt)
+
+        except Exception as e:
+            last_error = e
+            error_text = str(e).lower()
+
+            transient = any(
+                key in error_text
+                for key in [
+                    "502", "503", "bad gateway", "timeout",
+                    "cannot connect", "connection reset",
+                    "temporarily unavailable",
+                    "error uploading chunk",
+                    "unexpected mimetype",
+                ]
+            )
+
+            if transient and attempt < MAX_RETRIES:
+
+                if task and is_cancelled(task):
+                    raise RuntimeError("ارسال لغو شد.")
+
+                if task:
+                    push_status(
+                        task,
+                        f"ارتباط با روبیکا ناپایدار بود...\n"
+                        f"دوباره تلاش می‌کنم ({attempt + 1})",
+                        "uploading"
+                    )
+
+                time.sleep(3)
+                continue
+
+    raise last_error if last_error else RuntimeError("Upload failed.")
+
+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 requests.exceptions.Timeout:
+        raise RuntimeError("لینک جواب نداد")
+    except requests.exceptions.ConnectionError:
+        raise RuntimeError("مشکل شبکه")
+    except requests.exceptions.HTTPError as e:
+        code = e.response.status_code if e.response else "نامشخص"
+        raise RuntimeError(f"دانلود انجام نشد. کد خطا: {code}")
+    
+    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")
+
+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 = 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("ارسال لغو شد.")
+
+        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
+
+def worker_loop():
+    ensure_session()
+    print("Rubika worker started.")
+
+    while True:
+        task = pop_first_task()
+
+        if not task:
+            time.sleep(0.2)
+            continue
+
+        save_processing(task)
+
+        try:
+            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__":
+    worker_loop()