import os import pymysql from dotenv import load_dotenv load_dotenv() def get_connection(): return pymysql.connect( host=os.getenv("DB_HOST", "localhost"), port=int(os.getenv("DB_PORT", 3306)), user=os.getenv("DB_USER", "msmharbi2"), password=os.getenv("DB_PASSWORD", "8vGi)aiRm72WtH06"), database=os.getenv("DB_NAME", "my_media"), charset="utf8mb4", ) # ── Duplicate check ──────────────────────────────────────────────────────────── def check_duplicate(url: str, telegram_user_id: int = None): """Returns the existing row if URL already saved by this user, else None.""" conn = get_connection() try: with conn.cursor(pymysql.cursors.DictCursor) as cur: if telegram_user_id: cur.execute( "SELECT * FROM links WHERE url = %s AND telegram_user_id = %s", (url, telegram_user_id) ) else: cur.execute("SELECT * FROM links WHERE url = %s", (url,)) return cur.fetchone() finally: conn.close() # ── Save content ─────────────────────────────────────────────────────────────── def save_link( url: str, link_type: str, description: str = None, telegram_user_id: int = None, file_size_mb: float = 0.0, ) -> int: """Insert into links table. Returns new row id.""" conn = get_connection() try: with conn.cursor() as cur: cur.execute( "INSERT INTO links (url, type, description, telegram_user_id, file_size_mb) " "VALUES (%s, %s, %s, %s, %s)", (url, link_type, description, telegram_user_id, file_size_mb), ) conn.commit() return cur.lastrowid finally: conn.close() def save_media(link_id: int, file_path: str) -> int: """Insert into media table. Returns new row id.""" conn = get_connection() try: with conn.cursor() as cur: cur.execute( "INSERT INTO media (link_id, file_path) VALUES (%s, %s)", (link_id, file_path), ) conn.commit() return cur.lastrowid finally: conn.close() # ── Telegram user registration ───────────────────────────────────────────────── def ensure_telegram_user(tg_id: int, username: str | None, first_name: str) -> bool: """ Register Telegram user if not exists. Assigns Free plan automatically. Returns True if newly registered, False if already existed. """ conn = get_connection() try: with conn.cursor(pymysql.cursors.DictCursor) as cur: cur.execute( "INSERT IGNORE INTO telegram_users (telegram_user_id, username, first_name) " "VALUES (%s, %s, %s)", (tg_id, username, first_name or ""), ) is_new = cur.rowcount > 0 if is_new: # Get the free plan (price=0, lowest id) cur.execute( "SELECT id, storage_limit_mb, file_limit FROM subscription_plans " "WHERE price = 0 ORDER BY id ASC LIMIT 1" ) plan = cur.fetchone() plan_id = plan["id"] if plan else 1 storage_mb = plan["storage_limit_mb"] if plan else 512 file_limit = plan["file_limit"] if plan else 50 # Create active subscription (no expiry = permanent free) cur.execute( "INSERT INTO telegram_subscriptions " "(telegram_user_id, plan_id, storage_limit_mb, file_limit, expires_at, status) " "VALUES (%s, %s, %s, %s, NULL, 'active')", (tg_id, plan_id, storage_mb, file_limit), ) # Initialize storage stats cur.execute( "INSERT IGNORE INTO user_storage_stats (telegram_user_id) VALUES (%s)", (tg_id,), ) conn.commit() return is_new finally: conn.close() # ── Quota ────────────────────────────────────────────────────────────────────── def get_user_quota(tg_id: int) -> dict: """ Returns current quota and usage for a Telegram user. Keys: plan_name, storage_limit_mb, file_limit, file_count, storage_used_mb, expires_at, sub_status """ conn = get_connection() try: with conn.cursor(pymysql.cursors.DictCursor) as cur: cur.execute( """ SELECT sp.name AS plan_name, ts.storage_limit_mb, ts.file_limit, ts.expires_at, ts.status AS sub_status, COALESCE(uss.file_count, 0) AS file_count, COALESCE(uss.storage_used_mb, 0) AS storage_used_mb FROM telegram_subscriptions ts JOIN subscription_plans sp ON sp.id = ts.plan_id LEFT JOIN user_storage_stats uss ON uss.telegram_user_id = ts.telegram_user_id WHERE ts.telegram_user_id = %s ORDER BY ts.id DESC LIMIT 1 """, (tg_id,), ) row = cur.fetchone() if not row: return { "plan_name": "Free", "storage_limit_mb": 512, "file_limit": 50, "file_count": 0, "storage_used_mb": 0.0, "expires_at": None, "sub_status": "none", } row["storage_used_mb"] = float(row["storage_used_mb"]) return row finally: conn.close() def can_upload(tg_id: int, file_size_mb: float = 0.0) -> tuple[bool, str]: """ Check if user can upload. Returns (True, '') or (False, 'reason message'). """ q = get_user_quota(tg_id) # Check subscription status if q["sub_status"] == "expired": return False, "⚠️ انتهى اشتراكك. تواصل مع المسؤول للتجديد." # Check storage if q["storage_used_mb"] + file_size_mb > q["storage_limit_mb"]: used = round(q["storage_used_mb"], 1) limit = q["storage_limit_mb"] size = round(file_size_mb, 1) msg = ( f"❌ لا توجد مساحة كافية!\n\n" f"💾 التخزين: {used}MB / {limit}MB\n" ) if file_size_mb > 0: msg += f"📦 حجم الملف: {size}MB\n" msg += "\nتواصل مع المسؤول لزيادة المساحة." return False, msg # Check file count if q["file_limit"] is not None and q["file_count"] >= q["file_limit"]: return False, ( f"❌ وصلت إلى حد الملفات!\n\n" f"📁 الملفات: {q['file_count']} / {q['file_limit']}\n\n" f"تواصل مع المسؤول للترقية." ) return True, "" # ── Usage tracking ───────────────────────────────────────────────────────────── def record_upload(tg_id: int, file_size_mb: float) -> None: """Increment file count and storage used after a successful upload.""" conn = get_connection() try: with conn.cursor() as cur: cur.execute( "INSERT INTO user_storage_stats " "(telegram_user_id, file_count, storage_used_mb) " "VALUES (%s, 1, %s) " "ON DUPLICATE KEY UPDATE " " file_count = file_count + 1, " " storage_used_mb = storage_used_mb + %s", (tg_id, file_size_mb, file_size_mb), ) conn.commit() finally: conn.close() def record_delete(tg_id: int, file_size_mb: float) -> None: """Decrement file count and storage used after a deletion.""" conn = get_connection() try: with conn.cursor() as cur: cur.execute( "UPDATE user_storage_stats SET " " file_count = GREATEST(0, file_count - 1), " " storage_used_mb = GREATEST(0, storage_used_mb - %s) " "WHERE telegram_user_id = %s", (file_size_mb, tg_id), ) conn.commit() finally: conn.close()