default/telegram-agy-bot/auth_manager.py

import os
import re
import json
import time
import shutil
import asyncio
import logging
from pathlib import Path
from typing import Optional, Dict, Any, Tuple
from datetime import datetime

logger = logging.getLogger("AGYAuthManager")

PROFILES_BASE_DIR = Path("/root/.gemini_profiles")
MASTER_TOKEN_PATH = Path("/root/.gemini/antigravity-cli/antigravity-oauth-token")
MASTER_SETTINGS_PATH = Path("/root/.gemini/antigravity-cli/settings.json")

def escape_pre(text: str) -> str:
    """Escapes HTML entities for pre block."""
    if not text:
        return ""
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

class AuthManager:
    """
    Manages multi-account authentication, isolated user profiles,
    and OAuth login flows for Antigravity (AGY) CLI.
    """

    def __init__(self, base_dir: Path = PROFILES_BASE_DIR):
        self.base_dir = base_dir
        self.base_dir.mkdir(parents=True, exist_ok=True)
        # Stores active login subprocesses: {user_id: {"proc": Process, "url": str, "started_at": float}}
        self.active_logins: Dict[int, Dict[str, Any]] = {}

    def get_profile_dir(self, user_id: int) -> Path:
        """Returns the isolated home directory for a given user."""
        return self.base_dir / str(user_id)

    def get_config_dir(self, user_id: int) -> Path:
        """Returns the .gemini/antigravity-cli directory inside user profile."""
        return self.get_profile_dir(user_id) / ".gemini" / "antigravity-cli"

    def get_token_path(self, user_id: int) -> Path:
        """Returns the path to the user's antigravity-oauth-token."""
        return self.get_config_dir(user_id) / "antigravity-oauth-token"

    def has_custom_account(self, user_id: int) -> bool:
        """Checks if a user has a valid personal AGY account token configured."""
        token_file = self.get_token_path(user_id)
        if token_file.exists() and token_file.is_file() and token_file.stat().st_size > 10:
            return True
        return False

    def get_user_home(self, user_id: int) -> str:
        """Returns the isolated HOME directory path to use for this user."""
        return str(self.get_profile_dir(user_id))

    def get_user_env(self, user_id: int) -> Dict[str, str]:
        """Returns environment variables dictionary with the proper HOME set for the user."""
        env = os.environ.copy()
        env["HOME"] = self.get_user_home(user_id)
        return env

    def get_account_status(self, user_id: int) -> Dict[str, Any]:
        """Returns detailed status of the user's active account."""
        has_custom = self.has_custom_account(user_id)
        token_path = self.get_token_path(user_id)

        modified_time_str = "تنظیم نشده"
        if has_custom and token_path.exists():
            try:
                mtime = token_path.stat().st_mtime
                modified_time_str = datetime.fromtimestamp(mtime).strftime("%Y-%m-%d %H:%M:%S")
            except Exception:
                pass

        return {
            "user_id": user_id,
            "has_custom_account": has_custom,
            "is_default_server": False,
            "account_type": "حساب اختصاصی (Custom Account)" if has_custom else "حساب متصل نیست (Disconnected)",
            "token_path": str(token_path) if has_custom else "تنظیم نشده",
            "last_modified": modified_time_str,
            "is_login_in_progress": user_id in self.active_logins,
        }

    async def start_oauth_login(self, user_id: int, timeout: int = 25) -> Tuple[bool, str, Optional[str]]:
        """
        Starts headless OAuth login flow for a user via a pseudo-terminal (PTY)
        so that AGY CLI detects a controlling terminal and outputs the Google Auth URL.
        """
        import pty
        import select

        # Cancel any previous login session for this user
        self.cancel_oauth_login(user_id)

        config_dir = self.get_config_dir(user_id)
        config_dir.mkdir(parents=True, exist_ok=True)

        token_file = self.get_token_path(user_id)
        backup_token = None
        if token_file.exists():
            backup_token = token_file.read_text(encoding="utf-8", errors="ignore")
            try:
                token_file.unlink()
            except Exception:
                pass

        env = os.environ.copy()
        env["HOME"] = str(self.get_profile_dir(user_id))

        try:
            master_fd, slave_fd = pty.openpty()

            proc = await asyncio.create_subprocess_exec(
                "agy", "--print", "/usage",
                stdin=slave_fd,
                stdout=slave_fd,
                stderr=slave_fd,
                env=env,
                cwd="/root",
                close_fds=True,
                start_new_session=True,
            )
            os.close(slave_fd)

            auth_url: Optional[str] = None
            start_time = time.time()
            url_pattern = re.compile(r"https://accounts\.google\.com/o/oauth2/auth\S+")
            buffer = ""

            loop = asyncio.get_running_loop()

            def read_pty_chunk(fd: int) -> str:
                try:
                    r, _, _ = select.select([fd], [], [], 0.1)
                    if fd in r:
                        return os.read(fd, 2048).decode("utf-8", errors="replace")
                except Exception:
                    pass
                return ""

            while time.time() - start_time < timeout:
                chunk = await loop.run_in_executor(None, read_pty_chunk, master_fd)
                if chunk:
                    buffer += chunk
                    logger.debug(f"OAuth PTY chunk for user {user_id}: {chunk}")
                    m = url_pattern.search(buffer)
                    if m:
                        auth_url = m.group(0).rstrip(").,\r\n\t '\"")
                        break
                if proc.returncode is not None:
                    break
                await asyncio.sleep(0.15)

            if not auth_url:
                try:
                    proc.kill()
                except Exception:
                    pass
                try:
                    os.close(master_fd)
                except Exception:
                    pass
                if backup_token:
                    token_file.write_text(backup_token, encoding="utf-8")
                    token_file.chmod(0o600)
                return False, "عدم دریافت لینک احراز هویت از سرور گوگل. لطفاً مجدداً تلاش فرمایید.", None

            self.active_logins[user_id] = {
                "proc": proc,
                "master_fd": master_fd,
                "url": auth_url,
                "backup_token": backup_token,
                "started_at": time.time(),
            }

            return True, "لینک احراز هویت گوگل با موفقیت ایجاد شد.", auth_url

        except Exception as e:
            logger.error(f"Error starting OAuth login for user {user_id}: {e}", exc_info=True)
            if backup_token:
                try:
                    token_file.write_text(backup_token, encoding="utf-8")
                    token_file.chmod(0o600)
                except Exception:
                    pass
            return False, f"خطا در ایجاد فرآیند لاگین: {e}", None

    async def complete_oauth_login(self, user_id: int, auth_code: str, timeout: int = 40) -> Tuple[bool, str]:
        """
        Completes the OAuth login flow by feeding the authorization code into the waiting AGY PTY process.
        """
        import select

        if user_id not in self.active_logins:
            return False, "هیچ فرآیند لاگین فعالی برای شما یافت نشد. لطفاً دستور /login را مجدداً ارسال کنید."

        login_data = self.active_logins[user_id]
        proc = login_data["proc"]
        master_fd = login_data.get("master_fd")
        backup_token = login_data.get("backup_token")
        token_file = self.get_token_path(user_id)

        clean_code = auth_code.strip()
        clean_code = re.sub(r"^[`'\"]+|[`'\"]+$", "", clean_code)

        try:
            input_bytes = (clean_code + "\n").encode("utf-8")
            if master_fd is not None:
                os.write(master_fd, input_bytes)

            start_t = time.time()
            loop = asyncio.get_running_loop()
            out_buffer = ""

            def read_all_pty(fd: int) -> str:
                buf = ""
                while True:
                    r, _, _ = select.select([fd], [], [], 0.3)
                    if fd in r:
                        try:
                            c = os.read(fd, 2048).decode("utf-8", errors="replace")
                            if not c:
                                break
                            buf += c
                        except Exception:
                            break
                    else:
                        break
                return buf

            # Poll for token file creation or process termination
            while time.time() - start_t < timeout:
                if token_file.exists() and token_file.stat().st_size > 10:
                    token_file.chmod(0o600)
                    self.cancel_oauth_login(user_id)
                    return True, "🎉 تبریک! ورود با حساب گوگل اختصاصی شما با موفقیت انجام و ذخیره شد."

                if proc.returncode is not None:
                    break

                await asyncio.sleep(0.5)

            # Read any output from pty
            if master_fd is not None:
                out_buffer = await loop.run_in_executor(None, read_all_pty, master_fd)

            self.cancel_oauth_login(user_id)

            if token_file.exists() and token_file.stat().st_size > 10:
                token_file.chmod(0o600)
                return True, "🎉 تبریک! ورود با حساب گوگل اختصاصی شما با موفقیت انجام و ذخیره شد."

            if backup_token:
                token_file.write_text(backup_token, encoding="utf-8")
                token_file.chmod(0o600)

            err_msg = out_buffer.strip()
            return False, f"کد تایید معتبر نبود یا فرآیند با خطا مواجه شد:\n<pre>{escape_pre(err_msg or 'کد ارسالی نامعتبر بود یا منقضی شده است.')}</pre>"

        except Exception as e:
            logger.error(f"Error completing OAuth login for user {user_id}: {e}", exc_info=True)
            self.cancel_oauth_login(user_id)
            return False, f"خطا در ثبت کد احراز هویت: {e}"

    def cancel_oauth_login(self, user_id: int) -> bool:
        """Cancels any pending login flow for a user and cleans up PTY."""
        login_data = self.active_logins.pop(user_id, None)
        if login_data:
            proc = login_data.get("proc")
            master_fd = login_data.get("master_fd")
            if proc:
                try:
                    proc.kill()
                except Exception:
                    pass
            if master_fd is not None:
                try:
                    os.close(master_fd)
                except Exception:
                    pass
            backup_token = login_data.get("backup_token")
            token_file = self.get_token_path(user_id)
            if backup_token and not token_file.exists():
                try:
                    token_file.write_text(backup_token, encoding="utf-8")
                    token_file.chmod(0o600)
                except Exception:
                    pass
            return True
        return False

    def import_token(self, user_id: int, token_data_str: str) -> Tuple[bool, str]:
        """
        Directly imports an existing antigravity-oauth-token content for the user.
        """
        raw_text = token_data_str.strip()
        try:
            parsed = json.loads(raw_text)
            if not isinstance(parsed, dict) or ("token" not in parsed and "access_token" not in parsed and "refresh_token" not in parsed):
                return False, "فرمت فایل یا متن توکن معتبر نیست. توکن باید فایل JSON حاوی اطلاعات احراز هویت AGY باشد."
        except Exception:
            return False, "متن ارسالی یک ساختار JSON معتبر برای توکن AGY نیست."

        config_dir = self.get_config_dir(user_id)
        config_dir.mkdir(parents=True, exist_ok=True)
        token_file = self.get_token_path(user_id)

        try:
            token_file.write_text(raw_text, encoding="utf-8")
            token_file.chmod(0o600)
            return True, "✅ فایل توکن حساب اختصاصی شما با موفقیت ثبت و فعال شد."
        except Exception as e:
            return False, f"خطا در ذخیره‌سازی توکن: {e}"

    def logout_user(self, user_id: int) -> Tuple[bool, str]:
        """
        Logs out a user by removing their custom token.
        """
        self.cancel_oauth_login(user_id)
        token_file = self.get_token_path(user_id)
        if token_file.exists():
            try:
                token_file.unlink()
                return True, "حساب اختصاصی شما با موفقیت خارج شد. برای ارسال دستورات، ورود مجدد با حساب شخصی الزامی است."
            except Exception as e:
                return False, f"خطا در حذف توکن اختصاصی: {e}"
        return True, "شما در حال حاضر حساب متصلی ندارید."

    def copy_master_token_to_user(self, user_id: int) -> Tuple[bool, str]:
        """
        Disabled: Shared accounts are completely removed.
        """
        return False, "حساب‌های اشتراکی در این ربات غیرفعال شده‌اند و هر کاربر باید با حساب اختصاصی خود وارد شود."

auth_manager = AuthManager()