client.py

import logging
import struct
from typing import cast

from .base import BaseS7Comm
from .enums import HeaderError, ItemReturnCode, MessageType, SubfunctionCode, UserdataFunction, UserdataLastPDU
from .error_messages import ERROR_MESSAGES
from .exceptions import PacketLostError, StalePacketError
from .packets import (
    RequestPLCStop,
    S7AckDataHeader,
    S7Packet,
    SetupCommunicationRequest,
    UserDataContinuationRequest,
    UserDataRequest,
    UserDataResponse,
    VariableAddress,
    VariableReadRequest,
    VariableWriteRequest,
)
from .packets.exceptions import ReadVariableException
from .packets.parser import S7PacketParser
from .packets.rw_variable import ReadVariableResponse, WriteVariableResponse
from .transport import COTP, Transport


logger = logging.getLogger(__name__)


class S7Comm(BaseS7Comm):
    def __init__(
        self,
        tpdu_size: int = 1024,
        pdu_length: int = 480,
        source_tsap: int = 0x0100,
        dest_tsap: int = 0x0101,
        transport: Transport | None = None,
    ) -> None:
        super().__init__(pdu_length=pdu_length)
        self.transport: Transport
        if transport is None:
            self.transport = COTP(tpdu_size=tpdu_size, source_tsap=source_tsap, dest_tsap=dest_tsap)
        else:
            self.transport = transport

    def read_multi_vars(self, items: list[VariableAddress]) -> ReadVariableResponse:
        if len(items) > self.MAX_VARS:
            raise ValueError("Too many items")

        request = VariableReadRequest.create(items=items)
        response = self.send(request=request)
        if not isinstance(response, ReadVariableResponse):
            raise ValueError("Invalid response class")
        return response

    def write_multi_vars(self, items: list[tuple[VariableAddress, bytes]]) -> WriteVariableResponse:
        if len(items) > self.MAX_VARS:
            raise ValueError("Too many items")

        request = VariableWriteRequest.create(items=items)
        response = self.send(request=request)
        if not isinstance(response, WriteVariableResponse):
            raise ValueError("Invalid response class")
        return response

    def read_area(self, address: VariableAddress) -> bytes:
        result = bytes()
        main_request = VariableReadRequest.create(items=[address])

        for request in main_request.request_generator(pdu_length=self.pdu_length):
            response = self.send(request=request)
            if isinstance(response, ReadVariableResponse):
                for response_item in response.data:
                    if response_item.return_code != ItemReturnCode.SUCCESS:
                        raise ReadVariableException(
                            f"DataItem return code: {response_item.return_code}", response=response
                        )
                    result += response_item.data
        return result

    def write_area(self, address: VariableAddress, data: bytes) -> WriteVariableResponse:
        """
        Writes data to PLC memory area with automatic splitting into multiple requests
        if data exceeds PDU size.

        Args:
            address: Memory area address (e.g., "M0 INT 1")
            data: Data to write

        Returns:
            WriteVariableResponse with operation result
        """
        main_request = VariableWriteRequest.create(items=[(address, data)])
        for request in main_request.request_generator(pdu_length=self.pdu_length):
            response = self.send(request=request)
            if not isinstance(response, WriteVariableResponse):
                raise ValueError("Invalid response class")
            response.check_result()
            # TODO: combine all separate requests into one common initial request and return
        return cast(WriteVariableResponse, response)

    def connect(
        self,
        address: str,
        rack: int,
        slot: int,
        port: int = 102,
        max_amq_caller_ack: int = 0x0001,
        max_amq_callee_ack: int = 0x0001,
        pdu_length: int = 0x01E0,
    ) -> None:
        self.transport.connect(address=address, rack=rack, slot=slot, port=port)
        request = self._create_communication_request(
            max_amq_caller_ack=max_amq_caller_ack,
            max_amq_callee_ack=max_amq_callee_ack,
            pdu_length=pdu_length,
        )
        response = self.send(request=request)
        if not isinstance(response, SetupCommunicationRequest):
            raise ValueError("Invalid packet")
        self.pdu_length = response.parameter.pdu_length

    def _send_raw(self, request: S7Packet) -> None:
        packet = self._create_packet(request=request, pdu_reference=self.pdu_reference)
        self.transport.send(payload=packet)

    def send(self, request: S7Packet) -> S7Packet:
        self._send_raw(request)
        response = self._receive()

        # Validate PDU reference, retry receive if stale packet
        while True:
            if response.header is None:
                raise ValueError("Response header is None")
            try:
                self._validate_pdu_reference(response=response)
                break
            except StalePacketError:
                logger.warning("Dropping stale packet")
                response = self._receive()  # retry receive
                continue
            except PacketLostError:
                raise  # unrecoverable

        if response.header.message_type == MessageType.Userdata and isinstance(response, UserDataResponse):
            accumulated_data = response.data.data
            while response.parameter.last_data_unit == UserdataLastPDU.NO:
                continuation_request = UserDataContinuationRequest.from_response(response)
                self._send_raw(continuation_request)
                next_response = self._receive()
                if not isinstance(next_response, UserDataResponse):
                    raise ValueError("Expected UserDataResponse for continuation")
                accumulated_data += next_response.data.data
                response = next_response
            response.data.data = accumulated_data
        # Check Error
        if isinstance(response.header, S7AckDataHeader) and response.header.error_code != HeaderError.NO_ERROR:
            error_message = ERROR_MESSAGES.get(0x81 << 8 | 0x04, "Unknown Error")
            raise ValueError(
                f"Error class {response.header.error_class}, "
                f"error code: f{response.header.error_code}, "
                f"error message: {error_message}"
            )
        return response

    def _receive(self) -> S7Packet:
        """Receive and parse S7 packet from transport."""
        return S7PacketParser.parse(self.transport.receive())

    def close(self) -> None:
        self.transport.close()

    def read_szl(self, szl_id: int, szl_index: int) -> S7Packet:
        data = struct.pack("!HH", szl_id, szl_index)
        request = UserDataRequest.create(
            function_group=UserdataFunction.CPU_FUNCTION,
            subfunction=SubfunctionCode.READ_SZL,
            data=data,
        )
        response = self.send(request=request)
        return response

    def plc_stop(self) -> S7Packet:
        response = self.send(request=RequestPLCStop())
        return response