ForcePilot/backend/package/yuxi/channels/auth/secret_manager.py

from __future__ import annotations

import base64
import os
from enum import Enum
from typing import Any

from cryptography.fernet import Fernet

from yuxi.utils.logging_config import logger

_SENSITIVE_PATTERNS = [
    "token",
    "secret",
    "password",
    "key",
    "credential",
    "api_key",
    "app_secret",
    "bot_token",
    "access_token",
    "refresh_token",
    "private_key",
    "signature",
    "nsec",
    "code",
]

_MASK_LENGTH = 4
_MASK_CHAR = "*"


class SecretSource(Enum):
    CONFIG = "config"
    ENV = "env"
    FILE = "file"
    EXEC = "exec"
    KMS = "kms"
    SECRET_REF = "secret_ref"


class SecretManager:
    def __init__(self, fernet_key: str | None = None):
        self._fernet: Fernet | None = None
        if fernet_key:
            self._init_fernet(fernet_key)
        else:
            env_key = os.getenv("YUXI_FERNET_KEY")
            if env_key:
                self._init_fernet(env_key)
            else:
                logger.warning(
                    "[SecretManager] No Fernet key configured. "
                    "Credentials will be stored in plaintext. "
                    "Set YUXI_FERNET_KEY env var or provide fernet_key parameter."
                )

    def _init_fernet(self, key: str) -> None:
        try:
            self._fernet = Fernet(key.encode() if isinstance(key, str) else key)
            logger.info("[SecretManager] Fernet encryption initialized")
        except Exception as e:
            logger.error(f"[SecretManager] Failed to initialize Fernet: {e}")
            self._fernet = None

    @property
    def is_encryption_available(self) -> bool:
        return self._fernet is not None

    async def resolve_secret(
        self,
        key: str,
        sources: list[SecretSource] | None = None,
        config: dict[str, Any] | None = None,
    ) -> str | None:
        if sources is None:
            sources = [SecretSource.CONFIG, SecretSource.ENV]

        for source in sources:
            if source == SecretSource.CONFIG and config:
                value = config.get(key)
                if value:
                    return value

            if source == SecretSource.ENV:
                value = os.getenv(key) or os.getenv(key.upper())
                if value:
                    return value

            if source == SecretSource.FILE and config:
                file_path = config.get(f"{key}_file")
                if file_path and os.path.isfile(file_path):
                    try:
                        with open(file_path, encoding="utf-8") as f:
                            return f.read().strip()
                    except OSError as e:
                        logger.warning(f"[SecretManager] Failed to read secret file {file_path}: {e}")

            if source == SecretSource.SECRET_REF and config:
                ref = config.get(key)
                if isinstance(ref, str) and ref.startswith("secret://"):
                    resolved = await self._resolve_secret_ref(ref)
                    if resolved:
                        return resolved

            if source == SecretSource.EXEC and config:
                exec_cmd = config.get(f"{key}_exec")
                if exec_cmd:
                    import subprocess

                    try:
                        result = subprocess.run(
                            exec_cmd,
                            shell=True,
                            capture_output=True,
                            text=True,
                            timeout=10,
                        )
                        if result.returncode == 0:
                            return result.stdout.strip()
                        logger.warning(f"[SecretManager] Exec command for {key} failed: {result.stderr}")
                    except Exception as e:
                        logger.warning(f"[SecretManager] Failed to exec {exec_cmd}: {e}")

        return None

    async def _resolve_secret_ref(self, ref: str) -> str | None:
        path = ref.removeprefix("secret://")
        if os.path.isfile(path):
            try:
                with open(path, encoding="utf-8") as f:
                    decrypted = f.read().strip()
                    if self._fernet and self._looks_encrypted(decrypted):
                        decrypted = self.decrypt_secret(decrypted)
                    return decrypted
            except OSError as e:
                logger.warning(f"[SecretManager] Failed to read secret ref {path}: {e}")
        return None

    def encrypt_secret(self, value: str) -> str:
        if not self._fernet:
            logger.warning("[SecretManager] Encryption not available, storing plaintext")
            return value
        encrypted = self._fernet.encrypt(value.encode())
        return base64.urlsafe_b64encode(encrypted).decode()

    def decrypt_secret(self, encrypted: str) -> str:
        if not self._fernet:
            return encrypted
        try:
            decoded = base64.urlsafe_b64decode(encrypted.encode())
            return self._fernet.decrypt(decoded).decode()
        except Exception:
            return encrypted

    @staticmethod
    def redact(value: str | None) -> str:
        if not value:
            return "***"
        if len(value) <= _MASK_LENGTH * 2:
            return _MASK_CHAR * (len(value) if len(value) < 8 else 8)
        return value[:_MASK_LENGTH] + _MASK_CHAR * 8 + value[-_MASK_LENGTH:]

    @staticmethod
    def is_sensitive_key(key: str) -> bool:
        key_lower = key.lower()
        return any(pattern in key_lower for pattern in _SENSITIVE_PATTERNS)

    @staticmethod
    def _looks_encrypted(value: str) -> bool:
        try:
            decoded = base64.urlsafe_b64decode(value.encode())
            return len(decoded) > 32
        except Exception:
            return False

    @staticmethod
    def redact_config(config: dict[str, Any]) -> dict[str, Any]:
        result = {}
        for k, v in config.items():
            if SecretManager.is_sensitive_key(k):
                if isinstance(v, str):
                    result[k] = SecretManager.redact(v)
                else:
                    result[k] = "***"
            elif isinstance(v, dict):
                result[k] = SecretManager.redact_config(v)
            else:
                result[k] = v
        return result


_secret_manager: SecretManager | None = None


def get_secret_manager() -> SecretManager:
    global _secret_manager
    if _secret_manager is None:
        _secret_manager = SecretManager()
    return _secret_manager


def init_secret_manager(fernet_key: str | None = None) -> SecretManager:
    global _secret_manager
    _secret_manager = SecretManager(fernet_key=fernet_key)
    return _secret_manager