data-collector-debugtool/app/schemas/s7.py

from ipaddress import ip_address
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


SUPPORTED_POINT_TYPES = {
    "bool",
    "byte",
    "int8",
    "int16",
    "uint16",
    "int32",
    "uint32",
    "int64",
    "uint64",
    "float32",
    "float64",
}

class S7IpRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    ip: str = Field(min_length=1)

    @field_validator("ip")
    @classmethod
    def validate_ip(cls, value: str) -> str:
        try:
            ip_address(value)
        except ValueError as exc:
            raise ValueError("ip must be a valid IP address") from exc
        return value


class S7BaseRequest(S7IpRequest):
    device_type: Literal["S7-1200", "S7-Smart200"] = "S7-1200"
    port: int = Field(default=102, ge=1, le=65535)
    rock: int = Field(ge=0, le=31)
    slot: int = Field(ge=0, le=31)
    tsap_conn_type: Literal["PG", "OP", "BASIC"] = "PG"

    @field_validator("device_type", mode="before")
    @classmethod
    def normalize_device_type(cls, value: str) -> str:
        normalized = str(value).upper()
        if normalized == "S7-1200":
            return "S7-1200"
        if normalized == "S7-SMART200":
            return "S7-Smart200"
        return str(value)

    @field_validator("tsap_conn_type", mode="before")
    @classmethod
    def normalize_tsap_conn_type(cls, value: str) -> str:
        return str(value).upper()


class S7ReadSpec(BaseModel):
    model_config = ConfigDict(extra="forbid")

    area: Literal["DB", "M", "I", "Q", "V"]
    db: int = Field(default=0, ge=0)
    start: int = Field(ge=0)
    size: int = Field(ge=1, le=65535)

    @field_validator("area", mode="before")
    @classmethod
    def normalize_area(cls, value: str) -> str:
        return str(value).upper()

    @model_validator(mode="after")
    def validate_read(self) -> "S7ReadSpec":
        if self.area == "DB" and self.db <= 0:
            raise ValueError("db must be greater than 0 when area is DB")
        return self


class S7RawReadRequest(S7BaseRequest):
    read: S7ReadSpec

    @model_validator(mode="after")
    def validate_v_area(self) -> "S7RawReadRequest":
        if self.device_type != "S7-Smart200" and self.read.area == "V":
            raise ValueError("V area is only supported when device_type is S7-Smart200")
        return self


class S7PointSpec(BaseModel):
    model_config = ConfigDict(extra="forbid")

    area: Literal["DB", "M", "I", "Q", "V"]
    db: int = Field(default=0, ge=0)
    start: int = Field(ge=0)
    type: str
    bit: int | None = Field(default=None, ge=0, le=7)

    @field_validator("area", mode="before")
    @classmethod
    def normalize_area(cls, value: str) -> str:
        return str(value).upper()

    @field_validator("type")
    @classmethod
    def validate_type(cls, value: str) -> str:
        normalized = value.lower()
        if normalized not in SUPPORTED_POINT_TYPES:
            raise ValueError(
                "type must be one of bool, byte, int8, int16, uint16, int32, uint32, int64, uint64, float32, float64"
            )
        return normalized

    @model_validator(mode="after")
    def validate_point(self) -> "S7PointSpec":
        if self.area == "DB" and self.db <= 0:
            raise ValueError("db must be greater than 0 when area is DB")
        if self.type != "bool" and self.bit is not None:
            raise ValueError("bit is only supported for bool points")
        return self


class S7PointReadRequest(S7BaseRequest):
    points: list[S7PointSpec] = Field(min_length=1)

    @model_validator(mode="after")
    def validate_v_area(self) -> "S7PointReadRequest":
        if self.device_type != "S7-Smart200" and any(point.area == "V" for point in self.points):
            raise ValueError("V area is only supported when device_type is S7-Smart200")
        return self


class S7ConnectScanRequest(S7IpRequest):
    device_type: Literal["S7-1200", "S7-Smart200"] = "S7-1200"

    @field_validator("device_type", mode="before")
    @classmethod
    def normalize_device_type(cls, value: str) -> str:
        return S7BaseRequest.normalize_device_type(value)