iotop/netlink.py

   1 '''
   2 Netlink message generation/parsing
   3
   4 Copyright 2007        Johannes Berg <johannes@sipsolutions.net>
   5
   6 GPLv2+; See copying for details.
   7 '''
   8
   9 import struct, socket
  10
  11 try:
  12     # try to use python 2.5's netlink support
  13     _dummysock = socket.socket(socket.AF_NETLINK, socket.SOCK_RAW, 0)
  14     _dummysock.bind((0, 0))
  15     del _dummysock
  16     def _nl_bind(descriptor, addr):
  17         descriptor.bind(addr)
  18     def _nl_getsockname(descriptor):
  19         return descriptor.getsockname()
  20     def _nl_send(descriptor, msg):
  21         descriptor.send(msg)
  22     def _nl_recv(descriptor, bufs=16384):
  23         return descriptor.recvfrom(bufs)
  24 except socket.error:
  25     # or fall back to the _netlink C module
  26     try:
  27         import _netlink
  28         def _nl_bind(descriptor, addr):
  29             _netlink.bind(descriptor.fileno(), addr[1])
  30         def _nl_getsockname(descriptor):
  31             return _netlink.getsockname(descriptor.fileno())
  32         def _nl_send(descriptor, msg):
  33             _netlink.send(descriptor.fileno(), msg)
  34         def _nl_recv(descriptor, bufs=16384):
  35             return _netlink.recvfrom(descriptor.fileno(), bufs)
  36     except ImportError:
  37         # or fall back to the ctypes module
  38         import ctypes
  39         import os
  40
  41         libc = ctypes.CDLL(None)
  42
  43         class SOCKADDR_NL(ctypes.Structure):
  44             _fields_ = [("nl_family", ctypes.c_ushort),
  45                         ("nl_pad",    ctypes.c_ushort),
  46                         ("nl_pid",    ctypes.c_int),
  47                         ("nl_groups", ctypes.c_int)]
  48
  49         def _nl_bind(descriptor, addr):
  50             addr = SOCKADDR_NL(socket.AF_NETLINK, 0, os.getpid(), 0)
  51             return libc.bind(descriptor.fileno(),
  52                              ctypes.pointer(addr),
  53                              ctypes.sizeof(addr))
  54
  55         def _nl_getsockname(descriptor):
  56             addr = SOCKADDR_NL(0, 0, 0, 0)
  57             len = ctypes.c_int(ctypes.sizeof(addr));
  58             libc.getsockname(descriptor.fileno(),
  59                              ctypes.pointer(addr),
  60                              ctypes.pointer(len))
  61             return addr.nl_pid, addr.nl_groups;
  62
  63         def _nl_send(descriptor, msg):
  64             return libc.send(descriptor.fileno(), msg, len(msg), 0);
  65
  66         def _nl_recv(descriptor, bufs=16384):
  67             addr = SOCKADDR_NL(0, 0, 0, 0)
  68             len = ctypes.c_int(ctypes.sizeof(addr))
  69             buf = ctypes.create_string_buffer(bufs)
  70
  71             r = libc.recvfrom(descriptor.fileno(),
  72                               buf, bufs, 0,
  73                               ctypes.pointer(addr), ctypes.pointer(len))
  74
  75             ret = ctypes.string_at(ctypes.pointer(buf), r)
  76             return ret, (addr.nl_pid, addr.nl_groups)
  77
  78
  79 # flags
  80 NLM_F_REQUEST  = 1
  81 NLM_F_MULTI    = 2
  82 NLM_F_ACK      = 4
  83 NLM_F_ECHO     = 8
  84
  85 # types
  86 NLMSG_NOOP     = 1
  87 NLMSG_ERROR    = 2
  88 NLMSG_DONE     = 3
  89 NLMSG_OVERRUN  = 4
  90 NLMSG_MIN_TYPE = 0x10
  91
  92 class Attr:
  93     def __init__(self, attr_type, data, *values):
  94         self.type = attr_type
  95         if len(values):
  96             self.data = struct.pack(data, *values)
  97         else:
  98             self.data = data
  99
 100     def _dump(self):
 101         hdr = struct.pack("HH", len(self.data)+4, self.type)
 102         length = len(self.data)
 103         pad = ((length + 4 - 1) & ~3 ) - length
 104         return hdr + self.data + '\0' * pad
 105
 106     def __repr__(self):
 107         return '<Attr type %d, data "%s">' % (self.type, repr(self.data))
 108
 109     def u16(self):
 110         return struct.unpack('H', self.data)[0]
 111     def s16(self):
 112         return struct.unpack('h', self.data)[0]
 113     def u32(self):
 114         return struct.unpack('I', self.data)[0]
 115     def s32(self):
 116         return struct.unpack('i', self.data)[0]
 117     def str(self):
 118         return self.data
 119     def nulstr(self):
 120         return self.data.split('\0')[0]
 121     def nested(self):
 122         return parse_attributes(self.data)
 123
 124 class StrAttr(Attr):
 125     def __init__(self, attr_type, data):
 126         Attr.__init__(self, attr_type, "%ds" % len(data), data)
 127
 128 class NulStrAttr(Attr):
 129     def __init__(self, attr_type, data):
 130         Attr.__init__(self, attr_type, "%dsB" % len(data), data, 0)
 131
 132 class U32Attr(Attr):
 133     def __init__(self, attr_type, val):
 134         Attr.__init__(self, attr_type, "I", val)
 135
 136 class U8Attr(Attr):
 137     def __init__(self, attr_type, val):
 138         Attr.__init__(self, attr_type, "B", val)
 139
 140 class Nested(Attr):
 141     def __init__(self, attr_type, attrs):
 142         self.attrs = attrs
 143         self.type = attr_type
 144
 145     def _dump(self):
 146         contents = []
 147         for attr in self.attrs:
 148             contents.append(attr._dump())
 149         contents = ''.join(contents)
 150         length = len(contents)
 151         hdr = struct.pack("HH", length+4, self.type)
 152         return hdr + contents
 153
 154 NETLINK_ROUTE          = 0
 155 NETLINK_UNUSED         = 1
 156 NETLINK_USERSOCK       = 2
 157 NETLINK_FIREWALL       = 3
 158 NETLINK_INET_DIAG      = 4
 159 NETLINK_NFLOG          = 5
 160 NETLINK_XFRM           = 6
 161 NETLINK_SELINUX        = 7
 162 NETLINK_ISCSI          = 8
 163 NETLINK_AUDIT          = 9
 164 NETLINK_FIB_LOOKUP     = 10
 165 NETLINK_CONNECTOR      = 11
 166 NETLINK_NETFILTER      = 12
 167 NETLINK_IP6_FW         = 13
 168 NETLINK_DNRTMSG        = 14
 169 NETLINK_KOBJECT_UEVENT = 15
 170 NETLINK_GENERIC        = 16
 171
 172 class Message:
 173     def __init__(self, msg_type, flags=0, seq=-1, payload=None):
 174         self.type = msg_type
 175         self.flags = flags
 176         self.seq = seq
 177         self.pid = -1
 178         payload = payload or []
 179         if isinstance(payload, list):
 180             contents = []
 181             for attr in payload:
 182                 contents.append(attr._dump())
 183             self.payload = ''.join(contents)
 184         else:
 185             self.payload = payload
 186
 187     def send(self, conn):
 188         if self.seq == -1:
 189             self.seq = conn.seq()
 190
 191         self.pid = conn.pid
 192         length = len(self.payload)
 193
 194         hdr = struct.pack("IHHII", length + 4*4, self.type,
 195                           self.flags, self.seq, self.pid)
 196         conn.send(hdr + self.payload)
 197
 198     def __repr__(self):
 199         return '<netlink.Message type=%d, pid=%d, seq=%d, flags=0x%x "%s">' % (
 200             self.type, self.pid, self.seq, self.flags, repr(self.payload))
 201
 202 class Connection:
 203     def __init__(self, nltype, groups=0, unexpected_msg_handler=None):
 204         self.descriptor = socket.socket(socket.AF_NETLINK,
 205                                         socket.SOCK_RAW, nltype)
 206         self.descriptor.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
 207         self.descriptor.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
 208         _nl_bind(self.descriptor, (0, groups))
 209         self.pid, self.groups = _nl_getsockname(self.descriptor)
 210         self._seq = 0
 211         self.unexpected = unexpected_msg_handler
 212     def send(self, msg):
 213         _nl_send(self.descriptor, msg)
 214     def recv(self):
 215         contents, (nlpid, nlgrps) = _nl_recv(self.descriptor)
 216         # XXX: python doesn't give us message flags, check
 217         #      len(contents) vs. msglen for TRUNC
 218         msglen, msg_type, flags, seq, pid = struct.unpack("IHHII",
 219                                                           contents[:16])
 220         msg = Message(msg_type, flags, seq, contents[16:])
 221         msg.pid = pid
 222         if msg.type == NLMSG_ERROR:
 223             import os
 224             errno = -struct.unpack("i", msg.payload[:4])[0]
 225             if errno != 0:
 226                 err = OSError("Netlink error: %s (%d)" % (
 227                                                      os.strerror(errno), errno))
 228                 err.errno = errno
 229                 raise err
 230         return msg
 231     def seq(self):
 232         self._seq += 1
 233         return self._seq
 234
 235 def parse_attributes(data):
 236     attrs = {}
 237     while len(data):
 238         attr_len, attr_type = struct.unpack("HH", data[:4])
 239         attrs[attr_type] = Attr(attr_type, data[4:attr_len])
 240         attr_len = ((attr_len + 4 - 1) & ~3 )
 241         data = data[attr_len:]
 242     return attrs