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         Attr.__init__(self, attr_type, "%dsB" % len(data), data.encode('utf-8'), 0)
 132
 133 class U32Attr(Attr):
 134     def __init__(self, attr_type, val):
 135         Attr.__init__(self, attr_type, "I", val)
 136
 137 class U8Attr(Attr):
 138     def __init__(self, attr_type, val):
 139         Attr.__init__(self, attr_type, "B", val)
 140
 141 class Nested(Attr):
 142     def __init__(self, attr_type, attrs):
 143         self.attrs = attrs
 144         self.type = attr_type
 145
 146     def _dump(self):
 147         contents = []
 148         for attr in self.attrs:
 149             contents.append(attr._dump())
 150         contents = ''.join(contents)
 151         length = len(contents)
 152         hdr = struct.pack("HH", length+4, self.type)
 153         return hdr + contents
 154
 155 NETLINK_ROUTE          = 0
 156 NETLINK_UNUSED         = 1
 157 NETLINK_USERSOCK       = 2
 158 NETLINK_FIREWALL       = 3
 159 NETLINK_INET_DIAG      = 4
 160 NETLINK_NFLOG          = 5
 161 NETLINK_XFRM           = 6
 162 NETLINK_SELINUX        = 7
 163 NETLINK_ISCSI          = 8
 164 NETLINK_AUDIT          = 9
 165 NETLINK_FIB_LOOKUP     = 10
 166 NETLINK_CONNECTOR      = 11
 167 NETLINK_NETFILTER      = 12
 168 NETLINK_IP6_FW         = 13
 169 NETLINK_DNRTMSG        = 14
 170 NETLINK_KOBJECT_UEVENT = 15
 171 NETLINK_GENERIC        = 16
 172
 173 class Message:
 174     def __init__(self, msg_type, flags=0, seq=-1, payload=None):
 175         self.type = msg_type
 176         self.flags = flags
 177         self.seq = seq
 178         self.pid = -1
 179         payload = payload or []
 180         if isinstance(payload, list):
 181             contents = []
 182             for attr in payload:
 183                 contents.append(attr._dump())
 184             self.payload = b''.join(contents)
 185         else:
 186             self.payload = payload
 187
 188     def send(self, conn):
 189         if self.seq == -1:
 190             self.seq = conn.seq()
 191
 192         self.pid = conn.pid
 193         length = len(self.payload)
 194
 195         hdr = struct.pack("IHHII", length + 4*4, self.type,
 196                           self.flags, self.seq, self.pid)
 197         conn.send(hdr + self.payload)
 198
 199     def __repr__(self):
 200         return '<netlink.Message type=%d, pid=%d, seq=%d, flags=0x%x "%s">' % (
 201             self.type, self.pid, self.seq, self.flags, repr(self.payload))
 202
 203 class Connection:
 204     def __init__(self, nltype, groups=0, unexpected_msg_handler=None):
 205         self.descriptor = socket.socket(socket.AF_NETLINK,
 206                                         socket.SOCK_RAW, nltype)
 207         self.descriptor.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
 208         self.descriptor.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
 209         _nl_bind(self.descriptor, (0, groups))
 210         self.pid, self.groups = _nl_getsockname(self.descriptor)
 211         self._seq = 0
 212         self.unexpected = unexpected_msg_handler
 213     def send(self, msg):
 214         _nl_send(self.descriptor, msg)
 215     def recv(self):
 216         contents, (nlpid, nlgrps) = _nl_recv(self.descriptor)
 217         # XXX: python doesn't give us message flags, check
 218         #      len(contents) vs. msglen for TRUNC
 219         msglen, msg_type, flags, seq, pid = struct.unpack("IHHII",
 220                                                           contents[:16])
 221         msg = Message(msg_type, flags, seq, contents[16:])
 222         msg.pid = pid
 223         if msg.type == NLMSG_ERROR:
 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