1#!/usr/local/bin/python3
2import os
3import socket
4import sys
5from ctypes import c_int
6from ctypes import c_ubyte
7from ctypes import c_uint
8from ctypes import c_ushort
9from ctypes import sizeof
10from ctypes import Structure
11from enum import auto
12from enum import Enum
13
14from atf_python.sys.netlink.attrs import NlAttr
15from atf_python.sys.netlink.attrs import NlAttrStr
16from atf_python.sys.netlink.attrs import NlAttrU32
17from atf_python.sys.netlink.base_headers import GenlMsgHdr
18from atf_python.sys.netlink.base_headers import NlmBaseFlags
19from atf_python.sys.netlink.base_headers import Nlmsghdr
20from atf_python.sys.netlink.base_headers import NlMsgType
21from atf_python.sys.netlink.message import BaseNetlinkMessage
22from atf_python.sys.netlink.message import NlMsgCategory
23from atf_python.sys.netlink.message import NlMsgProps
24from atf_python.sys.netlink.message import StdNetlinkMessage
25from atf_python.sys.netlink.netlink_generic import GenlCtrlAttrType
26from atf_python.sys.netlink.netlink_generic import GenlCtrlMsgType
27from atf_python.sys.netlink.netlink_generic import handler_classes as genl_classes
28from atf_python.sys.netlink.netlink_route import handler_classes as rt_classes
29from atf_python.sys.netlink.utils import align4
30from atf_python.sys.netlink.utils import AttrDescr
31from atf_python.sys.netlink.utils import build_propmap
32from atf_python.sys.netlink.utils import enum_or_int
33from atf_python.sys.netlink.utils import get_bitmask_map
34from atf_python.sys.netlink.utils import NlConst
35from atf_python.sys.netlink.utils import prepare_attrs_map
36
37
38class SockaddrNl(Structure):
39    _fields_ = [
40        ("nl_len", c_ubyte),
41        ("nl_family", c_ubyte),
42        ("nl_pad", c_ushort),
43        ("nl_pid", c_uint),
44        ("nl_groups", c_uint),
45    ]
46
47
48class Nlmsgdone(Structure):
49    _fields_ = [
50        ("error", c_int),
51    ]
52
53
54class Nlmsgerr(Structure):
55    _fields_ = [
56        ("error", c_int),
57        ("msg", Nlmsghdr),
58    ]
59
60
61class NlErrattrType(Enum):
62    NLMSGERR_ATTR_UNUSED = 0
63    NLMSGERR_ATTR_MSG = auto()
64    NLMSGERR_ATTR_OFFS = auto()
65    NLMSGERR_ATTR_COOKIE = auto()
66    NLMSGERR_ATTR_POLICY = auto()
67
68
69class AddressFamilyLinux(Enum):
70    AF_INET = socket.AF_INET
71    AF_INET6 = socket.AF_INET6
72    AF_NETLINK = 16
73
74
75class AddressFamilyBsd(Enum):
76    AF_INET = socket.AF_INET
77    AF_INET6 = socket.AF_INET6
78    AF_NETLINK = 38
79
80
81class NlHelper:
82    def __init__(self):
83        self._pmap = {}
84        self._af_cls = self.get_af_cls()
85        self._seq_counter = 1
86        self.pid = os.getpid()
87
88    def get_seq(self):
89        ret = self._seq_counter
90        self._seq_counter += 1
91        return ret
92
93    def get_af_cls(self):
94        if sys.platform.startswith("freebsd"):
95            cls = AddressFamilyBsd
96        else:
97            cls = AddressFamilyLinux
98        return cls
99
100    def get_propmap(self, cls):
101        if cls not in self._pmap:
102            self._pmap[cls] = build_propmap(cls)
103        return self._pmap[cls]
104
105    def get_name_propmap(self, cls):
106        ret = {}
107        for prop in dir(cls):
108            if not prop.startswith("_"):
109                ret[prop] = getattr(cls, prop).value
110        return ret
111
112    def get_attr_byval(self, cls, attr_val):
113        propmap = self.get_propmap(cls)
114        return propmap.get(attr_val)
115
116    def get_af_name(self, family):
117        v = self.get_attr_byval(self._af_cls, family)
118        if v is not None:
119            return v
120        return "af#{}".format(family)
121
122    def get_af_value(self, family_str: str) -> int:
123        propmap = self.get_name_propmap(self._af_cls)
124        return propmap.get(family_str)
125
126    def get_bitmask_str(self, cls, val):
127        bmap = get_bitmask_map(self.get_propmap(cls), val)
128        return ",".join([v for k, v in bmap.items()])
129
130    @staticmethod
131    def get_bitmask_str_uncached(cls, val):
132        pmap = NlHelper.build_propmap(cls)
133        bmap = NlHelper.get_bitmask_map(pmap, val)
134        return ",".join([v for k, v in bmap.items()])
135
136
137nldone_attrs = prepare_attrs_map([])
138
139nlerr_attrs = prepare_attrs_map(
140    [
141        AttrDescr(NlErrattrType.NLMSGERR_ATTR_MSG, NlAttrStr),
142        AttrDescr(NlErrattrType.NLMSGERR_ATTR_OFFS, NlAttrU32),
143        AttrDescr(NlErrattrType.NLMSGERR_ATTR_COOKIE, NlAttr),
144    ]
145)
146
147
148class NetlinkDoneMessage(StdNetlinkMessage):
149    messages = [NlMsgProps(NlMsgType.NLMSG_DONE, NlMsgCategory.ACK)]
150    nl_attrs_map = nldone_attrs
151
152    @property
153    def error_code(self):
154        return self.base_hdr.error
155
156    def parse_base_header(self, data):
157        if len(data) < sizeof(Nlmsgdone):
158            raise ValueError("length less than nlmsgdone header")
159        done_hdr = Nlmsgdone.from_buffer_copy(data)
160        sz = sizeof(Nlmsgdone)
161        return (done_hdr, sz)
162
163    def print_base_header(self, hdr, prepend=""):
164        print("{}error={}".format(prepend, hdr.error))
165
166
167class NetlinkErrorMessage(StdNetlinkMessage):
168    messages = [NlMsgProps(NlMsgType.NLMSG_ERROR, NlMsgCategory.ACK)]
169    nl_attrs_map = nlerr_attrs
170
171    @property
172    def error_code(self):
173        return self.base_hdr.error
174
175    @property
176    def error_str(self):
177        nla = self.get_nla(NlErrattrType.NLMSGERR_ATTR_MSG)
178        if nla:
179            return nla.text
180        return None
181
182    @property
183    def error_offset(self):
184        nla = self.get_nla(NlErrattrType.NLMSGERR_ATTR_OFFS)
185        if nla:
186            return nla.u32
187        return None
188
189    @property
190    def cookie(self):
191        return self.get_nla(NlErrattrType.NLMSGERR_ATTR_COOKIE)
192
193    def parse_base_header(self, data):
194        if len(data) < sizeof(Nlmsgerr):
195            raise ValueError("length less than nlmsgerr header")
196        err_hdr = Nlmsgerr.from_buffer_copy(data)
197        sz = sizeof(Nlmsgerr)
198        if (self.nl_hdr.nlmsg_flags & 0x100) == 0:
199            sz += align4(err_hdr.msg.nlmsg_len - sizeof(Nlmsghdr))
200        return (err_hdr, sz)
201
202    def print_base_header(self, errhdr, prepend=""):
203        print("{}error={}, ".format(prepend, errhdr.error), end="")
204        hdr = errhdr.msg
205        print(
206            "{}len={}, type={}, flags={}(0x{:X}), seq={}, pid={}".format(
207                prepend,
208                hdr.nlmsg_len,
209                "msg#{}".format(hdr.nlmsg_type),
210                self.helper.get_bitmask_str(NlmBaseFlags, hdr.nlmsg_flags),
211                hdr.nlmsg_flags,
212                hdr.nlmsg_seq,
213                hdr.nlmsg_pid,
214            )
215        )
216
217
218core_classes = {
219    "netlink_core": [
220        NetlinkDoneMessage,
221        NetlinkErrorMessage,
222    ],
223}
224
225
226class Nlsock:
227    HANDLER_CLASSES = [core_classes, rt_classes, genl_classes]
228
229    def __init__(self, family, helper):
230        self.helper = helper
231        self.sock_fd = self._setup_netlink(family)
232        self._sock_family = family
233        self._data = bytes()
234        self.msgmap = self.build_msgmap()
235        self._family_map = {
236            NlConst.GENL_ID_CTRL: "nlctrl",
237        }
238
239    def build_msgmap(self):
240        handler_classes = {}
241        for d in self.HANDLER_CLASSES:
242            handler_classes.update(d)
243        xmap = {}
244        # 'family_name': [class.messages[MsgProps.msg],  ]
245        for family_id, family_classes in handler_classes.items():
246            xmap[family_id] = {}
247            for cls in family_classes:
248                for msg_props in cls.messages:
249                    xmap[family_id][enum_or_int(msg_props.msg)] = cls
250        return xmap
251
252    def _setup_netlink(self, netlink_family) -> int:
253        family = self.helper.get_af_value("AF_NETLINK")
254        s = socket.socket(family, socket.SOCK_RAW, netlink_family)
255        s.setsockopt(270, 10, 1)  # NETLINK_CAP_ACK
256        s.setsockopt(270, 11, 1)  # NETLINK_EXT_ACK
257        return s
258
259    def set_groups(self, mask: int):
260        self.sock_fd.setsockopt(socket.SOL_SOCKET, 1, mask)
261        # snl = SockaddrNl(nl_len = sizeof(SockaddrNl), nl_family=38,
262        #                  nl_pid=self.pid, nl_groups=mask)
263        # xbuffer = create_string_buffer(sizeof(SockaddrNl))
264        # memmove(xbuffer, addressof(snl), sizeof(SockaddrNl))
265        # k = struct.pack("@BBHII", 12, 38, 0, self.pid, mask)
266        # self.sock_fd.bind(k)
267
268    def write_message(self, msg, verbose=True):
269        if verbose:
270            print("vvvvvvvv OUT vvvvvvvv")
271            msg.print_message()
272        msg_bytes = bytes(msg)
273        try:
274            ret = os.write(self.sock_fd.fileno(), msg_bytes)
275            assert ret == len(msg_bytes)
276        except Exception as e:
277            print("write({}) -> {}".format(len(msg_bytes), e))
278
279    def parse_message(self, data: bytes):
280        if len(data) < sizeof(Nlmsghdr):
281            raise Exception("Short read from nl: {} bytes".format(len(data)))
282        hdr = Nlmsghdr.from_buffer_copy(data)
283        if hdr.nlmsg_type < 16:
284            family_name = "netlink_core"
285            nlmsg_type = hdr.nlmsg_type
286        elif self._sock_family == NlConst.NETLINK_ROUTE:
287            family_name = "netlink_route"
288            nlmsg_type = hdr.nlmsg_type
289        else:
290            # Genetlink
291            if len(data) < sizeof(Nlmsghdr) + sizeof(GenlMsgHdr):
292                raise Exception("Short read from genl: {} bytes".format(len(data)))
293            family_name = self._family_map.get(hdr.nlmsg_type, "")
294            ghdr = GenlMsgHdr.from_buffer_copy(data[sizeof(Nlmsghdr):])
295            nlmsg_type = ghdr.cmd
296        cls = self.msgmap.get(family_name, {}).get(nlmsg_type)
297        if not cls:
298            cls = BaseNetlinkMessage
299        return cls.from_bytes(self.helper, data)
300
301    def get_genl_family_id(self, family_name):
302        hdr = Nlmsghdr(
303            nlmsg_type=NlConst.GENL_ID_CTRL,
304            nlmsg_flags=NlmBaseFlags.NLM_F_REQUEST.value,
305            nlmsg_seq = self.helper.get_seq(),
306        )
307        ghdr = GenlMsgHdr(cmd=GenlCtrlMsgType.CTRL_CMD_GETFAMILY.value)
308        nla = NlAttrStr(GenlCtrlAttrType.CTRL_ATTR_FAMILY_NAME, family_name)
309        hdr.nlmsg_len = sizeof(Nlmsghdr) + sizeof(GenlMsgHdr) + len(bytes(nla))
310
311        msg_bytes = bytes(hdr) + bytes(ghdr) + bytes(nla)
312        self.write_data(msg_bytes)
313        while True:
314            rx_msg = self.read_message()
315            if hdr.nlmsg_seq == rx_msg.nl_hdr.nlmsg_seq:
316                if rx_msg.is_type(NlMsgType.NLMSG_ERROR):
317                    if rx_msg.error_code != 0:
318                        raise ValueError("unable to get family {}".format(family_name))
319                else:
320                    family_id = rx_msg.get_nla(GenlCtrlAttrType.CTRL_ATTR_FAMILY_ID).u16
321                    self._family_map[family_id] = family_name
322                    return family_id
323        raise ValueError("unable to get family {}".format(family_name))
324
325    def write_data(self, data: bytes):
326        self.sock_fd.send(data)
327
328    def read_data(self):
329        while True:
330            data = self.sock_fd.recv(65535)
331            self._data += data
332            if len(self._data) >= sizeof(Nlmsghdr):
333                break
334
335    def read_message(self) -> bytes:
336        if len(self._data) < sizeof(Nlmsghdr):
337            self.read_data()
338        hdr = Nlmsghdr.from_buffer_copy(self._data)
339        while hdr.nlmsg_len > len(self._data):
340            self.read_data()
341        raw_msg = self._data[: hdr.nlmsg_len]
342        self._data = self._data[hdr.nlmsg_len:]
343        return self.parse_message(raw_msg)
344
345
346class NetlinkMultipartIterator(object):
347    def __init__(self, obj, seq_number: int, msg_type):
348        self._obj = obj
349        self._seq = seq_number
350        self._msg_type = msg_type
351
352    def __iter__(self):
353        return self
354
355    def __next__(self):
356        msg = self._obj.read_message()
357        if self._seq != msg.nl_hdr.nlmsg_seq:
358            raise ValueError("bad sequence number")
359        if msg.is_type(NlMsgType.NLMSG_ERROR):
360            raise ValueError(
361                "error while handling multipart msg: {}".format(msg.error_code)
362            )
363        elif msg.is_type(NlMsgType.NLMSG_DONE):
364            if msg.error_code == 0:
365                raise StopIteration
366            raise ValueError(
367                "error listing some parts of the multipart msg: {}".format(
368                    msg.error_code
369                )
370            )
371        elif not msg.is_type(self._msg_type):
372            raise ValueError("bad message type: {}".format(msg))
373        return msg
374
375
376class NetlinkTestTemplate(object):
377    REQUIRED_MODULES = ["netlink"]
378
379    def setup_netlink(self, netlink_family: NlConst):
380        self.helper = NlHelper()
381        self.nlsock = Nlsock(netlink_family, self.helper)
382
383    def write_message(self, msg, silent=False):
384        if not silent:
385            print("")
386            print("============= >> TX MESSAGE =============")
387            msg.print_message()
388            msg.print_as_bytes(bytes(msg), "-- DATA --")
389        self.nlsock.write_data(bytes(msg))
390
391    def read_message(self, silent=False):
392        msg = self.nlsock.read_message()
393        if not silent:
394            print("")
395            print("============= << RX MESSAGE =============")
396            msg.print_message()
397        return msg
398
399    def get_reply(self, tx_msg):
400        self.write_message(tx_msg)
401        while True:
402            rx_msg = self.read_message()
403            if tx_msg.nl_hdr.nlmsg_seq == rx_msg.nl_hdr.nlmsg_seq:
404                return rx_msg
405
406    def read_msg_list(self, seq, msg_type):
407        return list(NetlinkMultipartIterator(self, seq, msg_type))
408