diff options
Diffstat (limited to 'python/libknot')
| -rw-r--r-- | python/libknot/__init__.py.in | 95 | ||||
| -rw-r--r-- | python/libknot/control.py | 374 | ||||
| -rw-r--r-- | python/libknot/dname.py | 70 | ||||
| -rw-r--r-- | python/libknot/probe.py | 278 |
4 files changed, 817 insertions, 0 deletions
diff --git a/python/libknot/__init__.py.in b/python/libknot/__init__.py.in new file mode 100644 index 0000000..554cbe6 --- /dev/null +++ b/python/libknot/__init__.py.in @@ -0,0 +1,95 @@ +"""Python libknot interface.""" + +import ctypes +import sys + + +class KnotLookup(ctypes.Structure): + """Libknot lookup return structure.""" + + _fields_ = [('id', ctypes.c_int), ('name', ctypes.c_char_p)] + + +class KnotRdataDescriptor(ctypes.Structure): + """Rdata descriptor structure.""" + + _fields_ = [('block_types', ctypes.c_int * 8), ('name', ctypes.c_char_p)] + + +class Knot(object): + """Basic libknot interface.""" + + LIBKNOT = None + LIBKNOT_VERSION = "@libknot_SOVERSION@" + + RCODE_NAMES = None + + STRERROR = None + RDATA_DESC = None + + @classmethod + def __init__(cls, path: str = None) -> None: + """Loads shared libknot library. + An explicit library path can be specified. + """ + + if cls.LIBKNOT: + return + + if path is None: + version = "" + try: + version = ".%u" % int(cls.LIBKNOT_VERSION) + except Exception: + pass + + if sys.platform == "darwin": + path = "libknot%s.dylib" % version + else: + path = "libknot.so%s" % version + + cls.LIBKNOT = ctypes.cdll.LoadLibrary(path) + + cls.RCODE_NAMES = (KnotLookup * 32).in_dll(cls.LIBKNOT, "knot_rcode_names") + + cls.STRERROR = cls.LIBKNOT.knot_strerror + cls.STRERROR.restype = ctypes.c_char_p + cls.STRERROR.argtypes = [ctypes.c_int] + + cls.RDATA_DESC = cls.LIBKNOT.knot_get_rdata_descriptor + cls.RDATA_DESC.restype = ctypes.POINTER(KnotRdataDescriptor) + cls.RDATA_DESC.argtypes = [ctypes.c_ushort] + + @classmethod + def rclass_str(cls, rclass: int) -> str: + """Returns RRCLASS in text form.""" + + if (rclass == 1): + return "IN" + elif (rclass == 3): + return "CH" + elif (rclass == 254): + return "NONE" + elif (rclass == 255): + return "ANY" + else: + return "CLASS%i" % rclass + + @classmethod + def rtype_str(cls, rtype: int) -> str: + """Returns RRTYPE in text form.""" + + descr = cls.RDATA_DESC(rtype).contents.name + if descr: + return descr.decode() + else: + return "TYPE%i" % rtype + + @classmethod + def rcode_str(cls, rcode: int) -> str: + """Returns RCODE in text form.""" + + for item in cls.RCODE_NAMES: + if item.name and item.id == rcode: + return item.name.decode() + return "RCODE%i" % rcode diff --git a/python/libknot/control.py b/python/libknot/control.py new file mode 100644 index 0000000..44aa516 --- /dev/null +++ b/python/libknot/control.py @@ -0,0 +1,374 @@ +"""Libknot server control interface wrapper.""" + +import ctypes +import enum +import warnings +import libknot + + +def load_lib(path: str = None) -> None: + """Compatibility wrapper.""" + + libknot.Knot(path) + warnings.warn("libknot.control.load_lib() is deprecated, use libknot.Knot() instead", \ + category=Warning, stacklevel=2) + + +class KnotCtlType(enum.IntEnum): + """Libknot server control data unit types.""" + + END = 0 + DATA = 1 + EXTRA = 2 + BLOCK = 3 + + +class KnotCtlDataIdx(enum.IntEnum): + """Libknot server control data unit indices.""" + + COMMAND = 0 + FLAGS = 1 + ERROR = 2 + SECTION = 3 + ITEM = 4 + ID = 5 + ZONE = 6 + OWNER = 7 + TTL = 8 + TYPE = 9 + DATA = 10 + FILTER = 11 + + +class KnotCtlData(object): + """Libknot server control data unit.""" + + DataArray = ctypes.c_char_p * len(KnotCtlDataIdx) + + def __init__(self) -> None: + self.data = self.DataArray() + + def __str__(self) -> str: + """Returns data unit in text form.""" + + string = str() + + for idx in KnotCtlDataIdx: + if self.data[idx]: + if string: + string += ", " + string += "%s = '%s'" % (idx.name, self.data[idx].decode()) + + return string + + def __getitem__(self, index: KnotCtlDataIdx) -> str: + """Data unit item getter.""" + + value = self.data[index] + return value.decode() if value else str() + + def __setitem__(self, index: KnotCtlDataIdx, value: str) -> None: + """Data unit item setter.""" + + self.data[index] = ctypes.c_char_p(value.encode()) if value else ctypes.c_char_p() + + +class KnotCtlError(Exception): + """Libknot server control error.""" + + def __init__(self, message: str, data: KnotCtlData = None) -> None: + super().__init__() + self.message = message + self.data = data + + def __str__(self) -> str: + out = "%s" % self.message + if self.data: + out += " (%s)" % self.data + return out + + +class KnotCtlErrorConnect(KnotCtlError): + """Control connection error.""" + + +class KnotCtlErrorSend(KnotCtlError): + """Control data send error.""" + + +class KnotCtlErrorReceive(KnotCtlError): + """Control data receive error.""" + + +class KnotCtlErrorRemote(KnotCtlError): + """Control error on the remote (server) side.""" + + +class KnotCtl(object): + """Libknot server control interface.""" + + ALLOC = None + FREE = None + SET_TIMEOUT = None + CONNECT = None + CLOSE = None + SEND = None + RECEIVE = None + + def __init__(self) -> None: + """Initializes a control interface instance.""" + + if not KnotCtl.ALLOC: + libknot.Knot() + + KnotCtl.ALLOC = libknot.Knot.LIBKNOT.knot_ctl_alloc + KnotCtl.ALLOC.restype = ctypes.c_void_p + + KnotCtl.FREE = libknot.Knot.LIBKNOT.knot_ctl_free + KnotCtl.FREE.argtypes = [ctypes.c_void_p] + + KnotCtl.SET_TIMEOUT = libknot.Knot.LIBKNOT.knot_ctl_set_timeout + KnotCtl.SET_TIMEOUT.argtypes = [ctypes.c_void_p, ctypes.c_int] + + KnotCtl.CONNECT = libknot.Knot.LIBKNOT.knot_ctl_connect + KnotCtl.CONNECT.restype = ctypes.c_int + KnotCtl.CONNECT.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + + KnotCtl.CLOSE = libknot.Knot.LIBKNOT.knot_ctl_close + KnotCtl.CLOSE.argtypes = [ctypes.c_void_p] + + KnotCtl.SEND = libknot.Knot.LIBKNOT.knot_ctl_send + KnotCtl.SEND.restype = ctypes.c_int + KnotCtl.SEND.argtypes = [ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p] + + KnotCtl.RECEIVE = libknot.Knot.LIBKNOT.knot_ctl_receive + KnotCtl.RECEIVE.restype = ctypes.c_int + KnotCtl.RECEIVE.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + + self.obj = KnotCtl.ALLOC() + + def __del__(self) -> None: + """Deallocates control interface instance.""" + + KnotCtl.FREE(self.obj) + + def set_timeout(self, timeout: int) -> None: + """Sets control socket operations timeout in seconds.""" + + KnotCtl.SET_TIMEOUT(self.obj, timeout * 1000) + + def connect(self, path: str) -> None: + """Connect to a specified control UNIX socket.""" + + ret = KnotCtl.CONNECT(self.obj, path.encode()) + if ret != 0: + err = libknot.Knot.STRERROR(ret) + raise KnotCtlErrorConnect(err.decode()) + + def close(self) -> None: + """Disconnects from the current control socket.""" + + KnotCtl.CLOSE(self.obj) + + def send(self, data_type: KnotCtlType, data: KnotCtlData = None) -> None: + """Sends a data unit to the connected control socket.""" + + ret = KnotCtl.SEND(self.obj, data_type, + data.data if data else ctypes.c_char_p()) + if ret != 0: + err = libknot.Knot.STRERROR(ret) + raise KnotCtlErrorSend(err.decode()) + + def receive(self, data: KnotCtlData = None) -> KnotCtlType: + """Receives a data unit from the connected control socket.""" + + data_type = ctypes.c_uint() + ret = KnotCtl.RECEIVE(self.obj, ctypes.byref(data_type), + data.data if data else ctypes.c_char_p()) + if ret != 0: + err = libknot.Knot.STRERROR(ret) + raise KnotCtlErrorReceive(err.decode()) + return KnotCtlType(data_type.value) + + def send_block(self, cmd: str, section: str = None, item: str = None, + identifier: str = None, zone: str = None, owner: str = None, + ttl: str = None, rtype: str = None, data: str = None, + flags: str = None, filters: str = None) -> None: + """Sends a control query block.""" + + query = KnotCtlData() + query[KnotCtlDataIdx.COMMAND] = cmd + query[KnotCtlDataIdx.SECTION] = section + query[KnotCtlDataIdx.ITEM] = item + query[KnotCtlDataIdx.ID] = identifier + query[KnotCtlDataIdx.ZONE] = zone + query[KnotCtlDataIdx.OWNER] = owner + query[KnotCtlDataIdx.TTL] = ttl + query[KnotCtlDataIdx.TYPE] = rtype + query[KnotCtlDataIdx.DATA] = data + query[KnotCtlDataIdx.FLAGS] = flags + query[KnotCtlDataIdx.FILTER] = filters + + self.send(KnotCtlType.DATA, query) + self.send(KnotCtlType.BLOCK) + + def _receive_conf(self, out, reply): + + section = reply[KnotCtlDataIdx.SECTION] + ident = reply[KnotCtlDataIdx.ID] + item = reply[KnotCtlDataIdx.ITEM] + data = reply[KnotCtlDataIdx.DATA] + + # Add the section if not exists. + if section not in out: + out[section] = dict() + + # Add the identifier if not exists. + if ident and ident not in out[section]: + out[section][ident] = dict() + + # Return if no item/value. + if not item: + return + + item_level = out[section][ident] if ident else out[section] + + # Treat alone identifier item differently. + if item in ["id", "domain", "target"]: + if data not in out[section]: + out[section][data] = dict() + else: + if item not in item_level: + item_level[item] = list() + + if data: + item_level[item].append(data) + + def _receive_zone_status(self, out, reply): + + zone = reply[KnotCtlDataIdx.ZONE] + rtype = reply[KnotCtlDataIdx.TYPE] + data = reply[KnotCtlDataIdx.DATA] + + # Add the zone if not exists. + if zone not in out: + out[zone] = dict() + + out[zone][rtype] = data + + def _receive_zone(self, out, reply): + + zone = reply[KnotCtlDataIdx.ZONE] + owner = reply[KnotCtlDataIdx.OWNER] + ttl = reply[KnotCtlDataIdx.TTL] + rtype = reply[KnotCtlDataIdx.TYPE] + data = reply[KnotCtlDataIdx.DATA] + + # Add the zone if not exists. + if zone not in out: + out[zone] = dict() + + if owner not in out[zone]: + out[zone][owner] = dict() + + if rtype not in out[zone][owner]: + out[zone][owner][rtype] = dict() + + # Add the key/value. + out[zone][owner][rtype]["ttl"] = ttl + + if not "data" in out[zone][owner][rtype]: + out[zone][owner][rtype]["data"] = [data] + else: + out[zone][owner][rtype]["data"].append(data) + + def _receive_stats(self, out, reply): + + zone = reply[KnotCtlDataIdx.ZONE] + section = reply[KnotCtlDataIdx.SECTION] + item = reply[KnotCtlDataIdx.ITEM] + idx = reply[KnotCtlDataIdx.ID] + data = int(reply[KnotCtlDataIdx.DATA]) + + # Add the zone if not exists. + if zone: + if "zone" not in out: + out["zone"] = dict() + + if zone not in out["zone"]: + out["zone"][zone] = dict() + + section_level = out["zone"][zone] if zone else out + + if section not in section_level: + section_level[section] = dict() + + if idx: + if item not in section_level[section]: + section_level[section][item] = dict() + + section_level[section][item][idx] = data + else: + section_level[section][item] = data + + def receive_stats(self) -> dict: + """Receives statistics answer and returns it as a structured dictionary.""" + + out = dict() + err_reply = None + + while True: + reply = KnotCtlData() + reply_type = self.receive(reply) + + # Stop if not data type. + if reply_type not in [KnotCtlType.DATA, KnotCtlType.EXTRA]: + break + + # Check for an error. + if reply[KnotCtlDataIdx.ERROR]: + err_reply = reply + continue + + self._receive_stats(out, reply) + + if err_reply: + raise KnotCtlErrorRemote(err_reply[KnotCtlDataIdx.ERROR], err_reply) + + return out + + def receive_block(self) -> dict: + """Receives a control answer and returns it as a structured dictionary.""" + + out = dict() + err_reply = None + + while True: + reply = KnotCtlData() + reply_type = self.receive(reply) + + # Stop if not data type. + if reply_type not in [KnotCtlType.DATA, KnotCtlType.EXTRA]: + break + + # Check for an error. + if reply[KnotCtlDataIdx.ERROR]: + err_reply = reply + continue + + # Check for config data. + if reply[KnotCtlDataIdx.SECTION]: + self._receive_conf(out, reply) + # Check for zone data. + elif reply[KnotCtlDataIdx.ZONE]: + if reply[KnotCtlDataIdx.OWNER]: + self._receive_zone(out, reply) + else: + self._receive_zone_status(out, reply) + else: + continue + + if err_reply: + raise KnotCtlErrorRemote(err_reply[KnotCtlDataIdx.ERROR], err_reply) + + return out diff --git a/python/libknot/dname.py b/python/libknot/dname.py new file mode 100644 index 0000000..4f585fe --- /dev/null +++ b/python/libknot/dname.py @@ -0,0 +1,70 @@ +"""Libknot dname interface wrapper.""" + +import ctypes +import libknot + + +class KnotDname(object): + """Libknot dname.""" + + CAPACITY = 255 + CAPACITY_TXT = 1004 + + DnameStorage = ctypes.c_char * CAPACITY + DnameTxtStorage = ctypes.c_char * CAPACITY_TXT + + SIZE = None + TO_STR = None + FROM_STR = None + + data = None + + def __init__(self, dname: str = None) -> None: + """Initializes a dname storage. Optionally initializes from a string.""" + + if not KnotDname.SIZE: + libknot.Knot() + + KnotDname.SIZE = libknot.Knot.LIBKNOT.knot_dname_size + KnotDname.SIZE.restype = ctypes.c_size_t + KnotDname.SIZE.argtypes = [KnotDname.DnameStorage] + + KnotDname.TO_STR = libknot.Knot.LIBKNOT.knot_dname_to_str + KnotDname.TO_STR.restype = ctypes.c_char_p + KnotDname.TO_STR.argtypes = [KnotDname.DnameTxtStorage, KnotDname.DnameStorage, ctypes.c_size_t] + + KnotDname.FROM_STR = libknot.Knot.LIBKNOT.knot_dname_from_str + KnotDname.FROM_STR.restype = ctypes.c_char_p + KnotDname.FROM_STR.argtypes = [KnotDname.DnameStorage, ctypes.c_char_p, ctypes.c_size_t] + + if dname: + self.data = KnotDname.DnameStorage() + if not KnotDname.FROM_STR(self.data, dname.encode('utf-8'), KnotDname.CAPACITY): + raise ValueError + + def size(self): + """Returns size of the stored dname.""" + + if self.data: + return KnotDname.SIZE(self.data) + else: + return 0 + + def str(self) -> str: + """Prints the stored dname in textual format.""" + + if self.data: + data_txt = KnotDname.DnameTxtStorage() + if not KnotDname.TO_STR(data_txt, self.data, KnotDname.CAPACITY_TXT): + raise ValueError + return data_txt.value.decode("utf-8") + else: + return "" + + def wire(self) -> bytes: + """Returns the dname in wire format.""" + + if self.data: + return self.data.value + b'\x00' + else: + return bytes() diff --git a/python/libknot/probe.py b/python/libknot/probe.py new file mode 100644 index 0000000..e6f09db --- /dev/null +++ b/python/libknot/probe.py @@ -0,0 +1,278 @@ +"""Libknot probe interface wrapper.""" + +import ctypes +import datetime +import enum +import socket +import libknot + + +class KnotProbeDataProto(enum.IntEnum): + """Libknot probe transport protocol types.""" + + UDP = 0 + TCP = 1 + QUIC = 3 + TLS = 4 + HTTPS = 5 + + +class KnotProbeDataDNSHdr(ctypes.BigEndianStructure): + """DNS message header.""" + + _fields_ = [('id', ctypes.c_ushort), + ('flag_qr', ctypes.c_ubyte, 1), + ('opcode', ctypes.c_ubyte, 4), + ('flag_aa', ctypes.c_ubyte, 1), + ('flag_tc', ctypes.c_ubyte, 1), + ('flag_rd', ctypes.c_ubyte, 1), + ('flag_ra', ctypes.c_ubyte, 1), + ('flag_z', ctypes.c_ubyte, 1), + ('flag_ad', ctypes.c_ubyte, 1), + ('flag_cd', ctypes.c_ubyte, 1), + ('rcode', ctypes.c_ubyte, 4), + ('questions', ctypes.c_ushort), + ('answers', ctypes.c_ushort), + ('authorities', ctypes.c_ushort), + ('additionals', ctypes.c_ushort)] + + +class KnotProbeData(ctypes.Structure): + """Libknot probe data unit.""" + + ADDR_MAX_SIZE = 16 + QNAME_MAX_SIZE = 255 + + EDE_NONE = 65535 + + _fields_ = [('ip', ctypes.c_ubyte), + ('proto', ctypes.c_ubyte), + ('local_addr', ctypes.c_ubyte * ADDR_MAX_SIZE), + ('local_port', ctypes.c_ushort), + ('remote_addr', ctypes.c_ubyte * ADDR_MAX_SIZE), + ('remote_port', ctypes.c_ushort), + ('reply_hdr', KnotProbeDataDNSHdr), + ('reply_size', ctypes.c_ushort), + ('reply_rcode', ctypes.c_ushort), + ('reply_ede', ctypes.c_ushort), + ('tcp_rtt', ctypes.c_uint), + ('edns_options', ctypes.c_uint), + ('edns_payload', ctypes.c_ushort), + ('edns_version', ctypes.c_ubyte), + ('edns_present', ctypes.c_ubyte, 1), + ('edns_flag_do', ctypes.c_ubyte, 1), + ('_reserved_', ctypes.c_ubyte, 6), + ('query_hdr', KnotProbeDataDNSHdr), + ('query_size', ctypes.c_ushort), + ('query_class', ctypes.c_ushort), + ('query_type', ctypes.c_ushort), + ('query_name_len', ctypes.c_ubyte), + ('query_name', ctypes.c_ubyte * (QNAME_MAX_SIZE))] + + def addr_str(self, addr: ctypes.c_ubyte * ADDR_MAX_SIZE) -> str: + """Converts IPv4 or IPv6 address from binary to text form.""" + + if self.ip == 4: + buffer = ctypes.create_string_buffer(4) + ctypes.memmove(buffer, ctypes.addressof(addr), 4) + return socket.inet_ntop(socket.AF_INET, buffer) + else: + return socket.inet_ntop(socket.AF_INET6, addr) + + def qname_str(self) -> str: + """Returns QNAME in text form.""" + + string = str() + pos = 0 + while pos < self.query_name_len: + label_len = self.query_name[pos] + if label_len == 0: + if self.query_name_len == 1: + string += "." + break + pos += 1 + label_end = pos + label_len + while pos < label_end: + string += chr(self.query_name[pos]) + pos += 1 + string += "." + return string + + def __str__(self) -> str: + """Returns the data unit in a pre-formatted text form.""" + + return self.str() + + def str(self, timestamp: bool = True, color: bool = True) -> str: + """Returns the data unit in a pre-formatted text form with customization.""" + + RST = "\x1B[0m" + BOLD = "\x1B[1m" + UNDR = "\x1B[4m" + RED = "\x1B[31m" + GRN = "\x1B[32m" + ORG = "\x1B[33m" + YELW = "\x1B[93m" + MGNT = "\x1B[35m" + CYAN = "\x1B[36m" + + def COL(string, color_str, active=color): + return str(string) if not active else color_str + str(string) + RST + + string = str() + if timestamp: + string += "%s " % COL(datetime.datetime.now().time(), YELW) + if self.ip != 0: + string += "%s -> %s, " % (COL(self.addr_str(self.remote_addr), UNDR), + COL(self.addr_str(self.local_addr), UNDR)) + string += "port %u -> %u " % (self.remote_port, self.local_port) + else: + string += "%s, " % COL("UNIX", UNDR) + if self.proto == KnotProbeDataProto.UDP: + string += COL("UDP", GRN) + elif self.proto == KnotProbeDataProto.TCP: + string += COL("TCP", RED) + else: + string += COL("QUIC", ORG) + if self.tcp_rtt > 0: + string += ", RTT %.2f ms" % (self.tcp_rtt / 1000) + string += "\n ID %u, " % self.query_hdr.id + if self.query_hdr.opcode == 0: + string += "QUERY" + elif self.query_hdr.opcode == 4: + string += COL("NOTIFY", MGNT) + elif self.query_hdr.opcode == 5: + string += COL("UPDATE", MGNT) + else: + string += COL("OPCODE%i" % self.query_hdr.opcode, MGNT) + string += ", " + string += COL("%s %s %s" % (self.qname_str(), + libknot.Knot.rclass_str(self.query_class), + libknot.Knot.rtype_str(self.query_type)), BOLD) + if self.edns_present == 1: + string += ", EDNS %i B" % self.edns_payload + if self.edns_flag_do == 1: + string += ", " + COL("DO", BOLD) + if (self.edns_options & (1 << 3)) != 0: + string += ", NSID" + if (self.edns_options & (1 << 8)) != 0: + string += ", ECS" + if (self.edns_options & (1 << 10)) != 0: + string += ", COOKIE" + string += ", " + COL("%u B" % self.query_size, CYAN) + if self.reply_size == 0: + string += " -> %s" % COL("DROPPED", RED) + return string + string += " -> %s" % COL(libknot.Knot.rcode_str(self.reply_rcode), BOLD) + if (self.reply_ede != libknot.probe.KnotProbeData.EDE_NONE): + string += ", EDE %u" % self.reply_ede + if self.reply_hdr.flag_aa != 0: + string += ", " + COL("AA", BOLD) + if self.reply_hdr.flag_tc != 0: + string += ", " + COL("TC", BOLD) + if self.reply_hdr.answers > 0: + string += ", %u ANS" % self.reply_hdr.answers + if self.reply_hdr.authorities > 0: + string += ", %u AUT" % self.reply_hdr.authorities + if self.reply_hdr.additionals > 0: + string += ", %u ADD" % self.reply_hdr.additionals + string += ", " + COL("%u B" % self.reply_size, CYAN) + return string + + +class KnotProbeDataArray(object): + """Libknot probe data unit array.""" + + def __init__(self, size: int = 1) -> None: + """Creates a data array of a given size.""" + + if size < 1 or size > 255: + raise ValueError + data_array = KnotProbeData * size + self.data = data_array() + self.capacity = size + self.used = 0 + self.pos = 0 + + def __getitem__(self, i: int) -> KnotProbeData: + """Returns a data unit at a specified position.""" + + if i < 0 or i >= self.capacity: + raise ValueError + return self.data[i] + + def __len__(self) -> int: + """Returns currently used size of the array.""" + + return self.used + + def __iter__(self): + """Initializes the array iterator.""" + + self.pos = 0 + return self + + def __next__(self) -> KnotProbeData: + """Increments the array iterator.""" + + if self.used == 0 or self.pos == self.used: + raise StopIteration + else: + data = self.data[self.pos] + self.pos += 1 + return data + + +class KnotProbe(object): + """Libknot probe consumer interface.""" + + ALLOC = None + FREE = None + CONSUME = None + SET_CONSUMER = None + + def __init__(self, path: str = "/run/knot", idx: int = 1) -> None: + """Initializes a probe channel at a specified path with a channel index.""" + + if not KnotProbe.ALLOC: + libknot.Knot() + + KnotProbe.ALLOC = libknot.Knot.LIBKNOT.knot_probe_alloc + KnotProbe.ALLOC.restype = ctypes.c_void_p + + KnotProbe.FREE = libknot.Knot.LIBKNOT.knot_probe_free + KnotProbe.FREE.argtypes = [ctypes.c_void_p] + + KnotProbe.CONSUME = libknot.Knot.LIBKNOT.knot_probe_consume + KnotProbe.CONSUME.restype = ctypes.c_int + KnotProbe.CONSUME.argtypes = [ctypes.c_void_p, ctypes.c_void_p, \ + ctypes.c_ubyte, ctypes.c_int] + + KnotProbe.SET_CONSUMER = libknot.Knot.LIBKNOT.knot_probe_set_consumer + KnotProbe.SET_CONSUMER.restype = ctypes.c_int + KnotProbe.SET_CONSUMER.argtypes = [ctypes.c_void_p, ctypes.c_char_p, \ + ctypes.c_ushort] + + self.obj = KnotProbe.ALLOC() + + ret = KnotProbe.SET_CONSUMER(self.obj, path.encode(), idx) + if ret != 0: + err = libknot.Knot.STRERROR(ret) + raise RuntimeError(err.decode()) + + def __del__(self) -> None: + """Deinitializes a probe channel.""" + + KnotProbe.FREE(self.obj) + + def consume(self, data: KnotProbeDataArray, timeout: int = 1000) -> int: + '''Consumes data units from a channel and stores them in data array. + Returns the number of consumed data units. + ''' + + ret = KnotProbe.CONSUME(self.obj, data.data, data.capacity, timeout) + if ret < 0: + err = libknot.Knot.STRERROR(ret) + raise RuntimeError(err.decode()) + data.used = ret + return ret |