ForcePilot/backend/package/yuxi/channels/adapters/zalo_user/credentials.py

from __future__ import annotations

import hashlib
import json
import os
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any

from yuxi.utils.logging_config import logger

from .bridge import BridgeClient
from .constants import ZaloCredentialStage


def _resolve_credential_dir(config: dict[str, Any], profile: str = "default") -> Path:
    data_dir = config.get("data_dir", config.get("dataDir", ""))
    if data_dir:
        cred_dir = Path(data_dir) / "channels" / "zalo_user" / profile
    else:
        cred_dir = Path(os.getcwd()) / "data" / "channels" / "zalo_user" / profile
    cred_dir.mkdir(parents=True, exist_ok=True)
    return cred_dir


def _credential_file_path(config: dict[str, Any], profile: str = "default") -> Path:
    return _resolve_credential_dir(config, profile) / "credentials.json"


def _atomic_write(file_path: Path, content: str) -> None:
    _check_not_symlink(file_path)

    tmp_fd, tmp_path = tempfile.mkstemp(dir=str(file_path.parent), prefix=".creds_")
    try:
        with os.fdopen(tmp_fd, "w", encoding="utf-8") as f:
            f.write(content)
        _set_restrictive_permissions(Path(tmp_path))
        os.replace(tmp_path, str(file_path))
    except Exception:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise


def _check_not_symlink(file_path: Path) -> None:
    try:
        resolved = file_path.resolve()
        if file_path.is_symlink() or file_path != resolved:
            logger.warning(f"[ZaloUser] Credential path appears to be a symlink, using resolved path: {resolved}")
    except OSError:
        pass


def _set_restrictive_permissions(file_path: Path) -> None:
    if os.name == "posix":
        os.chmod(str(file_path), 0o600)
    elif os.name == "nt":
        try:
            subprocess.run(
                ["icacls", str(file_path), "/inheritance:r", "/grant", f"{os.environ.get('USERNAME', 'SYSTEM')}:F"],
                check=False,
                capture_output=True,
                timeout=3,
            )
        except Exception:
            pass


