minerops/backend/app/agent_zones/service.py

"""Business rules for Agent Zone routing and transitional check-in."""

from typing import Any
from uuid import UUID

from app.agent_zones.exceptions import (
    AgentZoneNotFound,
    AgentZoneValidationError,
)
from app.agent_zones.models import (
    AgentCheckinRequest,
    AgentCheckinResponse,
    AgentZoneRoutingContext,
    AgentZoneCreateRequest,
    AgentZoneSummary,
    MinerZoneAssociationResult,
)
from app.agent_zones.repository import AgentZoneRepository


SENSITIVE_KEYS = {
    "agent_secret",
    "api3_account",
    "api3_password",
    "password",
    "salt",
    "token",
}


def _sanitize(value: Any) -> Any:
    if isinstance(value, dict):
        return {
            key: _sanitize(item)
            for key, item in value.items()
            if key.lower() not in SENSITIVE_KEYS
            and not any(
                marker in key.lower()
                for marker in ("secret", "password", "token", "salt")
            )
        }
    if isinstance(value, list):
        return [_sanitize(item) for item in value]
    return value


class AgentZoneService:
    @staticmethod
    def validate_agent_identity(
        authenticated_agent_id: str, claimed_agent_id: str | None
    ) -> None:
        if claimed_agent_id and claimed_agent_id.strip() != authenticated_agent_id:
            raise AgentZoneValidationError(
                "Payload agent_id does not match authenticated Agent identity"
            )

    def __init__(self, repository: AgentZoneRepository):
        self.repository = repository

    def list_zones(self, site_id: UUID | None = None) -> list[AgentZoneSummary]:
        return [
            AgentZoneSummary.model_validate(row)
            for row in self.repository.list_zones(site_id)
        ]

    def create_zone(self, request: AgentZoneCreateRequest) -> AgentZoneSummary:
        zone_id = self.repository.create_zone(
            site_id=request.site_id,
            name=request.name,
            description=request.description,
            networks=request.networks,
        )
        for zone in self.list_zones(request.site_id):
            if zone.id == zone_id:
                return zone
        raise AgentZoneNotFound("Created Agent Zone was not found")

    def checkin(
        self, authenticated_agent_id: str, request: AgentCheckinRequest
    ) -> AgentCheckinResponse:
        self.validate_agent_identity(authenticated_agent_id, request.agent_id)

        existing = self.repository.get_agent(authenticated_agent_id)
        requested_site = request.site_id
        requested_zone = request.agent_zone_id
        final_site = requested_site or (existing or {}).get("site_id")
        final_zone: UUID | None = None
        reason: str | None = None

        established_zone = None
        established_zone_id = (existing or {}).get("agent_zone_id")
        if established_zone_id is not None:
            established_zone = self.repository.get_zone(established_zone_id)
            if not established_zone:
                raise AgentZoneValidationError(
                    "Existing Agent Zone assignment is invalid"
                )
            if (
                requested_site is not None
                and established_zone["site_id"] != requested_site
            ):
                raise AgentZoneValidationError(
                    "Agent cannot change Sites while assigned to an Agent Zone"
                )

        if requested_zone is not None:
            zone = self.repository.get_zone(requested_zone)
            if not zone:
                raise AgentZoneNotFound("Agent Zone not found")
            if zone["site_id"] != requested_site:
                raise AgentZoneValidationError(
                    "Agent Zone does not belong to the supplied Site"
                )
            if not zone["is_enabled"]:
                raise AgentZoneValidationError("Agent Zone is disabled")
            if established_zone_id is not None and established_zone_id != requested_zone:
                raise AgentZoneValidationError(
                    "Agent is already assigned to a different Agent Zone"
                )
            final_zone = requested_zone
        elif established_zone is not None:
            final_site = established_zone["site_id"]
            final_zone = established_zone["id"]
            if not established_zone["is_enabled"]:
                reason = "agent_zone_disabled"
        elif final_site is not None:
            usable = self.repository.list_usable_zones(final_site)
            if len(usable) == 1:
                final_zone = usable[0]["id"]
            elif not usable:
                reason = "agent_zone_unavailable"
            else:
                reason = "agent_zone_ambiguous"
        else:
            reason = "site_unassigned"

        site_name = (request.site_name or "").strip()
        if not site_name:
            site_name = str((existing or {}).get("site_name") or "Unassigned")
        payload = _sanitize(request.model_dump(mode="json"))
        self.repository.save_checkin(
            agent_id=authenticated_agent_id,
            site_id=final_site,
            agent_zone_id=final_zone,
            site_name=site_name,
            hostname=(request.hostname or "").strip() or None,
            version=(request.version or "").strip() or None,
            ip_address=(request.ip_address or "").strip() or None,
            payload=payload,
        )
        routing_available = final_zone is not None and reason is None
        return AgentCheckinResponse(
            agent_id=authenticated_agent_id,
            site_id=final_site,
            agent_zone_id=final_zone,
            routing_available=routing_available,
            routing_reason=reason,
            site_name=site_name,
        )

    def routing_context(
        self, agent_id: str, recorded_zone_id: UUID | None = None
    ) -> AgentZoneRoutingContext:
        assignment = self.repository.get_agent_assignment(agent_id)
        if not assignment or assignment.get("agent_zone_id") is None:
            return AgentZoneRoutingContext(routing_reason="agent_zone_unassigned")
        zone_id = assignment["agent_zone_id"]
        if recorded_zone_id is not None and recorded_zone_id != zone_id:
            return AgentZoneRoutingContext(routing_reason="agent_zone_conflict")
        if assignment.get("zone_enabled") is not True:
            return AgentZoneRoutingContext(
                site_id=assignment.get("site_id"),
                agent_zone_id=zone_id,
                routing_reason="agent_zone_disabled",
            )
        return AgentZoneRoutingContext(
            site_id=assignment.get("site_id"), agent_zone_id=zone_id
        )

    def configuration(self, agent_id: str) -> dict[str, Any] | None:
        row = self.repository.get_agent_configuration(agent_id)
        if not row:
            return None
        return {
            "agent_id": row["agent_id"],
            "site_id": row["site_id"],
            "site_name": row["site_name"],
            "agent_zone_id": row["agent_zone_id"],
            "agent_zone_name": row["agent_zone_name"],
            "ip_ranges": [
                {
                    "name": row["agent_zone_name"],
                    "cidr": network,
                    "enabled": row["agent_zone_enabled"] is True,
                    "api_port": 4028,
                    "api3_port": 4433,
                }
                for network in row.get("networks") or []
            ],
        }

    def associate_miners_from_agent(
        self, agent_id: str, miner_ids: list[UUID]
    ) -> MinerZoneAssociationResult:
        routing = self.routing_context(agent_id)
        if routing.routing_reason:
            return MinerZoneAssociationResult(
                agent_zone_id=routing.agent_zone_id,
                routing_reason=routing.routing_reason,
            )
        return MinerZoneAssociationResult(
            agent_zone_id=routing.agent_zone_id,
            outcomes=self.repository.assign_miners_to_zone(
                miner_ids, routing.agent_zone_id
            ),
        )