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()