arsam
9ecc70d162
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
1368 lines
54 KiB
Python
1368 lines
54 KiB
Python
import os
|
||
import re
|
||
import sys
|
||
import json
|
||
import asyncio
|
||
import logging
|
||
import tempfile
|
||
import time
|
||
import collections
|
||
from datetime import datetime
|
||
from urllib.parse import urlparse, unquote
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
import psutil
|
||
|
||
from pyrogram import Client, filters
|
||
from pyrogram.types import Message
|
||
|
||
load_dotenv()
|
||
|
||
API_ID = int(os.getenv("API_ID", "0"))
|
||
API_HASH = os.getenv("API_HASH", "").strip()
|
||
BOT_TOKEN = os.getenv("BOT_TOKEN", "").strip()
|
||
RCLONE_REMOTE = "gdrive:"
|
||
|
||
def classify_to_folder(file_name: str, mime_type: str = "") -> str:
|
||
ext = file_name.rsplit(".", 1)[-1].lower() if "." in file_name else ""
|
||
if mime_type.startswith("video/") or ext in ("mp4","mkv","avi","mov","webm","flv"):
|
||
return "Videos"
|
||
if mime_type.startswith("audio/") or ext in ("mp3","flac","wav","ogg","m4a","aac"):
|
||
return "Music"
|
||
if mime_type.startswith("image/") or ext in ("jpg","jpeg","png","gif","webp","bmp","heic"):
|
||
return "Photos"
|
||
if ext in ("exe","msi","deb","rpm","dmg","appimage"):
|
||
return "Software"
|
||
if ext in ("zip","rar","7z","tar","gz","bz2","xz","zst"):
|
||
return "Archives"
|
||
if ext == "gguf" or ext in ("safetensors","pt","ckpt","pth"):
|
||
return "Models"
|
||
if ext == "zim":
|
||
return "ZIM"
|
||
if ext in ("conf","ovpn","cfg","ini"):
|
||
return "Config"
|
||
if ext in ("pdf","doc","docx","txt","md","csv","xlsx","pptx"):
|
||
return "Docs"
|
||
if ext == "bin":
|
||
return "Bin"
|
||
return "Other"
|
||
|
||
LOG_FILE = Path(__file__).parent / "bot.log"
|
||
USERS_FILE = Path(__file__).parent / "authorized_users.json"
|
||
|
||
STREAM_THRESHOLD = 2 * (1 << 30) # 2 GB
|
||
QUOTAS_FILE = Path(__file__).parent / "quotas.json"
|
||
|
||
# ── Persisted authorized users ────────────────────────────────────────────────
|
||
|
||
def _load_users() -> set:
|
||
try:
|
||
return set(json.loads(USERS_FILE.read_text()))
|
||
except Exception:
|
||
return {"coaxys"}
|
||
|
||
def _save_users(users: set):
|
||
USERS_FILE.write_text(json.dumps(sorted(users)))
|
||
|
||
AUTHORIZED_USERS: set = _load_users()
|
||
|
||
# ── Quota system ──────────────────────────────────────────────────────────────
|
||
# quotas.json schema:
|
||
# { "quotas": {"username": <bytes>}, "usage": {"username": {"YYYY-MM-DD": <bytes>}} }
|
||
|
||
def _load_quotas() -> dict:
|
||
try:
|
||
return json.loads(QUOTAS_FILE.read_text())
|
||
except Exception:
|
||
return {"quotas": {}, "usage": {}}
|
||
|
||
def _save_quotas(data: dict):
|
||
QUOTAS_FILE.write_text(json.dumps(data, indent=2))
|
||
|
||
def parse_size(s: str) -> int | None:
|
||
"""Parse '5GB', '500MB', '2.5GB', '100KB' → bytes. Returns None on failure."""
|
||
m = re.fullmatch(r'([\d.]+)\s*(B|KB|MB|GB|TB)', s.strip(), re.IGNORECASE)
|
||
if not m:
|
||
return None
|
||
n, unit = float(m.group(1)), m.group(2).upper()
|
||
mul = {"B": 1, "KB": 1<<10, "MB": 1<<20, "GB": 1<<30, "TB": 1<<40}
|
||
return int(n * mul[unit])
|
||
|
||
def today() -> str:
|
||
return datetime.utcnow().strftime("%Y-%m-%d")
|
||
|
||
def get_quota(username: str) -> int | None:
|
||
"""Daily quota in bytes, or None if not set."""
|
||
return _load_quotas()["quotas"].get(username.lower())
|
||
|
||
def get_usage_today(username: str) -> int:
|
||
return _load_quotas()["usage"].get(username.lower(), {}).get(today(), 0)
|
||
|
||
def quota_remaining(username: str) -> int | None:
|
||
"""Returns remaining bytes today, or None if no quota set."""
|
||
q = get_quota(username)
|
||
if q is None:
|
||
return None
|
||
return max(0, q - get_usage_today(username))
|
||
|
||
def record_upload(username: str, size_bytes: int):
|
||
"""Add size_bytes to today's usage for username."""
|
||
data = _load_quotas()
|
||
u = username.lower()
|
||
data["usage"].setdefault(u, {})
|
||
data["usage"][u][today()] = data["usage"][u].get(today(), 0) + size_bytes
|
||
_save_quotas(data)
|
||
|
||
def has_upload_access(message: Message) -> tuple[bool, str]:
|
||
"""
|
||
Returns (allowed, reason).
|
||
Auth users always pass. Others need a quota with remaining bytes.
|
||
"""
|
||
u = (message.from_user.username or "").lower() if message.from_user else ""
|
||
if u in AUTHORIZED_USERS:
|
||
return True, "auth"
|
||
q = get_quota(u)
|
||
if q is None:
|
||
return False, "no_access"
|
||
rem = quota_remaining(u)
|
||
if rem == 0:
|
||
return False, "quota_exceeded"
|
||
return True, "quota"
|
||
|
||
BOT_START_TIME = time.time()
|
||
session_uploads = 0
|
||
|
||
logging.basicConfig(
|
||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
level=logging.INFO
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
app = Client(
|
||
"gdrive_bot",
|
||
api_id=API_ID,
|
||
api_hash=API_HASH,
|
||
bot_token=BOT_TOKEN,
|
||
)
|
||
|
||
MEDIA_TYPES = (
|
||
"document", "video", "audio", "voice",
|
||
"photo", "animation", "video_note", "sticker"
|
||
)
|
||
|
||
URL_RE = re.compile(r'https?://\S+', re.IGNORECASE)
|
||
|
||
# ── Utilities ─────────────────────────────────────────────────────────────────
|
||
|
||
def get_media(message: Message):
|
||
for attr in MEDIA_TYPES:
|
||
m = getattr(message, attr, None)
|
||
if m:
|
||
return attr, m
|
||
return None, None
|
||
|
||
def sanitize_filename(name: str) -> str:
|
||
"""Clean up filenames from bots/channels: strip emoji, newlines, garbage chars."""
|
||
# Normalise unicode (NFKC collapses weird lookalikes)
|
||
import unicodedata
|
||
name = unicodedata.normalize("NFKC", name)
|
||
# Control chars and newlines → space
|
||
name = re.sub(r'[\x00-\x1f\x7f]', ' ', name)
|
||
# Common garbage from forwarded bot messages
|
||
name = name.replace('→', '').replace('↵', '').replace('␊', '')
|
||
# Characters illegal in filenames on most filesystems
|
||
name = re.sub(r'[<>:"/\\|?*]', '_', name)
|
||
# Emoji: remove everything in supplementary Unicode planes (U+1F000+) and misc symbol blocks
|
||
name = re.sub(
|
||
r'[\U0001F000-\U0010FFFF' # supplementary planes — most emoji live here
|
||
r'☀-➿' # misc symbols, dingbats
|
||
r'⌀-⏿' # misc technical
|
||
r'︀-️' # variation selectors
|
||
r'-' # zero-width / directional chars
|
||
r' -]', # general punctuation extras
|
||
'', name, flags=re.UNICODE
|
||
)
|
||
# Collapse runs of whitespace/underscores to a single underscore
|
||
name = re.sub(r'[\s_]+', '_', name).strip('_. ')
|
||
# Truncate stem to 180 chars but keep extension intact
|
||
root, ext = os.path.splitext(name)
|
||
root = root[:180].rstrip('_. ')
|
||
name = (root + ext).strip()
|
||
# Final fallback
|
||
return name or f"file_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||
|
||
def resolve_filename(media_type: str, media, caption: str | None) -> str:
|
||
# Caption overrides — rename before upload
|
||
if caption and caption.strip():
|
||
name = caption.strip()
|
||
orig = getattr(media, "file_name", None)
|
||
if orig and "." in orig and "." not in name:
|
||
name += os.path.splitext(orig)[1]
|
||
return sanitize_filename(name)
|
||
name = getattr(media, "file_name", None)
|
||
if name:
|
||
return sanitize_filename(name)
|
||
ext = {
|
||
"photo": ".jpg", "voice": ".ogg", "video_note": ".mp4",
|
||
"sticker": ".webp", "animation": ".mp4", "video": ".mp4", "audio": ".mp3",
|
||
}
|
||
return f"file_{datetime.now().strftime('%Y%m%d_%H%M%S')}{ext.get(media_type, '.bin')}"
|
||
|
||
def filename_from_url(url: str) -> str:
|
||
path = unquote(urlparse(url).path)
|
||
name = os.path.basename(path.rstrip("/"))
|
||
if not name or "." not in name:
|
||
name = f"download_{datetime.now().strftime('%Y%m%d_%H%M%S')}.bin"
|
||
return sanitize_filename(name)
|
||
|
||
def fmt_size(b: int) -> str:
|
||
if b >= 1 << 30: return f"{b / (1<<30):.2f} GB"
|
||
if b >= 1 << 20: return f"{b / (1<<20):.1f} MB"
|
||
if b >= 1 << 10: return f"{b / (1<<10):.1f} KB"
|
||
return f"{b} B"
|
||
|
||
def fmt_dur(s: float) -> str:
|
||
s = int(s)
|
||
if s < 60: return f"{s}s"
|
||
if s < 3600: return f"{s//60}m {s%60:02d}s"
|
||
return f"{s//3600}h {(s%3600)//60}m"
|
||
|
||
def fmt_ram() -> str:
|
||
vm = psutil.virtual_memory()
|
||
return f"{fmt_size(vm.used)} / {fmt_size(vm.total)} ({vm.percent:.0f}%)"
|
||
|
||
def spd(bps: float) -> str:
|
||
return f"{fmt_size(int(bps))}/s" if bps > 0 else "—"
|
||
|
||
def choose_chunk(file_bytes: int) -> str:
|
||
if file_bytes == 0: return "256M"
|
||
if file_bytes < 10 * (1<<20): return "8M"
|
||
if file_bytes < 100 * (1<<20): return "64M"
|
||
if file_bytes < 500 * (1<<20): return "256M"
|
||
return "512M"
|
||
|
||
async def get_drive_file_id(file_name: str) -> str | None:
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "lsjson", RCLONE_REMOTE,
|
||
"--recursive", "--files-only",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, _ = await proc.communicate()
|
||
files = json.loads(out.decode())
|
||
for f in files:
|
||
if f["Name"] == file_name:
|
||
return f["ID"]
|
||
return None
|
||
except Exception as e:
|
||
logger.warning(f"lsjson: {e}")
|
||
return None
|
||
|
||
async def send_links(message: Message, file_name: str):
|
||
file_id = await get_drive_file_id(file_name)
|
||
if file_id:
|
||
dl = f"https://drive.google.com/uc?export=download&id={file_id}"
|
||
view = f"https://drive.google.com/file/d/{file_id}/view"
|
||
wget = f"wget -O \"{file_name}\" \"{dl}\""
|
||
await message.reply_text(
|
||
f"🔗 **Links**\n\n"
|
||
f"📥 **Direct download:**\n`{dl}`\n\n"
|
||
f"👁 **View on Drive:**\n`{view}`\n\n"
|
||
f"🖥 **wget:**\n`{wget}`",
|
||
disable_web_page_preview=True,
|
||
)
|
||
|
||
# ── Auth ──────────────────────────────────────────────────────────────────────
|
||
|
||
def is_authorized(message: Message) -> bool:
|
||
u = message.from_user
|
||
if not u:
|
||
return False
|
||
return (u.username or "").lower() in AUTHORIZED_USERS
|
||
|
||
async def deny(message: Message):
|
||
await message.reply_text("⛔ You're not authorized to use this command.")
|
||
|
||
# ── /start ────────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("start"))
|
||
async def cmd_start(client: Client, message: Message):
|
||
vm = psutil.virtual_memory()
|
||
await message.reply_text(
|
||
"**GDrive Upload Bot** ⚡\n\n"
|
||
f"🖥 RAM: {fmt_size(vm.available)} free · 📡 MTProto + rclone\n\n"
|
||
"━━━━━━━━━━━━━━━━━━━━\n\n"
|
||
"**Send a file:**\n"
|
||
"📄 Documents 🖼 Photos 🎬 Videos\n"
|
||
"🎵 Audio 🎤 Voice 📹 Video notes 😀 Stickers\n\n"
|
||
"**Send a URL:**\n"
|
||
"🌐 Paste any http:// or https:// link — I'll download and upload it for you.\n\n"
|
||
"**✏️ Rename before upload:**\n"
|
||
"Add a caption to any file to rename it on Drive.\n"
|
||
"_Example: send a video with caption_ `lecture_01.mp4`\n\n"
|
||
"**🔗 Direct download links:**\n"
|
||
"After every upload I send a wget-ready link so you can pull the file anywhere.\n\n"
|
||
"No file size limit — uses MTProto for large Telegram downloads.\n\n"
|
||
"━━━━━━━━━━━━━━━━━━━━\n"
|
||
"Made by @Coaxys"
|
||
)
|
||
|
||
# ── /list ────────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("list"))
|
||
async def cmd_list(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
wait = await message.reply_text("📂 Fetching latest files from Drive...")
|
||
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "lsjson", RCLONE_REMOTE,
|
||
"--recursive", "--files-only",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(err.decode().strip())
|
||
|
||
files = json.loads(out.decode())
|
||
files.sort(key=lambda x: x["ModTime"], reverse=True)
|
||
top = files[:10]
|
||
total = len(files)
|
||
|
||
lines = []
|
||
for i, f in enumerate(top, 1):
|
||
size = f["Size"]
|
||
if size >= 1<<30: s = f"{size/(1<<30):.2f} GB"
|
||
elif size >= 1<<20: s = f"{size/(1<<20):.1f} MB"
|
||
elif size >= 1<<10: s = f"{size/(1<<10):.1f} KB"
|
||
else: s = f"{size} B"
|
||
date = f["ModTime"][:10]
|
||
name = f["Name"]
|
||
if len(name) > 40:
|
||
name = name[:37] + "..."
|
||
fid = f["ID"]
|
||
folder = f["Path"].split("/")[0] if "/" in f["Path"] else "root"
|
||
link = f"https://drive.google.com/file/d/{fid}/view"
|
||
lines.append(f"`{i:2}.` [{name}]({link})\n 📁 {folder} · 📦 {s} · 📅 {date}")
|
||
|
||
body = "\n\n".join(lines)
|
||
await wait.edit_text(
|
||
f"📂 **Latest {len(top)} of {total} files in Drive**\n\n{body}",
|
||
disable_web_page_preview=True,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"/list error: {e}")
|
||
await wait.edit_text(f"❌ Failed to fetch file list\n\n`{e}`")
|
||
|
||
# ── /listall ─────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("listall"))
|
||
async def cmd_listall(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
wait = await message.reply_text("📂 Fetching all files from Drive...")
|
||
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "lsjson", RCLONE_REMOTE,
|
||
"--recursive", "--files-only",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(err.decode().strip())
|
||
|
||
files = json.loads(out.decode())
|
||
files.sort(key=lambda x: x["ModTime"], reverse=True)
|
||
|
||
if not files:
|
||
return await wait.edit_text("📂 Drive is empty.")
|
||
|
||
total_size = sum(f["Size"] for f in files)
|
||
|
||
def human(b):
|
||
if b >= 1<<30: return f"{b/(1<<30):.2f} GB"
|
||
if b >= 1<<20: return f"{b/(1<<20):.1f} MB"
|
||
if b >= 1<<10: return f"{b/(1<<10):.1f} KB"
|
||
return f"{b} B"
|
||
|
||
# Build lines and split into ≤4000-char pages
|
||
header = (
|
||
f"📂 **All files in Drive** ({len(files)} files · {human(total_size)} total)\n"
|
||
f"{'━'*22}\n\n"
|
||
)
|
||
|
||
pages = []
|
||
current = header
|
||
for i, f in enumerate(files, 1):
|
||
name = f["Name"]
|
||
disp = name if len(name) <= 45 else name[:42] + "..."
|
||
link = f"https://drive.google.com/file/d/{f['ID']}/view"
|
||
line = f"`{i:3}.` [{disp}]({link})\n 📦 {human(f['Size'])} · 📅 {f['ModTime'][:10]}\n\n"
|
||
|
||
if len(current) + len(line) > 4000:
|
||
pages.append(current)
|
||
current = line
|
||
else:
|
||
current += line
|
||
|
||
if current.strip():
|
||
pages.append(current)
|
||
|
||
# Edit first message, send rest as new replies
|
||
await wait.edit_text(
|
||
pages[0] + (f"\n_Page 1 of {len(pages)}_" if len(pages) > 1 else ""),
|
||
disable_web_page_preview=True,
|
||
)
|
||
for idx, page in enumerate(pages[1:], 2):
|
||
await message.reply_text(
|
||
page + f"\n_Page {idx} of {len(pages)}_",
|
||
disable_web_page_preview=True,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"/listall error: {e}")
|
||
await wait.edit_text(f"❌ Failed to fetch file list\n\n`{e}`")
|
||
|
||
# ── /status ───────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("status"))
|
||
async def cmd_status(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
uptime = time.time() - BOT_START_TIME
|
||
vm = psutil.virtual_memory()
|
||
disk = psutil.disk_usage("/")
|
||
cpu = psutil.cpu_percent(interval=0.5)
|
||
|
||
def bar(pct: float, width: int = 10) -> str:
|
||
filled = int(pct / 100 * width)
|
||
return "█" * filled + "░" * (width - filled)
|
||
|
||
await message.reply_text(
|
||
f"📊 **Bot Status**\n\n"
|
||
f"⏱ **Uptime:** {fmt_dur(uptime)}\n"
|
||
f"📤 **Uploads this session:** {session_uploads}\n\n"
|
||
f"🖥 **CPU:** {cpu:.1f}% `[{bar(cpu)}]`\n\n"
|
||
f"💾 **RAM:**\n"
|
||
f" Used: {fmt_size(vm.used)} / {fmt_size(vm.total)}\n"
|
||
f" Free: {fmt_size(vm.available)}\n"
|
||
f" `[{bar(vm.percent)}]` {vm.percent:.1f}%\n\n"
|
||
f"🗄 **Disk (/):**\n"
|
||
f" Used: {fmt_size(disk.used)} / {fmt_size(disk.total)}\n"
|
||
f" Free: {fmt_size(disk.free)}\n"
|
||
f" `[{bar(disk.percent)}]` {disk.percent:.1f}%\n\n"
|
||
f"📡 **Method:** MTProto + rclone\n"
|
||
f"🔀 **Stream threshold:** ≥ 2 GB"
|
||
)
|
||
|
||
# ── /setquota ────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("setquota"))
|
||
async def cmd_setquota(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
parts = message.command
|
||
if len(parts) < 3:
|
||
return await message.reply_text(
|
||
"Usage: `/setquota @username <size>`\n"
|
||
"Examples: `/setquota @john 5GB` · `/setquota @john 500MB`"
|
||
)
|
||
|
||
username = parts[1].lstrip("@").lower()
|
||
size_bytes = parse_size(parts[2])
|
||
if size_bytes is None:
|
||
return await message.reply_text(
|
||
"❌ Invalid size. Use formats like `5GB`, `500MB`, `2.5GB`, `100KB`."
|
||
)
|
||
|
||
data = _load_quotas()
|
||
data["quotas"][username] = size_bytes
|
||
_save_quotas(data)
|
||
logger.info(f"Quota set: {username} → {fmt_size(size_bytes)}/day")
|
||
await message.reply_text(
|
||
f"✅ Daily quota set for `@{username}`\n\n"
|
||
f"📦 **Limit:** {fmt_size(size_bytes)} / day\n"
|
||
f"🔄 Resets at midnight UTC"
|
||
)
|
||
|
||
# ── /clearquota ───────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("clearquota"))
|
||
async def cmd_clearquota(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
parts = message.command
|
||
if len(parts) < 2:
|
||
return await message.reply_text("Usage: `/clearquota @username`")
|
||
|
||
username = parts[1].lstrip("@").lower()
|
||
data = _load_quotas()
|
||
|
||
if username not in data["quotas"]:
|
||
return await message.reply_text(f"ℹ️ `@{username}` has no quota set.")
|
||
|
||
del data["quotas"][username]
|
||
data["usage"].pop(username, None)
|
||
_save_quotas(data)
|
||
logger.info(f"Quota cleared: {username}")
|
||
await message.reply_text(f"✅ Quota and usage history cleared for `@{username}`.")
|
||
|
||
# ── /quotas ───────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("quotas"))
|
||
async def cmd_quotas(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
data = _load_quotas()
|
||
if not data["quotas"]:
|
||
return await message.reply_text("📋 No quotas configured yet.\n\nUse `/setquota @user <size>` to add one.")
|
||
|
||
lines = []
|
||
for username, limit in sorted(data["quotas"].items()):
|
||
used = get_usage_today(username)
|
||
remaining = max(0, limit - used)
|
||
pct = used / limit * 100 if limit else 0
|
||
bar_n = int(pct / 10)
|
||
bar = "█" * bar_n + "░" * (10 - bar_n)
|
||
lines.append(
|
||
f"👤 `@{username}`\n"
|
||
f" `[{bar}]` {pct:.1f}%\n"
|
||
f" 📦 Used: {fmt_size(used)} / {fmt_size(limit)}\n"
|
||
f" ✅ Left: {fmt_size(remaining)}"
|
||
)
|
||
|
||
await message.reply_text(
|
||
f"📋 **Quota Overview** — {today()} UTC\n\n"
|
||
+ "\n\n".join(lines) +
|
||
"\n\n🔄 Resets at midnight UTC"
|
||
)
|
||
|
||
# ── /myquota ──────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("myquota"))
|
||
async def cmd_myquota(client: Client, message: Message):
|
||
u = (message.from_user.username or "").lower() if message.from_user else ""
|
||
if not u:
|
||
return await message.reply_text("❌ Could not determine your username.")
|
||
|
||
if u in AUTHORIZED_USERS:
|
||
return await message.reply_text("✅ You're an **authorized user** — no quota limits apply to you.")
|
||
|
||
limit = get_quota(u)
|
||
if limit is None:
|
||
return await message.reply_text("🔒 You don't have a quota assigned. Contact an admin.")
|
||
|
||
used = get_usage_today(u)
|
||
remaining = max(0, limit - used)
|
||
pct = used / limit * 100 if limit else 0
|
||
bar_n = int(pct / 10)
|
||
bar = "█" * bar_n + "░" * (10 - bar_n)
|
||
|
||
await message.reply_text(
|
||
f"📊 **Your Quota — {today()} UTC**\n\n"
|
||
f"`[{bar}]` {pct:.1f}%\n\n"
|
||
f"📦 Used: {fmt_size(used)}\n"
|
||
f"✅ Remaining: {fmt_size(remaining)}\n"
|
||
f"🔒 Daily limit: {fmt_size(limit)}\n\n"
|
||
f"🔄 Resets at midnight UTC"
|
||
)
|
||
|
||
# ── /share ───────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("share"))
|
||
async def cmd_share(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
query = " ".join(message.command[1:]).strip()
|
||
if not query:
|
||
return await message.reply_text(
|
||
"Usage: `/share <filename or keyword>`\n\n"
|
||
"Makes the file publicly accessible and returns a shareable link.\n"
|
||
"Tip: partial name works — e.g. `/share Mitm` matches `Mitm-Yt.mp4`"
|
||
)
|
||
|
||
wait = await message.reply_text(f"🔍 Looking for `{query}` in Drive...")
|
||
|
||
# Find the file by exact or partial name match
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "lsjson", RCLONE_REMOTE,
|
||
"--recursive", "--files-only",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(err.decode().strip())
|
||
|
||
files = json.loads(out.decode())
|
||
# Prefer exact match, fall back to partial
|
||
exact = [f for f in files if f["Name"].lower() == query.lower()]
|
||
partial = [f for f in files if query.lower() in f["Name"].lower()]
|
||
matches = exact or partial
|
||
|
||
if not matches:
|
||
return await wait.edit_text(f"❌ No file matching `{query}` found in Drive.")
|
||
|
||
if len(matches) > 1 and not exact:
|
||
names = "\n".join(f"• `{f['Name']}`" for f in matches[:8])
|
||
more = f"\n_…and {len(matches)-8} more_" if len(matches) > 8 else ""
|
||
return await wait.edit_text(
|
||
f"⚠️ **Multiple matches for** `{query}`\n\n{names}{more}\n\n"
|
||
f"Be more specific to share a single file."
|
||
)
|
||
|
||
target = matches[0]
|
||
name = target["Name"]
|
||
path = target["Path"] # includes subfolder, e.g. "Videos/filename.mp4"
|
||
size = target["Size"]
|
||
|
||
except Exception as e:
|
||
logger.error(f"/share lookup error: {e}")
|
||
return await wait.edit_text(f"❌ Drive lookup failed\n\n`{e}`")
|
||
|
||
await wait.edit_text(f"🔗 Making `{name}` public...")
|
||
|
||
# Generate share link via rclone link
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "link", f"{RCLONE_REMOTE}{path}",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(err.decode().strip() or out.decode().strip())
|
||
|
||
share_url = out.decode().strip()
|
||
file_id = share_url.split("id=")[-1] if "id=" in share_url else None
|
||
|
||
except Exception as e:
|
||
logger.error(f"/share rclone link error: {e}")
|
||
return await wait.edit_text(f"❌ Failed to generate share link\n\n`{e}`")
|
||
|
||
def human(b):
|
||
if b >= 1<<30: return f"{b/(1<<30):.2f} GB"
|
||
if b >= 1<<20: return f"{b/(1<<20):.1f} MB"
|
||
if b >= 1<<10: return f"{b/(1<<10):.1f} KB"
|
||
return f"{b} B"
|
||
|
||
dl_link = f"https://drive.google.com/uc?export=download&id={file_id}" if file_id else None
|
||
wget_line = f"`wget -O \"{name}\" \"{dl_link}\"`" if dl_link else ""
|
||
|
||
await wait.edit_text(
|
||
f"✅ **File is now public**\n\n"
|
||
f"📎 **{name}**\n"
|
||
f"📦 {human(size)}\n\n"
|
||
f"🔗 **Share link:**\n`{share_url}`\n\n"
|
||
+ (f"📥 **Direct download:**\n`{dl_link}`\n\n" if dl_link else "")
|
||
+ (f"🖥 **wget:**\n{wget_line}" if wget_line else ""),
|
||
disable_web_page_preview=True,
|
||
)
|
||
|
||
# ── /search ───────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("search"))
|
||
async def cmd_search(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
query = " ".join(message.command[1:]).strip().lower()
|
||
if not query:
|
||
return await message.reply_text("Usage: `/search <filename or keyword>`")
|
||
|
||
wait = await message.reply_text(f"🔍 Searching for `{query}`...")
|
||
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "lsjson", RCLONE_REMOTE,
|
||
"--recursive", "--files-only",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(err.decode().strip())
|
||
|
||
files = json.loads(out.decode())
|
||
hits = [f for f in files if query in f["Name"].lower()]
|
||
hits.sort(key=lambda x: x["ModTime"], reverse=True)
|
||
|
||
if not hits:
|
||
return await wait.edit_text(f"🔍 No files matching `{query}` found in Drive.")
|
||
|
||
lines = []
|
||
for i, f in enumerate(hits[:15], 1):
|
||
size = f["Size"]
|
||
if size >= 1<<30: s = f"{size/(1<<30):.2f} GB"
|
||
elif size >= 1<<20: s = f"{size/(1<<20):.1f} MB"
|
||
elif size >= 1<<10: s = f"{size/(1<<10):.1f} KB"
|
||
else: s = f"{size} B"
|
||
name = f["Name"]
|
||
folder = f["Path"].split("/")[0] if "/" in f["Path"] else "root"
|
||
disp = name if len(name) <= 38 else name[:35] + "..."
|
||
link = f"https://drive.google.com/file/d/{f['ID']}/view"
|
||
lines.append(
|
||
f"`{i:2}.` [{disp}]({link})\n"
|
||
f" 📁 {folder} · 📦 {s} · 📅 {f['ModTime'][:10]}"
|
||
)
|
||
|
||
more = f"\n\n_…and {len(hits)-15} more_" if len(hits) > 15 else ""
|
||
await wait.edit_text(
|
||
f"🔍 **{len(hits)} result(s) for** `{query}`\n\n" + "\n\n".join(lines) + more,
|
||
disable_web_page_preview=True,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"/search error: {e}")
|
||
await wait.edit_text(f"❌ Search failed\n\n`{e}`")
|
||
|
||
# ── /adduser & /removeuser ────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("adduser"))
|
||
async def cmd_adduser(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
parts = message.command
|
||
if len(parts) < 2:
|
||
return await message.reply_text("Usage: `/adduser <username>`")
|
||
|
||
username = parts[1].lstrip("@").lower()
|
||
if username in AUTHORIZED_USERS:
|
||
return await message.reply_text(f"ℹ️ `@{username}` is already authorized.")
|
||
|
||
AUTHORIZED_USERS.add(username)
|
||
_save_users(AUTHORIZED_USERS)
|
||
logger.info(f"Added authorized user: {username}")
|
||
await message.reply_text(f"✅ `@{username}` added to authorized users.")
|
||
|
||
@app.on_message(filters.command("removeuser"))
|
||
async def cmd_removeuser(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
parts = message.command
|
||
if len(parts) < 2:
|
||
return await message.reply_text("Usage: `/removeuser <username>`")
|
||
|
||
username = parts[1].lstrip("@").lower()
|
||
if username not in AUTHORIZED_USERS:
|
||
return await message.reply_text(f"ℹ️ `@{username}` is not in the authorized list.")
|
||
|
||
AUTHORIZED_USERS.discard(username)
|
||
_save_users(AUTHORIZED_USERS)
|
||
logger.info(f"Removed authorized user: {username}")
|
||
await message.reply_text(f"✅ `@{username}` removed from authorized users.")
|
||
|
||
# ── /logs ─────────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("logs"))
|
||
async def cmd_logs(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
try:
|
||
lines = LOG_FILE.read_text(errors="replace").splitlines()
|
||
tail = "\n".join(lines[-50:])
|
||
# Send as file if content is large, otherwise inline
|
||
if len(tail) > 3800:
|
||
with tempfile.NamedTemporaryFile("w", suffix=".log",
|
||
delete=False, prefix="botlog_") as f:
|
||
f.write(tail)
|
||
tmp = f.name
|
||
await message.reply_document(tmp, caption="📄 Last 50 log lines")
|
||
os.remove(tmp)
|
||
else:
|
||
await message.reply_text(f"📄 **Last 50 log lines**\n\n```\n{tail}\n```")
|
||
except Exception as e:
|
||
await message.reply_text(f"❌ Could not read log\n\n`{e}`")
|
||
|
||
# ── /restart ──────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("restart"))
|
||
async def cmd_restart(client: Client, message: Message):
|
||
if not is_authorized(message):
|
||
return await deny(message)
|
||
|
||
await message.reply_text("♻️ Restarting bot... brb.")
|
||
logger.info("Restart requested via /restart")
|
||
|
||
async def do_restart():
|
||
await asyncio.sleep(1)
|
||
await client.stop()
|
||
os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
asyncio.create_task(do_restart())
|
||
|
||
# ── /help ─────────────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.command("help"))
|
||
async def cmd_help(client: Client, message: Message):
|
||
auth_note = (
|
||
"\n✅ You are an **authorized user**."
|
||
if is_authorized(message) else
|
||
"\n🔒 Commands require authorization to run."
|
||
)
|
||
await message.reply_text(
|
||
"**GDrive Bot — Commands**\n\n"
|
||
"━━━━━━━━━━━━━━━━━━━━\n"
|
||
"📂 **File & Transfer**\n"
|
||
"`/list` — Show the 10 most recent files in the Drive folder\n"
|
||
"`/listall` — List **every** file in the Drive folder (paginated)\n"
|
||
"`/search <query>` — Search files by name or keyword\n"
|
||
"`/share <filename>` — Make a file public and get a shareable link\n\n"
|
||
"📊 **Info**\n"
|
||
"`/status` — Bot uptime, RAM, CPU, disk usage\n"
|
||
"`/logs` — Tail the last 50 lines of the bot log\n\n"
|
||
"👤 **User Management**\n"
|
||
"`/adduser <@username>` — Grant full (unlimited) upload access\n"
|
||
"`/removeuser <@username>` — Revoke full access\n\n"
|
||
"📊 **Quotas** _(for non-auth users)_\n"
|
||
"`/setquota @user <size>` — Set a daily upload quota (e.g. `5GB`, `500MB`)\n"
|
||
"`/clearquota @user` — Remove quota and reset usage\n"
|
||
"`/quotas` — Overview of all quotas and today's usage\n"
|
||
"`/myquota` — Check your own quota (any user)\n\n"
|
||
"⚙️ **System**\n"
|
||
"`/restart` — Gracefully restart the bot\n"
|
||
"`/help` — Show this message\n"
|
||
"`/start` — Show upload instructions\n\n"
|
||
"━━━━━━━━━━━━━━━━━━━━\n"
|
||
"**How to upload:**\n"
|
||
"📁 Send any file · 🌐 Paste a URL · ✏️ Add caption to rename\n\n"
|
||
"Made by @Coaxys"
|
||
+ auth_note
|
||
)
|
||
|
||
# ── URL handler ───────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.text & ~filters.command(
|
||
["start", "help", "list", "listall", "status", "search", "share",
|
||
"adduser", "removeuser", "logs", "restart",
|
||
"setquota", "clearquota", "quotas", "myquota"]
|
||
))
|
||
async def handle_url(client: Client, message: Message):
|
||
allowed, reason = has_upload_access(message)
|
||
if not allowed:
|
||
if reason == "quota_exceeded":
|
||
u = (message.from_user.username or "").lower()
|
||
limit = get_quota(u)
|
||
await message.reply_text(
|
||
f"🚫 **Daily quota exceeded**\n\n"
|
||
f"You've used your full {fmt_size(limit)} for today.\n"
|
||
f"🔄 Resets at midnight UTC."
|
||
)
|
||
return
|
||
match = URL_RE.search(message.text or "")
|
||
if not match:
|
||
return
|
||
|
||
url = match.group(0).rstrip(".,;)")
|
||
file_name = filename_from_url(url)
|
||
chunk = "512M" # URL downloads: always use max chunks
|
||
total_start = time.time()
|
||
|
||
status = await message.reply_text(
|
||
f"🌐 **URL received**\n\n"
|
||
f"`{url}`\n\n"
|
||
f"📎 Saving as: **{file_name}**\n"
|
||
f"🧩 Chunk: {chunk}\n\n"
|
||
f"⏳ Starting rclone stream (URL → RAM → Drive)..."
|
||
)
|
||
|
||
folder = classify_to_folder(file_name)
|
||
rclone_cmd = [
|
||
"rclone", "copyurl", url,
|
||
f"{RCLONE_REMOTE}{folder}/{file_name}",
|
||
"--drive-chunk-size", chunk,
|
||
"--buffer-size", "512M",
|
||
]
|
||
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
*rclone_cmd,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
except Exception as e:
|
||
await status.edit_text(f"❌ **rclone failed to start**\n\n`{e}`")
|
||
return
|
||
|
||
# Pulse-update while rclone works (copyurl gives no byte-level progress)
|
||
start_ts = time.time()
|
||
dots = ["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
|
||
dot_i = 0
|
||
done_event = asyncio.Event()
|
||
|
||
async def pulse():
|
||
nonlocal dot_i
|
||
while not done_event.is_set():
|
||
elapsed = time.time() - start_ts
|
||
try:
|
||
await status.edit_text(
|
||
f"🌐 **Downloading & uploading...**\n\n"
|
||
f"`{url}`\n\n"
|
||
f"📎 **{file_name}**\n"
|
||
f"{dots[dot_i % len(dots)]} Elapsed: {fmt_dur(elapsed)}\n\n"
|
||
f"🧩 Chunk: {chunk} · 💾 Disk: **0 B**\n"
|
||
f"🖥 RAM: {fmt_ram()}"
|
||
)
|
||
except Exception:
|
||
pass
|
||
dot_i += 1
|
||
await asyncio.sleep(4)
|
||
|
||
pulse_task = asyncio.create_task(pulse())
|
||
out, err = await proc.communicate()
|
||
done_event.set()
|
||
pulse_task.cancel()
|
||
try:
|
||
await pulse_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
total_elapsed = time.time() - total_start
|
||
|
||
if proc.returncode != 0:
|
||
error = err.decode().strip() or out.decode().strip()
|
||
logger.error(f"rclone copyurl failed: {error}")
|
||
await status.edit_text(
|
||
f"❌ **URL download failed**\n\n"
|
||
f"`{url}`\n\n"
|
||
f"💬 `{error}`"
|
||
)
|
||
return
|
||
|
||
global session_uploads
|
||
session_uploads += 1
|
||
await status.edit_text(
|
||
f"✅ **Done!**\n\n"
|
||
f"📎 **{file_name}**\n"
|
||
f"🌐 Source: `{url[:60]}{'...' if len(url)>60 else ''}`\n\n"
|
||
f"🕐 Total: {fmt_dur(total_elapsed)} · 💾 Disk: **0 B**\n"
|
||
f"📁 → **Google Drive / {folder}**"
|
||
)
|
||
|
||
await send_links(message, file_name)
|
||
|
||
# ── File handler ──────────────────────────────────────────────────────────────
|
||
|
||
@app.on_message(filters.media)
|
||
async def handle_file(client: Client, message: Message):
|
||
allowed, reason = has_upload_access(message)
|
||
if not allowed:
|
||
if reason == "quota_exceeded":
|
||
u = (message.from_user.username or "").lower()
|
||
limit = get_quota(u)
|
||
await message.reply_text(
|
||
f"🚫 **Daily quota exceeded**\n\n"
|
||
f"You've used your full {fmt_size(limit)} for today.\n"
|
||
f"🔄 Resets at midnight UTC."
|
||
)
|
||
return
|
||
media_type, media = get_media(message)
|
||
if not media:
|
||
return
|
||
|
||
caption = (message.caption or "").strip() or None
|
||
file_name = resolve_filename(media_type, media, caption)
|
||
raw_size = getattr(media, "file_size", 0) or 0
|
||
size_str = fmt_size(raw_size) if raw_size else "unknown"
|
||
mime = getattr(media, "mime_type", None) or media_type
|
||
dc_id = getattr(media, "dc_id", None)
|
||
dc_note = f"DC{dc_id}" if dc_id else "DC?"
|
||
chunk_label = choose_chunk(raw_size)
|
||
use_stream = raw_size >= STREAM_THRESHOLD
|
||
method = "🌊 RAM stream" if use_stream else "💾 Download → Upload"
|
||
total_start = time.time()
|
||
|
||
username = (message.from_user.username or "").lower() if message.from_user else ""
|
||
rename_note = f"\n✏️ Renamed to: **{file_name}**" if caption else ""
|
||
|
||
# Quota pre-check: reject if file would exceed remaining quota
|
||
if username not in AUTHORIZED_USERS and raw_size:
|
||
rem = quota_remaining(username)
|
||
if rem is not None and raw_size > rem:
|
||
return await message.reply_text(
|
||
f"🚫 **Not enough quota**\n\n"
|
||
f"📦 File size: {fmt_size(raw_size)}\n"
|
||
f"✅ Remaining: {fmt_size(rem)}\n\n"
|
||
f"🔄 Quota resets at midnight UTC."
|
||
)
|
||
|
||
status = await message.reply_text(
|
||
f"📥 **Received**\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {size_str} · 🗂 `{mime}` · 📡 {dc_note}{rename_note}\n"
|
||
f"🔧 Method: {method} · 🧩 Chunk: {chunk_label}\n\n"
|
||
f"⏳ Starting..."
|
||
)
|
||
|
||
folder = classify_to_folder(file_name, mime)
|
||
if use_stream:
|
||
await _stream_to_drive(client, message, status, file_name, raw_size,
|
||
size_str, mime, dc_note, chunk_label, total_start,
|
||
username=username, folder=folder)
|
||
else:
|
||
await _download_then_upload(client, message, status, file_name, raw_size,
|
||
size_str, mime, dc_note, chunk_label, total_start,
|
||
username=username, folder=folder)
|
||
|
||
# ── Method A: download to disk → rclone copy (< 2 GB) ───────────────────────
|
||
|
||
async def _download_then_upload(
|
||
client, message, status, file_name, raw_size,
|
||
size_str, mime, dc_note, chunk_label, total_start, username="", folder="Other"
|
||
):
|
||
temp_dir = tempfile.mkdtemp()
|
||
local_path = os.path.join(temp_dir, file_name)
|
||
|
||
dl_start = time.time()
|
||
last_edit = 0.0
|
||
peak_dl = 0.0
|
||
speed_win = collections.deque(maxlen=6)
|
||
prev_curr = 0
|
||
|
||
async def progress(current: int, total: int):
|
||
nonlocal last_edit, peak_dl, prev_curr
|
||
now = time.time()
|
||
if now - last_edit < 3:
|
||
return
|
||
last_edit = now
|
||
elapsed = now - dl_start
|
||
avg = current / elapsed if elapsed > 0 else 0
|
||
speed_win.append(avg)
|
||
peak_dl = max(peak_dl, avg)
|
||
smooth = sum(speed_win) / len(speed_win)
|
||
pct = current / total * 100 if total else 0
|
||
eta = (total - current) / smooth if smooth > 0 else 0
|
||
bar = "█" * int(pct / 5) + "░" * (20 - int(pct / 5))
|
||
try:
|
||
await status.edit_text(
|
||
f"⬇️ **Downloading** · {dc_note}\n\n"
|
||
f"**{file_name}**\n"
|
||
f"`[{bar}]` **{pct:.1f}%**\n"
|
||
f"📦 {fmt_size(current)} / {fmt_size(total)}\n\n"
|
||
f"⚡ Speed: {spd(smooth)}\n"
|
||
f"🏆 Peak: {spd(peak_dl)}\n"
|
||
f"⏱ ETA: {fmt_dur(eta)} · Elapsed: {fmt_dur(elapsed)}\n\n"
|
||
f"🧩 Chunk: {chunk_label} · 🖥 RAM: {fmt_ram()}"
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
MAX_RETRIES = 3
|
||
RETRY_DELAYS = [5, 15, 30]
|
||
last_error = None
|
||
|
||
for attempt in range(1, MAX_RETRIES + 1):
|
||
try:
|
||
# Reset progress state for each attempt
|
||
dl_start = time.time()
|
||
last_edit = 0.0
|
||
peak_dl = 0.0
|
||
speed_win.clear()
|
||
|
||
if attempt > 1:
|
||
try:
|
||
await status.edit_text(
|
||
f"🔄 **Retrying download** ({attempt}/{MAX_RETRIES}) · {dc_note}\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {size_str}\n\n"
|
||
f"⚠️ Previous error: `{last_error}`\n"
|
||
f"⏳ Starting attempt {attempt}..."
|
||
)
|
||
except Exception:
|
||
pass
|
||
# Clean up partial file from last attempt
|
||
try: os.remove(local_path)
|
||
except Exception: pass
|
||
|
||
downloaded = await client.download_media(message, file_name=local_path, progress=progress)
|
||
if not downloaded:
|
||
raise RuntimeError("download_media returned None")
|
||
break # success — exit retry loop
|
||
|
||
except Exception as e:
|
||
last_error = str(e)
|
||
logger.warning(f"Download attempt {attempt}/{MAX_RETRIES} failed: {e}")
|
||
if attempt == MAX_RETRIES:
|
||
logger.error(f"All {MAX_RETRIES} download attempts failed for {file_name}")
|
||
await status.edit_text(
|
||
f"❌ **Download failed** after {MAX_RETRIES} attempts\n\n"
|
||
f"**{file_name}**\n\n"
|
||
f"💬 Last error: `{last_error}`"
|
||
)
|
||
try: os.remove(local_path)
|
||
except Exception: pass
|
||
try: os.rmdir(temp_dir)
|
||
except Exception: pass
|
||
return
|
||
delay = RETRY_DELAYS[attempt - 1]
|
||
try:
|
||
await status.edit_text(
|
||
f"⚠️ **Connection dropped** (attempt {attempt}/{MAX_RETRIES}) · {dc_note}\n\n"
|
||
f"**{file_name}**\n\n"
|
||
f"💬 `{last_error}`\n\n"
|
||
f"⏳ Retrying in {delay}s..."
|
||
)
|
||
except Exception:
|
||
pass
|
||
await asyncio.sleep(delay)
|
||
|
||
dl_elapsed = time.time() - dl_start
|
||
avg_dl = raw_size / dl_elapsed if dl_elapsed > 0 and raw_size else 0
|
||
|
||
await status.edit_text(
|
||
f"☁️ **Uploading to Drive...**\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {size_str} downloaded in {fmt_dur(dl_elapsed)} ({spd(avg_dl)})\n\n"
|
||
f"⏳ rclone copying to Google Drive..."
|
||
)
|
||
|
||
up_start = time.time()
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
"rclone", "copy", local_path, f"{RCLONE_REMOTE}{folder}/",
|
||
"--drive-chunk-size", chunk_label,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
out, err = await proc.communicate()
|
||
up_elapsed = time.time() - up_start
|
||
total_elapsed = time.time() - total_start
|
||
avg_up = raw_size / up_elapsed if up_elapsed > 0 and raw_size else 0
|
||
|
||
if proc.returncode != 0:
|
||
error = err.decode().strip() or out.decode().strip()
|
||
logger.error(f"rclone copy failed: {error}")
|
||
await status.edit_text(
|
||
f"❌ **Upload failed**\n\n**{file_name}**\n\n💬 `{error}`\n\n"
|
||
f"💾 Local copy at:\n`{local_path}`"
|
||
)
|
||
return
|
||
|
||
global session_uploads
|
||
session_uploads += 1
|
||
if username and username not in AUTHORIZED_USERS:
|
||
record_upload(username, raw_size or up_elapsed * 1) # track actual bytes
|
||
quota_note = ""
|
||
if username and username not in AUTHORIZED_USERS:
|
||
rem = quota_remaining(username)
|
||
if rem is not None:
|
||
quota_note = f"\n📊 Quota left today: {fmt_size(rem)}"
|
||
await status.edit_text(
|
||
f"✅ **Done!**\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {size_str} · 🗂 `{mime}`\n\n"
|
||
f"📊 **Stats**\n"
|
||
f"⬇️ Download: {fmt_dur(dl_elapsed)} @ {spd(avg_dl)}\n"
|
||
f"☁️ Upload: {fmt_dur(up_elapsed)} @ {spd(avg_up)}\n"
|
||
f"🕐 Total: {fmt_dur(total_elapsed)}\n\n"
|
||
f"🧩 Chunk: {chunk_label} · 📡 {dc_note}\n"
|
||
f"📁 → **Google Drive / {folder}**{quota_note}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"rclone copy error: {e}")
|
||
await status.edit_text(f"❌ **Upload error**\n\n**{file_name}**\n\n💬 `{e}`")
|
||
return
|
||
finally:
|
||
try: os.remove(local_path)
|
||
except Exception: pass
|
||
try: os.rmdir(temp_dir)
|
||
except Exception: pass
|
||
|
||
await send_links(message, file_name)
|
||
|
||
# ── Method B: pure RAM stream via rclone rcat (≥ 2 GB) ──────────────────────
|
||
|
||
async def _stream_to_drive(
|
||
client, message, status, file_name, raw_size,
|
||
size_str, mime, dc_note, chunk_label, total_start, username="", folder="Other"
|
||
):
|
||
rclone_cmd = [
|
||
"rclone", "rcat", f"{RCLONE_REMOTE}{folder}/{file_name}",
|
||
"--drive-chunk-size", chunk_label,
|
||
]
|
||
if raw_size:
|
||
rclone_cmd += ["--size", str(raw_size)]
|
||
|
||
MAX_RETRIES = 3
|
||
RETRY_DELAYS = [5, 15, 30]
|
||
last_error = None
|
||
bytes_sent = 0
|
||
peak_spd = 0.0
|
||
stream_start = time.time()
|
||
|
||
for attempt in range(1, MAX_RETRIES + 1):
|
||
if attempt > 1:
|
||
delay = RETRY_DELAYS[attempt - 2]
|
||
try:
|
||
await status.edit_text(
|
||
f"⚠️ **Stream dropped** (attempt {attempt - 1}/{MAX_RETRIES}) · {dc_note}\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {fmt_size(bytes_sent)} received so far\n\n"
|
||
f"💬 `{last_error}`\n\n"
|
||
f"⏳ Retrying from 0% in {delay}s... (stream must restart)"
|
||
)
|
||
except Exception:
|
||
pass
|
||
await asyncio.sleep(delay)
|
||
# Reset counters — stream restarts from beginning
|
||
bytes_sent = 0
|
||
stream_start = time.time()
|
||
|
||
# Launch a fresh rclone process for each attempt
|
||
try:
|
||
rclone = await asyncio.create_subprocess_exec(
|
||
*rclone_cmd,
|
||
stdin=asyncio.subprocess.PIPE,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"rclone launch failed: {e}")
|
||
await status.edit_text(f"❌ **rclone failed to start**\n\n`{e}`")
|
||
return
|
||
|
||
last_edit = 0.0
|
||
speed_win = collections.deque(maxlen=6)
|
||
win_start = stream_start
|
||
win_bytes = 0
|
||
stream_err = None
|
||
|
||
try:
|
||
async for chunk in client.stream_media(message):
|
||
rclone.stdin.write(chunk)
|
||
await rclone.stdin.drain()
|
||
bytes_sent += len(chunk)
|
||
win_bytes += len(chunk)
|
||
|
||
now = time.time()
|
||
win_age = now - win_start
|
||
if win_age >= 3:
|
||
inst = win_bytes / win_age
|
||
speed_win.append(inst)
|
||
peak_spd = max(peak_spd, inst)
|
||
win_bytes = 0
|
||
win_start = now
|
||
|
||
avg = sum(speed_win) / len(speed_win) if speed_win else 0
|
||
|
||
if now - last_edit >= 3:
|
||
last_edit = now
|
||
elapsed = now - stream_start
|
||
pct = bytes_sent / raw_size * 100 if raw_size else 0
|
||
eta = (raw_size - bytes_sent) / avg if avg > 0 and raw_size else 0
|
||
bar = "█" * int(pct / 5) + "░" * (20 - int(pct / 5))
|
||
retry_note = f" · 🔁 attempt {attempt}/{MAX_RETRIES}" if attempt > 1 else ""
|
||
try:
|
||
await status.edit_text(
|
||
f"🌊 **Streaming** · {dc_note}{retry_note}\n\n"
|
||
f"**{file_name}**\n"
|
||
f"`[{bar}]` **{pct:.1f}%**\n"
|
||
f"📦 {fmt_size(bytes_sent)} / {size_str}\n\n"
|
||
f"⚡ Speed: {spd(avg)}\n"
|
||
f"🏆 Peak: {spd(peak_spd)}\n"
|
||
f"⏱ ETA: {fmt_dur(eta)} · Elapsed: {fmt_dur(elapsed)}\n\n"
|
||
f"🧩 Chunk: {chunk_label} · 💾 Disk: **0 B**\n"
|
||
f"🖥 RAM: {fmt_ram()}"
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
rclone.stdin.close()
|
||
|
||
except Exception as e:
|
||
stream_err = e
|
||
last_error = str(e)
|
||
logger.warning(f"Stream attempt {attempt}/{MAX_RETRIES} failed: {e}")
|
||
try: rclone.stdin.close()
|
||
except Exception: pass
|
||
try: rclone.kill()
|
||
except Exception: pass
|
||
|
||
if stream_err is None:
|
||
break # stream completed cleanly
|
||
|
||
if attempt == MAX_RETRIES:
|
||
await status.edit_text(
|
||
f"❌ **Stream failed** after {MAX_RETRIES} attempts\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {fmt_size(bytes_sent)} received before final failure\n\n"
|
||
f"💬 `{last_error}`"
|
||
)
|
||
return
|
||
|
||
stream_elapsed = time.time() - stream_start
|
||
try:
|
||
await status.edit_text(
|
||
f"☁️ **Flushing to Drive API...**\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {fmt_size(bytes_sent)} streamed in {fmt_dur(stream_elapsed)}\n"
|
||
f"⚡ Avg: {spd(bytes_sent / stream_elapsed if stream_elapsed else 0)}"
|
||
f" · 🏆 Peak: {spd(peak_spd)}\n\n"
|
||
f"⏳ rclone flushing last {chunk_label} chunk..."
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
out, err = await rclone.communicate()
|
||
total_elapsed = time.time() - total_start
|
||
|
||
if rclone.returncode != 0:
|
||
error = err.decode().strip() or out.decode().strip()
|
||
logger.error(f"rclone rcat exit {rclone.returncode}: {error}")
|
||
await status.edit_text(
|
||
f"❌ **Drive upload failed**\n\n**{file_name}**\n"
|
||
f"📦 {fmt_size(bytes_sent)} streamed — rclone exited {rclone.returncode}\n\n"
|
||
f"💬 `{error}`"
|
||
)
|
||
return
|
||
|
||
global session_uploads
|
||
session_uploads += 1
|
||
if username and username not in AUTHORIZED_USERS:
|
||
record_upload(username, bytes_sent)
|
||
quota_note = ""
|
||
if username and username not in AUTHORIZED_USERS:
|
||
rem = quota_remaining(username)
|
||
if rem is not None:
|
||
quota_note = f"\n📊 Quota left today: {fmt_size(rem)}"
|
||
avg_final = bytes_sent / stream_elapsed if stream_elapsed > 0 else 0
|
||
await status.edit_text(
|
||
f"✅ **Done!**\n\n"
|
||
f"**{file_name}**\n"
|
||
f"📦 {fmt_size(bytes_sent)} · 🗂 `{mime}`\n\n"
|
||
f"📊 **Stats**\n"
|
||
f"⚡ Avg speed: {spd(avg_final)}\n"
|
||
f"🏆 Peak speed: {spd(peak_spd)}\n"
|
||
f"⏱ Stream: {fmt_dur(stream_elapsed)}\n"
|
||
f"🕐 Total: {fmt_dur(total_elapsed)}\n\n"
|
||
f"🧩 Chunk: {chunk_label} · 📡 {dc_note} · 💾 Disk: **0 B**\n"
|
||
f"📁 → **Google Drive / {folder}**{quota_note}"
|
||
)
|
||
|
||
await send_links(message, file_name)
|
||
|
||
# ── Entry ─────────────────────────────────────────────────────────────────────
|
||
|
||
if __name__ == "__main__":
|
||
if not API_ID or not API_HASH or not BOT_TOKEN:
|
||
logger.error("API_ID, API_HASH, BOT_TOKEN must be set in .env")
|
||
raise SystemExit(1)
|
||
logger.info("Bot starting — manual < 2 GB · stream ≥ 2 GB · URL download · caption rename")
|
||
app.run()
|