ForcePilot/backend/package/yuxi/channel/security/fix.py

from __future__ import annotations

import logging
import os
import stat
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

_O_NOFOLLOW = os.O_NOFOLLOW if hasattr(os, "O_NOFOLLOW") else 0
_HAS_FCHMOD = hasattr(os, "fchmod")


@dataclass
class SecurityFixResult:
    status: str = "ok"
    path: str = ""
    changes: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return self.status == "ok" and len(self.errors) == 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status,
            "path": self.path,
            "changes": self.changes,
            "errors": self.errors,
        }


@dataclass
class SecurityPermissionTarget:
    path: Path
    mode: int
    expected_type: str


def _safe_chmod(path: Path, mode: int, expected_type: str) -> SecurityFixResult:
    result = SecurityFixResult(path=str(path))

    try:
        fd = os.open(path, os.O_RDONLY | _O_NOFOLLOW)
    except FileNotFoundError:
        result.status = "skip"
        result.errors.append(f"{path}: file not found")
        return result
    except PermissionError:
        result.status = "skip"
        result.errors.append(f"{path}: permission denied reading stat")
        return result
    except OSError as e:
        result.status = "error"
        result.errors.append(f"{path}: open failed ({e})")
        return result

    try:
        st = os.fstat(fd)
    except OSError as e:
        os.close(fd)
        result.status = "error"
        result.errors.append(f"{path}: fstat failed ({e})")
        return result

    is_dir = stat.S_ISDIR(st.st_mode)
    is_file = stat.S_ISREG(st.st_mode)

    if expected_type == "dir" and not is_dir:
        os.close(fd)
        result.status = "skip"
        result.errors.append(f"{path}: expected directory but is not")
        return result

    if expected_type == "file" and not is_file:
        os.close(fd)
        result.status = "skip"
        result.errors.append(f"{path}: expected file but is not")
        return result

    current_mode = stat.S_IMODE(st.st_mode)
    if (current_mode & 0o777) == mode:
        os.close(fd)
        result.status = "ok"
        result.changes.append(f"{path}: mode already {oct(mode)}")
        return result

    try:
        if _HAS_FCHMOD:
            os.fchmod(fd, mode)
        else:
            os.chmod(path, mode)
        result.status = "ok"
        result.changes.append(f"{path}: mode {oct(current_mode)} -> {oct(mode)}")
    except PermissionError:
        result.status = "error"
        result.errors.append(f"{path}: permission denied setting mode {oct(mode)}")
    except OSError as e:
        result.status = "error"
        result.errors.append(f"{path}: fchmod failed ({e})")
    finally:
        os.close(fd)

    return result


def collect_security_permission_targets(
    state_dir: str | Path | None = None,
    config_path: str | Path | None = None,
) -> list[SecurityPermissionTarget]:
    targets: list[SecurityPermissionTarget] = []

    if state_dir:
        p = Path(state_dir)
        targets.append(SecurityPermissionTarget(path=p, mode=0o700, expected_type="dir"))

        creds_dir = p / "credentials"
        if creds_dir.exists() and creds_dir.is_dir():
            targets.append(SecurityPermissionTarget(path=creds_dir, mode=0o700, expected_type="dir"))

        oauth_dir = p / "oauth"
        if oauth_dir.exists() and oauth_dir.is_dir():
            targets.append(SecurityPermissionTarget(path=oauth_dir, mode=0o700, expected_type="dir"))

        agents_dir = p / "agents"
        if agents_dir.exists() and agents_dir.is_dir():
            for agent_dir in agents_dir.iterdir():
                if agent_dir.is_dir():
                    targets.append(SecurityPermissionTarget(path=agent_dir, mode=0o700, expected_type="dir"))
                    auth_file = agent_dir / "auth-profiles.json"
                    if auth_file.exists():
                        targets.append(SecurityPermissionTarget(path=auth_file, mode=0o600, expected_type="file"))

    if config_path:
        p = Path(config_path)
        targets.append(SecurityPermissionTarget(path=p, mode=0o600, expected_type="file"))

    return targets


def fix_security_permissions(
    targets: list[SecurityPermissionTarget],
) -> list[SecurityFixResult]:
    results: list[SecurityFixResult] = []
    for target in targets:
        result = _safe_chmod(target.path, target.mode, target.expected_type)
        results.append(result)

        if result.status == "error":
            logger.warning("Security fix failed: %s", result.errors)
        elif result.status == "ok" and any("->" in c for c in result.changes):
            logger.info("Security fix applied: %s", result.changes)

    return results


def fix_group_policy_in_config(config: dict[str, Any]) -> bool:
    if not isinstance(config.get("channels"), dict):
        return False

    fixed = 0
    channels = config["channels"]

    def _fix(scope: dict):
        nonlocal fixed
        if not isinstance(scope, dict):
            return
        for key, value in scope.items():
            if key == "groupPolicy" and value == "open":
                scope[key] = "allowlist"
                fixed += 1
            elif isinstance(value, dict):
                _fix(value)

    for channel_id, channel_value in list(channels.items()):
        if isinstance(channel_value, dict):
            _fix(channel_value)

    if fixed > 0:
        logger.info("Fixed %d groupPolicy=open -> allowlist", fixed)

    return fixed > 0


def fix_logging_redact_in_config(config: dict[str, Any]) -> bool:
    logging_config = config.get("logging", {})
    if isinstance(logging_config, dict) and logging_config.get("redactSensitive") == "off":
        logging_config["redactSensitive"] = "tools"
        logger.info("Fixed logging.redactSensitive: off -> tools")
        return True
    return False


def apply_security_fix_config_mutations(config: dict[str, Any]) -> bool:
    mutated = fix_logging_redact_in_config(config)
    mutated |= fix_group_policy_in_config(config)
    return mutated


async def fix_security(
    config: dict[str, Any],
    state_dir: str | Path | None = None,
    config_path: str | Path | None = None,
) -> list[SecurityFixResult]:
    results: list[SecurityFixResult] = []

    if apply_security_fix_config_mutations(config):
        results.append(
            SecurityFixResult(
                status="ok",
                path="<config>",
                changes=["Applied config mutations for security hardening"],
            )
        )

    targets = collect_security_permission_targets(
        state_dir=state_dir,
        config_path=config_path,
    )
    perm_results = fix_security_permissions(targets)
    results.extend(perm_results)

    return results


__all__ = [
    "SecurityFixResult",
    "SecurityPermissionTarget",
    "collect_security_permission_targets",
    "fix_security_permissions",
    "fix_group_policy_in_config",
    "fix_logging_redact_in_config",
    "apply_security_fix_config_mutations",
    "fix_security",
]