Merge pull request svinota#3 from superbobry/netlink_send

Minor changes to `netlink` API

cxnet/netlink/generic.py

@@ -26,9 +26,6 @@
 from core import nlmsghdr, nl_socket, nlattr, attr_msg, NLM_F_REQUEST, NLMSG_MAX_LEN, NLMSG_MIN_TYPE, NLMSG_ALIGN
 import sys
 
-GENL_NAMSIZ   = 16    # length of family name
-GENL_MIN_ID   = NLMSG_MIN_TYPE
-GENL_MAX_ID   = 1023
 
 class genlmsghdr(Structure):
     _fields_ = [
@@ -37,13 +34,19 @@ class genlmsghdr(Structure):
         ("reserved",    c_uint16),
     ]
 
-class genlmsg(Structure,attr_msg):
+
+class genlmsg(Structure, attr_msg):
     _fields_ = [
         ("hdr",         nlmsghdr),
         ("genlmsghdr",  genlmsghdr),
         ("data",        c_byte * (NLMSG_MAX_LEN - sizeof(nlmsghdr) - sizeof(genlmsghdr))),
     ]
 
+
+GENL_NAMSIZ   = 16    # length of family name
+GENL_MIN_ID   = NLMSG_MIN_TYPE
+GENL_MAX_ID   = 1023
+
 GENL_HDRLEN         = NLMSG_ALIGN(sizeof(genlmsghdr))
 GENL_ADMIN_PERM     = 0x01
 GENL_CMD_CAP_DO     = 0x02
@@ -89,30 +92,32 @@ class genlmsg(Structure,attr_msg):
 CTRL_ATTR_MCAST_GRP_NAME    = 0x1
 CTRL_ATTR_MCAST_GRP_ID      = 0x2
 
+
 class genl_socket(nl_socket):
 
     msg = genlmsg
 
-    def get_protocol_id(self,name):
-        if sys.version_info >= (3,0):
-            buf = create_string_buffer(bytes(name,"ascii"))
+    def get_protocol_id(self, name):
+        if sys.version_info >= (3, 0):
+            buf = create_string_buffer(name.encode("utf-8"))
         else:
             buf = create_string_buffer(name)
-        (l,msg) = self.send_cmd(GENL_ID_CTRL,CTRL_CMD_GETFAMILY,CTRL_ATTR_FAMILY_NAME,buf)
+
+        self.send_cmd(GENL_ID_CTRL, CTRL_CMD_GETFAMILY,
+                      CTRL_ATTR_FAMILY_NAME, buf)
+        l, msg = self.recv()
         name = nlattr.from_address(addressof(msg.data))
-        prid = nlattr.from_address(addressof(msg.data) + NLMSG_ALIGN(name.nla_len))
+        prid = nlattr.from_address(addressof(msg.data) +
+                                   NLMSG_ALIGN(name.nla_len))
         assert prid.nla_type == CTRL_ATTR_FAMILY_ID
         return c_uint16.from_address(addressof(prid) + sizeof(prid)).value
 
-
-    def send_cmd(self,prid,cmd,nla_type,nla_data,seq=0):
+    def send_cmd(self, prid, cmd, nla_type, nla_data, seq=0):
         msg = genlmsg()
         msg.hdr.type = prid
         msg.hdr.flags = NLM_F_REQUEST
         msg.hdr.sequence_number = seq
         msg.genlmsghdr.cmd = cmd
         msg.genlmsghdr.version = 0x1
-        msg.set_attr(nla_type,nla_data)
-        self.send(msg,msg.size())
-        return  self.recv()
-
+        msg.set_attr(nla_type, nla_data)
+        return self.send(msg, msg.size())

cxnet/netlink/taskstats.py

@@ -20,17 +20,19 @@
 #     along with Connexion; if not, write to the Free Software
 #     Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 
+from ctypes import Structure
+from ctypes import create_string_buffer, sizeof, addressof
+from ctypes import c_uint8, c_uint16, c_uint32, c_uint64, c_char, c_ubyte
+
 from cxnet.utils import hprint
 from cxnet.netlink.core import nlattr, NLMSG_ALIGN
 from cxnet.netlink.generic import genl_socket
 
-from ctypes import Structure
-from ctypes import create_string_buffer, sizeof, addressof
-from ctypes import c_uint8, c_uint16, c_uint32, c_uint64, c_char, c_ubyte
 
 TASKSTATS_VERSION = 6
 TS_COMM_LEN = 32
 
+
 class taskstatsmsg(Structure):
     _pack_ = 8
     _fields_ = [
@@ -193,17 +195,14 @@ class taskstatsmsg(Structure):
         "cpu_scaled_run_real_total":"Scaled cpu_run_real_total",
     }
 
-    def pprint(self,attr=None):
-        print(self.sprint(attr))
-
-    def sprint(self,attr=None):
-        if attr is not None:
-            return "%-26s%-32s%s" % (attr,getattr(self,attr),self.descriptions[attr])
-        else:
-            s = ""
-            for x in sorted(self.descriptions.keys()):
-                s += "%s\n" % (self.sprint(x))
-            return s
+    def __str__(self):
+        def inner(attr):
+            return "{0:<26}{1:<32}{2}".format(
+                attr, getattr(self, attr), self.descriptions[attr])
+
+        return "\n".join(map(inner, sorted(self.descriptions)))
+
+
 #
 # Commands sent from userspace
 # Not versioned.
@@ -226,8 +225,6 @@ def sprint(self,attr=None):
 TASKSTATS_CMD_ATTR_DEREGISTER_CPUMASK = 4
 
 
-
-
 if __name__ == "__main__":
     ###
     #
@@ -247,25 +244,30 @@ def sprint(self,attr=None):
     if len(sys.argv) > 1:
         try:
             pid = int(sys.argv[1])
-        except:
+        except ValueError:
             mask = sys.argv[1][1:]
     else:
         import os
         pid = os.getpid()
 
-
     if pid:
-        (l,msg) = s.send_cmd(prid,TASKSTATS_CMD_GET,TASKSTATS_CMD_ATTR_PID,c_uint32(pid))
+        s.send_cmd(prid, TASKSTATS_CMD_GET, TASKSTATS_CMD_ATTR_PID, c_uint32(pid))
+        l, msg = s.recv()
+
     if mask:
-        print "PRID: %x" % (prid)
-        (l,msg) = s.send_cmd(prid,TASKSTATS_CMD_GET,TASKSTATS_CMD_ATTR_REGISTER_CPUMASK,create_string_buffer(mask,8))
+        s.send_cmd(prid,
+                   TASKSTATS_CMD_GET,
+                   TASKSTATS_CMD_ATTR_REGISTER_CPUMASK,
+                   create_string_buffer(mask, 8))
+        l, msg = s.recv()
+
     a = nlattr.from_address(addressof(msg.data))
     assert a.nla_type == TASKSTATS_TYPE_AGGR_PID
     pid = nlattr.from_address(addressof(msg.data) + sizeof(a))
     assert pid.nla_type == TASKSTATS_TYPE_PID
-    stats = taskstatsmsg.from_address(addressof(msg.data) + sizeof(a) + NLMSG_ALIGN(pid.nla_len) + sizeof(nlattr))
-
-    stats.pprint()
+    stats = taskstatsmsg.from_address(addressof(msg.data) + sizeof(a) +
+                                      NLMSG_ALIGN(pid.nla_len) + sizeof(nlattr))
 
-    print("\nraw packet dump:")
-    hprint(msg,l)
+    print(stats)    # Taskstats structure.
+    print("\n")
+    hprint(msg, l)  # Raw packet dump.