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