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