xref: /freebsd/tests/atf_python/sys/netlink/attrs.py (revision 6563dcb6b1f57e51db63854f3774b52e672232ed) 
1import socket
2import struct
3from enum import Enum
4
5from atf_python.sys.netlink.utils import align4
6from atf_python.sys.netlink.utils import enum_or_int
7
8
9class NlAttr(object):
10    HDR_LEN = 4  # sizeof(struct nlattr)
11
12    def __init__(self, nla_type, data):
13        if isinstance(nla_type, Enum):
14            self._nla_type = nla_type.value
15            self._enum = nla_type
16        else:
17            self._nla_type = nla_type
18            self._enum = None
19        self.nla_list = []
20        self._data = data
21
22    @property
23    def nla_type(self):
24        return self._nla_type & 0x3FFF
25
26    @property
27    def nla_len(self):
28        return len(self._data) + 4
29
30    def add_nla(self, nla):
31        self.nla_list.append(nla)
32
33    def print_attr(self, prepend=""):
34        if self._enum is not None:
35            type_str = self._enum.name
36        else:
37            type_str = "nla#{}".format(self.nla_type)
38        print(
39            "{}len={} type={}({}){}".format(
40                prepend, self.nla_len, type_str, self.nla_type, self._print_attr_value()
41            )
42        )
43
44    @staticmethod
45    def _validate(data):
46        if len(data) < 4:
47            raise ValueError("attribute too short")
48        nla_len, nla_type = struct.unpack("@HH", data[:4])
49        if nla_len > len(data):
50            raise ValueError("attribute length too big")
51        if nla_len < 4:
52            raise ValueError("attribute length too short")
53
54    @classmethod
55    def _parse(cls, data):
56        nla_len, nla_type = struct.unpack("@HH", data[:4])
57        return cls(nla_type, data[4:])
58
59    @classmethod
60    def from_bytes(cls, data, attr_type_enum=None):
61        cls._validate(data)
62        attr = cls._parse(data)
63        attr._enum = attr_type_enum
64        return attr
65
66    def _to_bytes(self, data: bytes):
67        ret = data
68        if align4(len(ret)) != len(ret):
69            ret = data + bytes(align4(len(ret)) - len(ret))
70        return struct.pack("@HH", len(data) + 4, self._nla_type) + ret
71
72    def __bytes__(self):
73        return self._to_bytes(self._data)
74
75    def _print_attr_value(self):
76        return " " + " ".join(["x{:02X}".format(b) for b in self._data])
77
78
79class NlAttrNested(NlAttr):
80    def __init__(self, nla_type, val):
81        super().__init__(nla_type, b"")
82        self.nla_list = val
83
84    def get_nla(self, nla_type):
85        nla_type_raw = enum_or_int(nla_type)
86        for nla in self.nla_list:
87            if nla.nla_type == nla_type_raw:
88                return nla
89        return None
90
91    @property
92    def nla_len(self):
93        return align4(len(b"".join([bytes(nla) for nla in self.nla_list]))) + 4
94
95    def print_attr(self, prepend=""):
96        if self._enum is not None:
97            type_str = self._enum.name
98        else:
99            type_str = "nla#{}".format(self.nla_type)
100        print(
101            "{}len={} type={}({}) {{".format(
102                prepend, self.nla_len, type_str, self.nla_type
103            )
104        )
105        for nla in self.nla_list:
106            nla.print_attr(prepend + "  ")
107        print("{}}}".format(prepend))
108
109    def __bytes__(self):
110        return self._to_bytes(b"".join([bytes(nla) for nla in self.nla_list]))
111
112
113class NlAttrU32(NlAttr):
114    def __init__(self, nla_type, val):
115        self.u32 = enum_or_int(val)
116        super().__init__(nla_type, b"")
117
118    @property
119    def nla_len(self):
120        return 8
121
122    def _print_attr_value(self):
123        return " val={}".format(self.u32)
124
125    @staticmethod
126    def _validate(data):
127        assert len(data) == 8
128        nla_len, nla_type = struct.unpack("@HH", data[:4])
129        assert nla_len == 8
130
131    @classmethod
132    def _parse(cls, data):
133        nla_len, nla_type, val = struct.unpack("@HHI", data)
134        return cls(nla_type, val)
135
136    def __bytes__(self):
137        return self._to_bytes(struct.pack("@I", self.u32))
138
139
140class NlAttrU64(NlAttr):
141    def __init__(self, nla_type, val):
142        self.u64 = enum_or_int(val)
143        super().__init__(nla_type, b"")
144
145    @property
146    def nla_len(self):
147        return 12
148
149    def _print_attr_value(self):
150        return " val={}".format(self.u64)
151
152    @staticmethod
153    def _validate(data):
154        assert len(data) == 12
155        nla_len, nla_type = struct.unpack("@HH", data[:4])
156        assert nla_len == 12
157
158    @classmethod
159    def _parse(cls, data):
160        nla_len, nla_type = struct.unpack("@HH", data[:4])
161        val = struct.unpack("@Q", data[4:])[0]
162        return cls(nla_type, val)
163
164    def __bytes__(self):
165        return self._to_bytes(struct.pack("@Q", self.u64))
166
167
168class NlAttrS32(NlAttr):
169    def __init__(self, nla_type, val):
170        self.s32 = enum_or_int(val)
171        super().__init__(nla_type, b"")
172
173    @property
174    def nla_len(self):
175        return 8
176
177    def _print_attr_value(self):
178        return " val={}".format(self.s32)
179
180    @staticmethod
181    def _validate(data):
182        assert len(data) == 8
183        nla_len, nla_type = struct.unpack("@HH", data[:4])
184        assert nla_len == 8
185
186    @classmethod
187    def _parse(cls, data):
188        nla_len, nla_type, val = struct.unpack("@HHi", data)
189        return cls(nla_type, val)
190
191    def __bytes__(self):
192        return self._to_bytes(struct.pack("@i", self.s32))
193
194
195class NlAttrU16(NlAttr):
196    def __init__(self, nla_type, val):
197        self.u16 = enum_or_int(val)
198        super().__init__(nla_type, b"")
199
200    @property
201    def nla_len(self):
202        return 6
203
204    def _print_attr_value(self):
205        return " val={}".format(self.u16)
206
207    @staticmethod
208    def _validate(data):
209        assert len(data) == 6
210        nla_len, nla_type = struct.unpack("@HH", data[:4])
211        assert nla_len == 6
212
213    @classmethod
214    def _parse(cls, data):
215        nla_len, nla_type, val = struct.unpack("@HHH", data)
216        return cls(nla_type, val)
217
218    def __bytes__(self):
219        return self._to_bytes(struct.pack("@H", self.u16))
220
221
222class NlAttrU8(NlAttr):
223    def __init__(self, nla_type, val):
224        self.u8 = enum_or_int(val)
225        super().__init__(nla_type, b"")
226
227    @property
228    def nla_len(self):
229        return 5
230
231    def _print_attr_value(self):
232        return " val={}".format(self.u8)
233
234    @staticmethod
235    def _validate(data):
236        assert len(data) == 5
237        nla_len, nla_type = struct.unpack("@HH", data[:4])
238        assert nla_len == 5
239
240    @classmethod
241    def _parse(cls, data):
242        nla_len, nla_type, val = struct.unpack("@HHB", data)
243        return cls(nla_type, val)
244
245    def __bytes__(self):
246        return self._to_bytes(struct.pack("@B", self.u8))
247
248
249class NlAttrIp(NlAttr):
250    def __init__(self, nla_type, addr: str):
251        super().__init__(nla_type, b"")
252        self.addr = addr
253        if ":" in self.addr:
254            self.family = socket.AF_INET6
255        else:
256            self.family = socket.AF_INET
257
258    @staticmethod
259    def _validate(data):
260        nla_len, nla_type = struct.unpack("@HH", data[:4])
261        data_len = nla_len - 4
262        if data_len != 4 and data_len != 16:
263            raise ValueError(
264                "Error validating attr {}: nla_len is not valid".format(  # noqa: E501
265                    nla_type
266                )
267            )
268
269    @property
270    def nla_len(self):
271        if self.family == socket.AF_INET6:
272            return 20
273        else:
274            return 8
275        return align4(len(self._data)) + 4
276
277    @classmethod
278    def _parse(cls, data):
279        nla_len, nla_type = struct.unpack("@HH", data[:4])
280        data_len = len(data) - 4
281        if data_len == 4:
282            addr = socket.inet_ntop(socket.AF_INET, data[4:8])
283        else:
284            addr = socket.inet_ntop(socket.AF_INET6, data[4:20])
285        return cls(nla_type, addr)
286
287    def __bytes__(self):
288        return self._to_bytes(socket.inet_pton(self.family, self.addr))
289
290    def _print_attr_value(self):
291        return " addr={}".format(self.addr)
292
293
294class NlAttrIp4(NlAttrIp):
295    def __init__(self, nla_type, addr: str):
296        super().__init__(nla_type, addr)
297        assert self.family == socket.AF_INET
298
299
300class NlAttrIp6(NlAttrIp):
301    def __init__(self, nla_type, addr: str):
302        super().__init__(nla_type, addr)
303        assert self.family == socket.AF_INET6
304
305
306class NlAttrStr(NlAttr):
307    def __init__(self, nla_type, text):
308        super().__init__(nla_type, b"")
309        self.text = text
310
311    @staticmethod
312    def _validate(data):
313        NlAttr._validate(data)
314        try:
315            data[4:].decode("utf-8")
316        except Exception as e:
317            raise ValueError("wrong utf-8 string: {}".format(e))
318
319    @property
320    def nla_len(self):
321        return len(self.text) + 5
322
323    @classmethod
324    def _parse(cls, data):
325        text = data[4:-1].decode("utf-8")
326        nla_len, nla_type = struct.unpack("@HH", data[:4])
327        return cls(nla_type, text)
328
329    def __bytes__(self):
330        return self._to_bytes(bytes(self.text, encoding="utf-8") + bytes(1))
331
332    def _print_attr_value(self):
333        return ' val="{}"'.format(self.text)
334
335
336class NlAttrStrn(NlAttr):
337    def __init__(self, nla_type, text):
338        super().__init__(nla_type, b"")
339        self.text = text
340
341    @staticmethod
342    def _validate(data):
343        NlAttr._validate(data)
344        try:
345            data[4:].decode("utf-8")
346        except Exception as e:
347            raise ValueError("wrong utf-8 string: {}".format(e))
348
349    @property
350    def nla_len(self):
351        return len(self.text) + 4
352
353    @classmethod
354    def _parse(cls, data):
355        text = data[4:].decode("utf-8")
356        nla_len, nla_type = struct.unpack("@HH", data[:4])
357        return cls(nla_type, text)
358
359    def __bytes__(self):
360        return self._to_bytes(bytes(self.text, encoding="utf-8"))
361
362    def _print_attr_value(self):
363        return ' val="{}"'.format(self.text)
364