user_data.py

import struct

from ..enums import (
    DataTransportSize,
    MessageType,
    SubfunctionCode,
    TransmissionType,
    UserdataFunction,
    UserdataLastPDU,
    UserdataMethod,
)
from .data_item import DataItem
from .packet import S7Packet


class UserDataRequestParameter:
    HEAD = b"\x00\x01\x12"

    def __init__(
        self,
        function_group: UserdataFunction,
        subfunction: SubfunctionCode,
        method: UserdataMethod = UserdataMethod.REQUEST,
        type: TransmissionType = TransmissionType.REQUEST,
        sequence_number: int = 0,
        length: int = 4,
    ):
        self.length = length
        self.method = method
        self.type = type  # 4 bits
        self.function_group = function_group  # 4 bits
        self.subfunction = subfunction
        self.sequence_number = sequence_number
        self.data_length = None

    def serialize(self) -> bytes:
        type_funcgroup = self.type << 4 | self.function_group
        param = struct.pack("BBBBB", self.length, self.method, type_funcgroup, self.subfunction, self.sequence_number)
        return self.HEAD + param


class UserDataContinuationParameter:
    """Parameter for requesting next data unit in multi-packet Userdata response."""

    HEAD = b"\x00\x01\x12"

    def __init__(
        self,
        function_group: UserdataFunction,
        subfunction: SubfunctionCode,
        sequence_number: int,
    ):
        self.length = 8  # Extended parameter length
        self.method = UserdataMethod.RESPONSE  # Uses RESPONSE method for continuation
        self.type = TransmissionType.REQUEST
        self.function_group = function_group
        self.subfunction = subfunction
        self.sequence_number = sequence_number
        self.data_unit_reference = 0
        self.last_data_unit = UserdataLastPDU.YES
        self.error_code = 0

    def serialize(self) -> bytes:
        type_funcgroup = self.type << 4 | self.function_group
        param = struct.pack(
            "!BBBBBBBH",
            self.length,
            self.method,
            type_funcgroup,
            self.subfunction,
            self.sequence_number,
            self.data_unit_reference,
            self.last_data_unit,
            self.error_code,
        )
        return self.HEAD + param


class UserDataResponseParameter(UserDataRequestParameter):
    def __init__(
        self,
        data_unit_reference: int,
        last_data_unit: UserdataLastPDU,
        error_code: int,
        method: UserdataMethod,
        type: TransmissionType,
        function_group: UserdataFunction,
        subfunction: SubfunctionCode,
        sequence_number: int = 0,
        length: int = 4,
    ):
        super().__init__(
            method=method,
            type=type,
            function_group=function_group,
            subfunction=subfunction,
            sequence_number=sequence_number,
            length=length,
        )
        self.data_unit_reference = data_unit_reference
        self.last_data_unit = last_data_unit
        self.error_code = error_code

    @classmethod
    def parse(cls, packet: bytes) -> "UserDataResponseParameter":
        _, length = struct.unpack_from("3sB", packet)
        (
            method,
            type_funcgroup,
            subfunction,
            sequence_number,
            data_unit_reference,
            last_data_unit,
            error_code,
        ) = struct.unpack_from("BBBBBBH", packet, offset=4)
        type_ = type_funcgroup >> 4
        function_group = type_funcgroup & 0x0F
        return cls(
            length=length,
            method=method,
            type=type_,
            function_group=function_group,
            subfunction=subfunction,
            sequence_number=sequence_number,
            data_unit_reference=data_unit_reference,
            last_data_unit=UserdataLastPDU(last_data_unit),
            error_code=error_code,
        )


class UserDataRequest(S7Packet):
    MESSAGE_TYPE = MessageType.Userdata

    def __init__(self, parameter: UserDataRequestParameter, data: DataItem) -> None:
        self.parameter = parameter
        self.data = data

    @classmethod
    def create(
        cls,
        function_group: UserdataFunction,
        subfunction: SubfunctionCode,
        method: UserdataMethod = UserdataMethod.REQUEST,
        type: TransmissionType = TransmissionType.REQUEST,
        sequence_number: int = 0,
        length: int = 4,
        transport_size: DataTransportSize = DataTransportSize.OCTET,
        return_code: int = 0xFF,
        data: bytes = b"",
    ) -> "UserDataRequest":
        parameter = UserDataRequestParameter(
            function_group=function_group,
            subfunction=subfunction,
            method=method,
            type=type,
            sequence_number=sequence_number,
            length=length,
        )
        data_item = DataItem(transport_size=transport_size, return_code=return_code, data_length=len(data), data=data)
        return cls(parameter=parameter, data=data_item)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return bytes(self.data.serialize())


class UserDataResponse(S7Packet):
    PARAMETER_LENGTH: int = 12

    def __init__(self, parameter: UserDataResponseParameter, data: DataItem) -> None:
        self.parameter = parameter
        self.data = data

    @classmethod
    def parse(cls, packet: bytes) -> "UserDataResponse":
        parameter = UserDataResponseParameter.parse(packet)
        data = DataItem.parse(packet[cls.PARAMETER_LENGTH :])
        return cls(parameter=parameter, data=data)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return bytes(self.data.serialize())


class UserDataContinuationRequest(S7Packet):
    """Request for next data unit in multi-packet Userdata response."""

    MESSAGE_TYPE = MessageType.Userdata

    def __init__(self, parameter: UserDataContinuationParameter) -> None:
        self.parameter = parameter
        # Minimal data: return_code=0x0a, transport_size=0, data_length=0
        self.data = DataItem(transport_size=DataTransportSize.NULL, data_length=0, return_code=0x0A)

    @classmethod
    def from_response(cls, response: "UserDataResponse") -> "UserDataContinuationRequest":
        """Create continuation request from a Userdata response that has more data."""
        parameter = UserDataContinuationParameter(
            function_group=response.parameter.function_group,
            subfunction=response.parameter.subfunction,
            sequence_number=response.parameter.sequence_number,
        )
        return cls(parameter=parameter)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return self.data.serialize()