construct_typing_typed_tcp_header_dataclass_build_parse.py

1import dataclasses
2import construct
3from construct_typing import DataclassStruct, DataclassBitStruct, TBitStructField, TStructField
4
5@dataclasses.dataclass
6class TcpHeader:
7    source_port: int = TStructField(construct.Int16ub)
8    destination_port: int = TStructField(construct.Int16ub)
9    sequence_number: int = TStructField(construct.Int32ub)
10    acknowledgment_number: int = TStructField(construct.Int32ub)
11    # The following is a bitstruct
12    @dataclasses.dataclass
13    class Flags:
14        data_offset: int = TBitStructField(construct.BitsInteger(4))
15        reserved: int = TBitStructField(construct.BitsInteger(3))
16        ns: int = TBitStructField(construct.Flag)
17        cwr: int = TBitStructField(construct.Flag)
18        ece: int = TBitStructField(construct.Flag)
19        urg: int = TBitStructField(construct.Flag)
20        ack: int = TBitStructField(construct.Flag)
21        psh: int = TBitStructField(construct.Flag)
22        rst: int = TBitStructField(construct.Flag)
23        syn: int = TBitStructField(construct.Flag)
24        fin: int = TBitStructField(construct.Flag)
25
26    flags: Flags = TStructField(DataclassBitStruct(Flags))
27    window_size: int = TStructField(construct.Int16ub)
28    checksum: int = TStructField(construct.Int16ub)
29    urgent_pointer: int = TStructField(construct.Int16ub)
30
31# Create the construct
32tcp_header_struct = DataclassStruct(TcpHeader)
33
34# Example: Build from object
35header_obj = TcpHeader(
36    source_port=80,
37    destination_port=12345,
38    sequence_number=1,
39    acknowledgment_number=0,
40    flags=TcpHeader.Flags(
41        data_offset=5, reserved=0, ns=0, cwr=0, ece=0,
42        urg=0, ack=0, psh=0, rst=0, syn=1, fin=0
43    ),
44    window_size=8192,
45    checksum=0,
46    urgent_pointer=0
47)
48
49data = tcp_header_struct.build(header_obj)
50print(f"Serialized data: {data.hex()}")
51
52# Example: Parse from bytes
53parsed_obj = tcp_header_struct.parse(data)
54print(f"Parsed Source Port: {parsed_obj.source_port}")