From: Daniel Salzman <daniel.salzman@nic.cz>
Date: Fri, 19 Nov 2021 10:21:15 +0000 (+0100)
Subject: python: various probe improvements
X-Git-Tag: v3.3.dev~271
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=9bf810a2c5b164de5c07d005ab31dca3d6506bce;p=thirdparty%2Fknot-dns.git

python: various probe improvements
---

diff --git a/python/libknot/__init__.py.in b/python/libknot/__init__.py.in
index 6faa5f5e91..554cbe6f69 100644
--- a/python/libknot/__init__.py.in
+++ b/python/libknot/__init__.py.in
@@ -60,6 +60,21 @@ class Knot(object):
         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."""
diff --git a/python/libknot/probe.py b/python/libknot/probe.py
index 08565cd866..97dd3861c9 100644
--- a/python/libknot/probe.py
+++ b/python/libknot/probe.py
@@ -1,6 +1,7 @@
 """Libknot probe interface wrapper."""
 
 import ctypes
+import datetime
 import enum
 import socket
 import libknot
@@ -98,16 +99,75 @@ class KnotProbeData(ctypes.Structure):
         return string
 
     def __str__(self) -> str:
-        """Returns data unit in pre-formatted simple text form."""
+        """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"
+        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()
-        string += "%s@%u > " % (self.addr_str(self.remote_addr), self.remote_port)
-        string += "%s@%u "   % (self.addr_str(self.local_addr), self.local_port)
-        string += "%s, " % ("UDP" if self.proto == 0 else "TCP")
-        string += "%s %s" % (self.qname_str(), libknot.Knot.rtype_str(self.query_type))
-        if self.edns_present == 1 and self.edns_flag_do == 1:
-            string += " DO"
-        string += ", %s" % libknot.Knot.rcode_str(self.reply_rcode)
+        if timestamp:
+            string += "%s " % COL(datetime.datetime.now().time(), YELW)
+        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)
+        string += (COL("UDP", GRN) if self.proto == 0 else COL("TCP", RED))
+        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
 
 
diff --git a/src/libknot/probe/data.h b/src/libknot/probe/data.h
index 1546f1298b..8bc7d1e439 100644
--- a/src/libknot/probe/data.h
+++ b/src/libknot/probe/data.h
@@ -125,7 +125,7 @@ int knot_probe_data_set(knot_probe_data_t *data, knot_probe_proto_t proto,
  *
  * \param sockfd  Socket descriptor of a TCP connection.
  *
- * \return Average TCP RTT.
+ * \return Average TCP RTT in microseconds.
  */
 uint32_t knot_probe_tcp_rtt(int sockfd);