netx/netx_api/ne_schemas.py

from __future__ import annotations

from datetime import datetime
from typing import Literal

from pydantic import BaseModel, Field, field_validator

from .device_types import SUPPORTED_VENDORS
from .ne_exec_guard import EXEC_POLICIES, EXEC_POLICY_READONLY

ConnectStatus = Literal["unknown", "testing", "pass", "fail"]
ExecPolicy = Literal["readonly", "linux_shell", "unrestricted"]


class ManagedNeCreate(BaseModel):
    name: str = ""
    vendor: str
    device_type: str
    ip_address: str
    port: int = 22
    protocol: str = "ssh"
    username: str
    password: str = ""
    tags: str = ""
    remark: str = ""
    exec_policy: ExecPolicy = "readonly"
    hop_enabled: bool = False
    hop_vendor: str = "zte"
    hop_host: str = ""
    hop_port: int = 22
    hop_protocol: str = "ssh"
    hop_username: str = ""
    hop_password: str = ""
    hop_command_template: str = ""
    hop_vrf: str = ""
    hop_target_auth_mode: str = "bastion_managed"
    hop_enter_system_view: bool = False

    @field_validator("vendor")
    @classmethod
    def normalize_vendor(cls, v: str) -> str:
        raw = str(v or "").strip()
        if not raw:
            raise ValueError("vendor_required")
        for item in SUPPORTED_VENDORS:
            if item.lower() == raw.lower():
                return item
        return "Other"

    @field_validator("exec_policy", mode="before")
    @classmethod
    def normalize_exec_policy_create(cls, v: object) -> str:
        if v is None or str(v).strip() == "":
            return EXEC_POLICY_READONLY
        raw = str(v).strip().lower()
        if raw not in EXEC_POLICIES:
            raise ValueError("unsupported_exec_policy")
        return raw


class ManagedNeUpdate(BaseModel):
    name: str | None = None
    vendor: str | None = None
    device_type: str | None = None
    ip_address: str | None = None
    port: int | None = None
    protocol: str | None = None
    username: str | None = None
    password: str | None = None
    tags: str | None = None
    remark: str | None = None
    exec_policy: ExecPolicy | None = None
    hop_enabled: bool | None = None
    hop_vendor: str | None = None
    hop_host: str | None = None
    hop_port: int | None = None
    hop_protocol: str | None = None
    hop_username: str | None = None
    hop_password: str | None = None
    hop_command_template: str | None = None
    hop_vrf: str | None = None
    hop_target_auth_mode: str | None = None
    hop_enter_system_view: bool | None = None

    @field_validator("vendor")
    @classmethod
    def normalize_vendor_update(cls, v: str | None) -> str | None:
        if v is None:
            return None
        raw = str(v).strip()
        if not raw:
            raise ValueError("vendor_required")
        for item in SUPPORTED_VENDORS:
            if item.lower() == raw.lower():
                return item
        return "Other"

    @field_validator("exec_policy", mode="before")
    @classmethod
    def normalize_exec_policy_update(cls, v: object) -> str | None:
        if v is None:
            return None
        raw = str(v).strip().lower()
        if not raw:
            return EXEC_POLICY_READONLY
        if raw not in EXEC_POLICIES:
            raise ValueError("unsupported_exec_policy")
        return raw


class ManagedNeOut(BaseModel):
    id: str
    name: str
    vendor: str
    device_type: str
    ip_address: str
    port: int
    protocol: str
    username: str
    # True when a non-empty password ciphertext is stored (never returns the secret).
    has_password: bool = False
    connect_status: ConnectStatus
    connect_message: str
    connect_detail: str = ""
    connect_tested_at: datetime | None
    tags: str
    remark: str
    # Provenance: "" | ume_sync | webcrt | lldp | …
    source: str = ""
    source_ref: str = ""
    exec_policy: ExecPolicy = "readonly"
    hop_enabled: bool = False
    hop_vendor: str = "zte"
    hop_host: str = ""
    hop_port: int = 22
    hop_protocol: str = "ssh"
    hop_username: str = ""
    hop_command_template: str = ""
    hop_vrf: str = ""
    hop_target_auth_mode: str = "bastion_managed"
    hop_enter_system_view: bool = False
    created_at: datetime
    updated_at: datetime


class ConnectTestRequest(BaseModel):
    ids: list[str] = Field(min_length=1)


class ManagedNeExecRequest(BaseModel):
    """Run CLI on a managed NE or UME inventory NE (gates follow managed NE exec_policy)."""

    ne_id: str | None = None
    ume_ne_id: str | None = None
    # Absolute ceiling; runtime limit is settings.ne_exec_max_commands (default 5).
    commands: list[str] = Field(min_length=1, max_length=50)
    read_timeout_sec: int | None = Field(default=None, ge=10, le=120)


class ManagedNeExecBatchTarget(BaseModel):
    """One NE in a concurrent exec batch (exactly one of ne_id / ume_ne_id)."""

    ne_id: str | None = None
    ume_ne_id: str | None = None
    commands: list[str] | None = Field(default=None, max_length=50)


class ManagedNeExecBatchRequest(BaseModel):
    """Run the same (or per-target) read-only CLI on many NEs concurrently."""

    targets: list[ManagedNeExecBatchTarget] | None = Field(default=None, max_length=20)
    ne_ids: list[str] | None = Field(default=None, max_length=20)
    ume_ne_ids: list[str] | None = Field(default=None, max_length=20)
    commands: list[str] | None = Field(default=None, max_length=50)
    read_timeout_sec: int | None = Field(default=None, ge=10, le=120)
    concurrency: int | None = Field(default=4, ge=1, le=8)


class HopProxyConfig(BaseModel):
    """Shared jump-host (proxy) settings applied to one or many NEs."""

    hop_vendor: str = "zte"
    hop_host: str
    hop_port: int = 22
    hop_protocol: str = "ssh"
    hop_username: str
    hop_password: str = ""
    hop_command_template: str = ""
    hop_vrf: str = ""
    hop_target_auth_mode: str = "bastion_managed"
    hop_enter_system_view: bool = False


class BatchHopApplyRequest(BaseModel):
    ids: list[str] = Field(min_length=1)
    hop: HopProxyConfig


class BatchAccountConfig(BaseModel):
    username: str = ""
    password: str = ""


class BatchAccountApplyRequest(BaseModel):
    ids: list[str] = Field(min_length=1)
    account: BatchAccountConfig


class ImportFailure(BaseModel):
    row: int
    reason: str


class ImportResult(BaseModel):
    inserted: int
    updated: int
    failed: list[ImportFailure]


class UmeManagedSyncResult(BaseModel):
    inserted: int
    updated: int
    deleted: int
    total_inventory: int


class UmeManagedDeleteResult(BaseModel):
    deleted: int