data-collector-debugtool/app/services/s7_service.py

from __future__ import annotations

from itertools import product
from typing import Any

from snap7.client import Client

try:
    from snap7.type import Area
except ImportError:  # pragma: no cover - compatibility for older python-snap7 releases
    Area = None

from app.response import response_payload
from app.schemas.s7 import S7ConnectScanRequest, S7PointReadRequest, S7PointSpec, S7RawReadRequest, S7ReadSpec
from app.services.s7_codec import byte_size_for_type, bytes_to_hex, convert_s7_value


CONNECTION_TYPES = {
    "PG": 0x01,
    "OP": 0x02,
    "BASIC": 0x03,
}

AREA_CODES: dict[str, Any] = {
    "I": 0x81,
    "Q": 0x82,
    "M": 0x83,
    "DB": 0x84,
}

if Area is not None:
    AREA_CODES.update({"I": Area.PE, "Q": Area.PA, "M": Area.MK, "DB": Area.DB})


class S7CommunicationError(RuntimeError):
    def __init__(self, message: str) -> None:
        super().__init__(message)


class S7CommunicationTrace:
    def __init__(self) -> None:
        self._messages: list[str] = []

    def append(self, message: str) -> None:
        self._messages.append(message)

    def as_list(self) -> list[str]:
        return list(self._messages)


def device_payload(request: S7RawReadRequest | S7PointReadRequest) -> dict[str, Any]:
    return {
        "device_type": request.device_type,
        "ip": request.ip,
        "port": request.port,
        "rock": request.rock,
        "slot": request.slot,
        "tsap_conn_type": request.tsap_conn_type,
    }


def connect_scan_device_payload(request: S7ConnectScanRequest) -> dict[str, Any]:
    return {"device_type": "S7TCP", "ip": request.ip, "port": 102}


def connect_client(ip: str, port: int, rock: int, slot: int, tsap_conn_type: str) -> Client:
    client = Client()
    try:
        client.set_connection_type(CONNECTION_TYPES[tsap_conn_type])
        client.connect(ip, rock, slot, tcp_port=port)
    except TypeError:
        try:
            client.connect(ip, rock, slot, port)
        except Exception:
            close_client(client)
            raise
    except Exception:
        close_client(client)
        raise

    if not client.get_connected():
        close_client(client)
        raise S7CommunicationError(f"failed to connect to {ip}:{port}")
    return client


def close_client(client: Client | None) -> None:
    if client is None:
        return
    try:
        if client.get_connected():
            client.disconnect()
    finally:
        destroy = getattr(client, "destroy", None)
        if callable(destroy):
            destroy()


def pdu_length_message(client: Client) -> str:
    try:
        return f" pdu_length={client.get_pdu_length()}"
    except Exception:
        return ""


def read_s7_bytes(client: Client, read: S7ReadSpec | S7PointSpec, size: int) -> bytearray:
    if read.area == "DB":
        return client.db_read(read.db, read.start, size)
    return client.read_area(AREA_CODES[read.area], 0, read.start, size)


def trace_connect(trace: S7CommunicationTrace, request: S7RawReadRequest | S7PointReadRequest) -> None:
    trace.append(
        "Tx:S7_CONNECT "
        f"ip={request.ip} port={request.port} rock={request.rock} slot={request.slot} "
        f"tsap_conn_type={request.tsap_conn_type}"
    )


def trace_read(trace: S7CommunicationTrace, read: S7ReadSpec | S7PointSpec, size: int) -> None:
    trace.append(f"Tx:S7_READ area={read.area} db={read.db} start={read.start} size={size}")


def read_raw(request: S7RawReadRequest) -> dict[str, Any]:
    trace = S7CommunicationTrace()
    client: Client | None = None
    error: Exception | None = None

    try:
        trace_connect(trace, request)
        client = connect_client(request.ip, request.port, request.rock, request.slot, request.tsap_conn_type)
        trace.append(f"Rx:S7_CONNECTED{pdu_length_message(client)}")

        trace_read(trace, request.read, request.read.size)
        payload = read_s7_bytes(client, request.read, request.read.size)
        trace.append(f"Rx:{bytes_to_hex(payload)}")
    except Exception as exc:
        error = exc
        trace.append(f"Rx:S7_ERROR message={exc}")
    finally:
        if client is not None:
            trace.append("Tx:S7_DISCONNECT")
            close_client(client)
            trace.append("Rx:S7_DISCONNECTED")

    if error is not None:
        return response_payload(1, str(error), {"device": device_payload(request), "communication": trace.as_list()})
    return response_payload(0, "success", {"device": device_payload(request), "communication": trace.as_list()})


def read_points(request: S7PointReadRequest) -> dict[str, Any]:
    trace = S7CommunicationTrace()
    client: Client | None = None
    points: list[dict[str, Any]] = []
    error: Exception | None = None

    try:
        trace_connect(trace, request)
        client = connect_client(request.ip, request.port, request.rock, request.slot, request.tsap_conn_type)
        trace.append(f"Rx:S7_CONNECTED{pdu_length_message(client)}")

        for point in request.points:
            size = byte_size_for_type(point.type)
            trace_read(trace, point, size)
            payload = read_s7_bytes(client, point, size)
            trace.append(f"Rx:{bytes_to_hex(payload)}")

            result = {
                "area": point.area,
                "db": point.db,
                "start": point.start,
                "type": point.type,
                "value": convert_s7_value(payload, point.type, point.bit),
            }
            if point.bit is not None:
                result["bit"] = point.bit
            points.append(result)
    except Exception as exc:
        error = exc
        trace.append(f"Rx:S7_ERROR message={exc}")
    finally:
        if client is not None:
            trace.append("Tx:S7_DISCONNECT")
            close_client(client)
            trace.append("Rx:S7_DISCONNECTED")

    if error is not None:
        return response_payload(
            1,
            str(error),
            {"device": device_payload(request), "points": points},
        )
    return response_payload(
        0,
        "success",
        {"device": device_payload(request), "points": points},
    )


def connect_scan(request: S7ConnectScanRequest) -> dict[str, Any]:
    available: list[dict[str, Any]] = []

    for rock, slot, tsap_conn_type in product(range(3), range(3), CONNECTION_TYPES):
        client: Client | None = None
        try:
            client = connect_client(request.ip, 102, rock, slot, tsap_conn_type)
            available.append({"rock": rock, "slot": slot, "tsap_conn_type": tsap_conn_type})
        except Exception:
            pass
        finally:
            close_client(client)

    return response_payload(0, "success", {"device": connect_scan_device_payload(request), "available": available})