src/agents/executor/handlers.py - master - azathd/mutiagent

src/agents/executor/handlers.py
57 строк2 KB
Your Name
init
15 май 2026, 17:47
15 май 2026, 17:47359c81e
"""Action handlers for approved actions."""

from __future__ import annotations

import json
from datetime import UTC, datetime
from typing import Any

from minio import Minio

from src.agents.executor.marzban import MarzbanClient
from src.agents.executor.nftables import NftablesManager
from src.common.models import ActionType
from src.common.settings import settings


class ActionHandlers:
    def __init__(self) -> None:
        self.marzban = MarzbanClient()
        self.nft = NftablesManager()
        self.minio = Minio(
            settings.minio_endpoint.replace("http://", "").replace("https://", ""),
            access_key=settings.minio_access_key,
            secret_key=settings.minio_secret_key,
            secure=settings.minio_endpoint.startswith("https"),
        )

    async def snapshot(self, action_id: str, state: dict[str, Any]) -> str:
        key = f"snapshots/{datetime.now(tz=UTC).strftime('%Y%m%dT%H%M%S')}-{action_id}.json"
        data = json.dumps(state).encode()
        self.minio.put_object("xray", key, data=data, length=len(data), content_type="application/json")
        return key

    async def apply(self, action_type: str, target: str, payload: dict[str, Any]) -> dict[str, Any]:
        dry = settings.executor_dry_run
        if action_type == ActionType.USER_SUSPEND.value:
            if not dry:
                await self.marzban.user_disable(target)
            return {"status": "suspended", "dry_run": dry}
        if action_type == ActionType.USER_UNSUSPEND.value:
            if not dry:
                await self.marzban.user_enable(target)
            return {"status": "enabled", "dry_run": dry}
        if action_type == ActionType.USER_SET_LIMIT.value:
            limit = int(payload.get("data_limit", 0))
            if not dry:
                await self.marzban.user_set_limit(target, limit)
            return {"status": "limit_set", "limit": limit, "dry_run": dry}
        if action_type == ActionType.USER_DELETE.value:
            if not dry:
                await self.marzban.user_delete(target)
            return {"status": "deleted", "dry_run": dry}
        if action_type == ActionType.FIREWALL_BLOCK_IP.value:
            ttl = int(payload.get("ttl_sec", 3600))
            ok = await self.nft.block_ip(target, ttl)
            return {"status": "blocked" if ok else "failed", "dry_run": dry}
        return {"status": "noop"}