class CredentialManager:
    def __init__(self, config: dict[str, Any] | None = None, profile: str = "default"):
        self._stage = ZaloCredentialStage.UNINITIALIZED
        self._credential_signature: str | None = None
        self._credential_data: dict[str, Any] = {}
        self._last_check_time: float = 0.0
        self._error_message: str | None = None
        self._config = config or {}
        self._profile = profile

    @property
    def stage(self) -> ZaloCredentialStage:
        return self._stage

    @property
    def is_authenticated(self) -> bool:
        return self._stage == ZaloCredentialStage.CONNECTED

    @property
    def credential_signature(self) -> str | None:
        return self._credential_signature

    @property
    def error_message(self) -> str | None:
        return self._error_message

    def mark_qr_pending(self, qr_id: str = "") -> None:
        self._stage = ZaloCredentialStage.QR_PENDING
        self._credential_data = {"qr_id": qr_id}
        self._error_message = None
        logger.info(f"[ZaloUser] Credential stage -> QR_PENDING (qr_id={qr_id})")

    def mark_qr_scanned(self) -> None:
        self._stage = ZaloCredentialStage.QR_SCANNED
        self._error_message = None
        logger.info("[ZaloUser] Credential stage -> QR_SCANNED")

    def mark_qr_declined(self) -> None:
        self._stage = ZaloCredentialStage.QR_DECLINED
        self._error_message = "QR code login was declined by user"
        logger.warning("[ZaloUser] Credential stage -> QR_DECLINED")

    def mark_qr_expired(self) -> None:
        self._stage = ZaloCredentialStage.QR_EXPIRED
        self._error_message = "QR code expired"
        logger.warning("[ZaloUser] Credential stage -> QR_EXPIRED")

    def mark_connected(self, credential_data: dict[str, Any] | None = None) -> None:
        self._stage = ZaloCredentialStage.CONNECTED
        if credential_data:
            self._credential_data = credential_data
            self._credential_signature = _compute_signature(credential_data)
            self._save()
        self._error_message = None
        logger.info("[ZaloUser] Credential stage -> CONNECTED")

    def mark_logged_out(self) -> None:
        self._stage = ZaloCredentialStage.LOGGED_OUT
        self._credential_data = {}
        self._credential_signature = None
        self._error_message = None
        self._delete_file()
        logger.info("[ZaloUser] Credential stage -> LOGGED_OUT")

    def mark_error(self, error: str) -> None:
        self._stage = ZaloCredentialStage.ERROR
        self._error_message = error
        logger.error(f"[ZaloUser] Credential stage -> ERROR: {error}")

    def has_credential_changed(self, new_data: dict[str, Any]) -> bool:
        if self._credential_signature is None:
            return True
        new_sig = _compute_signature(new_data)
        return new_sig != self._credential_signature

    def update_credentials(self, data: dict[str, Any]) -> bool:
        if not self.has_credential_changed(data):
            return False
        self._credential_data = data
        self._credential_signature = _compute_signature(data)
        self._save()
        return True

    async def check_session_status(self, bridge: BridgeClient) -> bool:
        try:
            status = await bridge.check_login_status()
            now = time.monotonic()
            self._last_check_time = now

            if status.get("logged_in"):
                if self._stage != ZaloCredentialStage.CONNECTED:
                    self.mark_connected(status)
                return True
            else:
                if self._stage == ZaloCredentialStage.CONNECTED:
                    self.mark_logged_out()
                    return False
                return False
        except Exception as e:
            logger.warning(f"[ZaloUser] Session status check failed: {e}")
            return self._stage == ZaloCredentialStage.CONNECTED

    async def logout(self, bridge: BridgeClient) -> bool:
        try:
            await bridge.post("/login/logout")
            self.mark_logged_out()
            return True
        except Exception as e:
            logger.warning(f"[ZaloUser] Logout request failed: {e}")
            self.mark_logged_out()
            return False

    def reset(self) -> None:
        self._stage = ZaloCredentialStage.UNINITIALIZED
        self._credential_signature = None
        self._credential_data = {}
        self._last_check_time = 0.0
        self._error_message = None
        self._delete_file()

    def to_dict(self) -> dict[str, Any]:
        return {
            "stage": self._stage.value,
            "is_authenticated": self.is_authenticated,
            "last_check_time": self._last_check_time,
            "error": self._error_message,
        }

    def _save(self) -> None:
        if not self._credential_data:
            return
        try:
            file_path = _credential_file_path(self._config, self._profile)
            content = json.dumps(self._credential_data, ensure_ascii=False, default=str)
            _atomic_write(file_path, content)
            logger.debug(f"[ZaloUser] Credentials persisted to {file_path}")
        except Exception as e:
            logger.warning(f"[ZaloUser] Failed to persist credentials: {e}")

    def _delete_file(self) -> None:
        try:
            file_path = _credential_file_path(self._config, self._profile)
            if file_path.exists():
                file_path.unlink()
        except Exception as e:
            logger.debug(f"[ZaloUser] Failed to delete credential file: {e}")

    def load_from_file(self) -> bool:
        try:
            file_path = _credential_file_path(self._config, self._profile)
            if not file_path.exists():
                return False
            with open(file_path, encoding="utf-8") as f:
                data = json.load(f)
            self._credential_data = data
            self._credential_signature = _compute_signature(data)
            self._stage = ZaloCredentialStage.CONNECTED
            logger.info(f"[ZaloUser] Credentials loaded from {file_path}")
            return True
        except Exception as e:
            logger.warning(f"[ZaloUser] Failed to load credentials from file: {e}")
            return False


def _compute_signature(data: dict[str, Any]) -> str:
    canonical = json.dumps(data, sort_keys=True, default=str)
    return hashlib.sha256(canonical.encode()).hexdigest()