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

from __future__ import annotations

import base64
import hashlib
import json
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from cryptography.fernet import Fernet

from yuxi.utils.logging_config import logger

if TYPE_CHECKING:
    from yuxi.channels.services.plugin_state_store import PluginStateStore

_CREDENTIALS_DIR = ".yuxi"
_CREDENTIALS_SUBDIR = "credentials"
_CREDENTIALS_FILE = "matrix_credentials.json"


@dataclass
class MatrixCredentials:
    user_id: str = ""
    homeserver: str = ""
    access_token: str = ""
    device_id: str = ""
    display_name: str = ""

    def to_dict(self) -> dict[str, str]:
        return {
            "user_id": self.user_id,
            "homeserver": self.homeserver,
            "access_token": self.access_token,
            "device_id": self.device_id,
            "display_name": self.display_name,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> MatrixCredentials:
        return cls(
            user_id=data.get("user_id", ""),
            homeserver=data.get("homeserver", ""),
            access_token=data.get("access_token", ""),
            device_id=data.get("device_id", ""),
            display_name=data.get("display_name", ""),
        )

    @property
    def is_valid(self) -> bool:
        return bool(self.user_id and self.homeserver and self.access_token)


class CredentialStore:
    def __init__(self, base_dir: str | None = None):
        if base_dir:
            self._base_dir = base_dir
        else:
            home = os.path.expanduser("~")
            self._base_dir = os.path.join(home, _CREDENTIALS_DIR)
        self._store_dir = os.path.join(self._base_dir, _CREDENTIALS_SUBDIR, "matrix")

    def _get_path(self, account_id: str) -> str:
        safe_id = account_id.replace("@", "_").replace(":", "_").replace("/", "_")
        return os.path.join(self._store_dir, f"{safe_id}.json")

    @staticmethod
    def _get_fernet() -> Fernet | None:
        key = os.environ.get("MATRIX_CREDENTIAL_KEY")
        if not key:
            return None
        digest = hashlib.sha256(key.encode()).digest()
        return Fernet(base64.urlsafe_b64encode(digest))

    def save(self, credentials: MatrixCredentials) -> bool:
        os.makedirs(self._store_dir, exist_ok=True)
        filepath = self._get_path(credentials.user_id or "default")
        try:
            data = credentials.to_dict()
            fernet = self._get_fernet()
            payload = json.dumps(data)
            if fernet:
                payload = fernet.encrypt(payload.encode()).decode()
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(payload)
            os.chmod(filepath, 0o600)
            safe_data = {k: "***" if k == "access_token" else v for k, v in data.items()}
            logger.info(f"Matrix credentials saved: {filepath} ({safe_data})")
            return True
        except OSError as e:
            logger.error(f"Matrix credentials save failed: {e}")
            return False

    def load(self, account_id: str = "default") -> MatrixCredentials | None:
        filepath = self._get_path(account_id)
        if not os.path.exists(filepath):
            return None
        try:
            with open(filepath, encoding="utf-8") as f:
                raw = f.read()
            fernet = self._get_fernet()
            if fernet:
                try:
                    raw = fernet.decrypt(raw.encode()).decode()
                except Exception:
                    pass
            data = json.loads(raw)
            return MatrixCredentials.from_dict(data)
        except (json.JSONDecodeError, OSError) as e:
            logger.error(f"Matrix credentials load failed: {e}")
            return None

    def delete(self, account_id: str) -> bool:
        filepath = self._get_path(account_id)
        try:
            if os.path.exists(filepath):
                os.remove(filepath)
                logger.info(f"Matrix credentials deleted: {filepath}")
            return True
        except OSError as e:
            logger.error(f"Matrix credentials delete failed: {e}")
            return False

    def list_accounts(self) -> list[str]:
        if not os.path.isdir(self._store_dir):
            return []
        accounts = []
        for filename in os.listdir(self._store_dir):
            if filename.endswith(".json"):
                accounts.append(filename.rsplit(".", 1)[0])
        return accounts

    def load_all(self) -> list[MatrixCredentials]:
        accounts = []
        for account_id in self.list_accounts():
            creds = self.load(account_id)
            if creds and creds.is_valid:
                accounts.append(creds)
        return accounts


class DbCredentialStore:
    def __init__(self, state_store: PluginStateStore, channel_id: str = "matrix"):
        self._store = state_store
        self._channel = channel_id

    async def save(self, credentials: MatrixCredentials) -> bool:
        account_id = credentials.user_id or "default"
        try:
            await self._store.set(
                self._channel,
                f"cred:{account_id}",
                credentials.to_dict(),
                namespace="credentials",
            )
            return True
        except Exception as e:
            logger.error(f"DbCredentialStore: save failed: {e}")
            return False

    async def load(self, account_id: str = "default") -> MatrixCredentials | None:
        try:
            data = await self._store.get(self._channel, f"cred:{account_id}", namespace="credentials")
            return MatrixCredentials.from_dict(data) if data else None
        except Exception as e:
            logger.error(f"DbCredentialStore: load failed: {e}")
            return None

    async def delete(self, account_id: str) -> bool:
        try:
            await self._store.delete(self._channel, f"cred:{account_id}", namespace="credentials")
            return True
        except Exception as e:
            logger.error(f"DbCredentialStore: delete failed: {e}")
            return False

    async def list_accounts(self) -> list[str]:
        try:
            keys = await self._store.list_keys(self._channel, namespace="credentials")
            return [k[len("cred:") :] for k in keys if k.startswith("cred:")]
        except Exception as e:
            logger.error(f"DbCredentialStore: list_accounts failed: {e}")
            return []

    async def load_all(self) -> list[MatrixCredentials]:
        accounts = []
        for account_id in await self.list_accounts():
            creds = await self.load(account_id)
            if creds and creds.is_valid:
                accounts.append(creds)
        return accounts