nl80211 named, pythonic queries

nl80211.py

@@ -173,7 +173,114 @@ u32 = SingleStructParser('L')
 u64 = SingleStructParser('Q')
 
 
-rate_info = Attributes({
+class Netlink(object):
+  _NLMSG_F_REQUEST = 0x01
+  _NLMSG_F_MULTI = 0x02
+  _NLMSG_F_ACK = 0x04
+  _NLMSG_F_ECHO = 0x08
+  _NLMSG_F_DUMP_INTR = 0x10
+
+  NLMSG_F_ROOT = 0x100
+  NLMSG_F_MATCH = 0x200
+  NLMSG_F_ATOMIC = 0x400
+  NLMSG_F_DUMP = NLMSG_F_ROOT | NLMSG_F_MATCH
+
+  _NLMSG_DONE = 3
+
+  _nlmsghdr = StructParser('LHHLL', ('length', 'type', 'flags', 'seq', 'pid'))
+
+  def __init__(self):
+    self._sock = socket.socket(socket.AF_NETLINK, socket.SOCK_DGRAM, 16)
+    self._sock.bind((0, 0))
+
+  def Send(self, msgtype, flags, msg):
+    flags |= self._NLMSG_F_REQUEST
+    accumulator = Accumulator()
+    self._nlmsghdr.Pack(
+      accumulator,
+      length=len(msg) + self._nlmsghdr.size,
+      type=msgtype,
+      flags=flags,
+      seq=random.randint(0, 2 ** 32 - 1),
+      pid=os.getpid())
+    accumulator.Append(msg)
+    self._sock.send(str(accumulator))
+
+  def Recv(self):
+    while True:
+      data = self._sock.recv(4096)
+      iterator = Iterator(data)
+      while not iterator.AtEnd():
+        myhdr = self._nlmsghdr.Unpack(iterator)
+        if myhdr['type'] == self._NLMSG_DONE:
+          return
+        yield (myhdr['type'], iterator.ExtractIterator(myhdr['length'] - self._nlmsghdr.size))
+        if not myhdr['flags'] & self._NLMSG_F_MULTI:
+          return
+
+
+class GenericNetlink(object):
+  _genlmsghdr = StructParser('BBH', ('cmd', 'version', 'reserved'))
+
+  _op_attr = Attributes({
+    1: ('id', u32),
+    2: ('flags', u32),
+  })
+
+  _mcast_grp_attr = Attributes({
+    1: ('name', string),
+    2: ('id', u32),
+  })
+
+  _ctrl_attr = Attributes({
+    1: ('family_id', u16),
+    2: ('family_name', string),
+    3: ('version', u32),
+    4: ('hdrsize', u32),
+    5: ('maxattr', u32),
+    6: ('ops', Array(_op_attr)),
+    7: ('mcast_groups', Array(_mcast_grp_attr)),
+  })
+
+  _msgtypes = {
+    'nlctrl': [_ctrl_attr, 0x10],
+  }
+
+  CTRL_CMD_NEWFAMILY = 0x01
+  CTRL_CMD_GETFAMILY = 0x03
+
+  def __init__(self):
+    self._netlink = Netlink()
+    self.Send('nlctrl', self._netlink.NLMSG_F_DUMP, self.CTRL_CMD_GETFAMILY, 1, '')
+    for msg in self.Recv():
+      msgtype, attrs = msg
+      assert msgtype == self.CTRL_CMD_NEWFAMILY, msgtype
+      family_name = attrs['family_name'].rstrip('\0')
+      self._msgtypes.setdefault(family_name, [None, None])[1] = attrs['family_id']
+
+  def RegisterMsgType(self, family_name, parser):
+    self._msgtypes[family_name][0] = parser
+
+  def Send(self, msgtype, flags, cmd, version, msg):
+    accumulator = Accumulator()
+    self._genlmsghdr.Pack(
+      accumulator,
+      cmd=cmd,
+      version=version,
+      reserved=0)
+    accumulator.Append(msg)
+    msgtype_id = self._msgtypes[msgtype][1]
+    self._netlink.Send(msgtype_id, flags, str(accumulator))
+
+  def Recv(self):
+    for msgtype, iterator in self._netlink.Recv():
+      genlhdr = self._genlmsghdr.Unpack(iterator)
+      parser = [v[0] for v in self._msgtypes.itervalues() if v[1] == msgtype][0]
+      yield (genlhdr['cmd'], parser.Unpack(iterator))
+
+
+class NL80211(object):
+  _rate_info = Attributes({
     1: ('bitrate', u16),
     2: ('mcs', u8),
     4: ('short_gi', flag),
@@ -182,26 +289,26 @@ rate_info = Attributes({
     10: ('160_mhz_width', u32),
   })
 
-bss_param = Attributes({
+  _bss_param = Attributes({
     2: ('short_preamble', flag),
     3: ('short_slot_time', flag),
     4: ('dtim_period', u8),
     5: ('beacon_interval', u16),
   })
 
-sta_info = Attributes({
+  _sta_info = Attributes({
     1: ('inactive_time', u32),
     2: ('rx_bytes', u32),
     3: ('tx_bytes', u32),
     7: ('signal', u8),
-  8: ('tx_bitrate', rate_info),
+    8: ('tx_bitrate', _rate_info),
     9: ('rx_packets', u32),
     10: ('tx_packets', u32),
     11: ('tx_retries', u32),
     12: ('tx_failed', u32),
     13: ('signal_avg', u8),
-  14: ('rx_bitrate', rate_info),
+    14: ('rx_bitrate', _rate_info),
-  15: ('bss_param', bss_param),
+    15: ('bss_param', _bss_param),
     16: ('connected_time', u32),
     17: ('sta_flags', StructParser('LL', ('mask', 'values'))),
     18: ('beacon_loss', u32),
@@ -209,47 +316,13 @@ sta_info = Attributes({
     24: ('tx_bytes_64', u64),
   })
 
-nl80211_attr = Attributes({
+  _nl80211_attr = Attributes({
     3: ('ifindex', u32),
     6: ('mac', string),
-  21: ('sta_info', sta_info),
+    21: ('sta_info', _sta_info),
     46: ('generation', u32),
   })
 
-
-op_attr = Attributes({
-  1: ('id', u32),
-  2: ('flags', u32),
-})
-
-mcast_grp_attr = Attributes({
-  1: ('name', string),
-  2: ('id', u32),
-})
-
-ctrl_attr = Attributes({
-  1: ('family_id', u16),
-  2: ('family_name', string),
-  3: ('version', u32),
-  4: ('hdrsize', u32),
-  5: ('maxattr', u32),
-  6: ('ops', Array(op_attr)),
-  7: ('mcast_groups', Array(mcast_grp_attr)),
-})
-
-NLMSG_F_REQUEST = 1 << 0
-NLMSG_F_MULTI = 1 << 1
-NLMSG_F_ACK = 1 << 2
-NLMSG_F_ECHO = 1 << 3
-NLMSG_F_DUMP_INTR = 1 << 4
-
-NLMSG_F_ROOT = 1 << 8
-NLMSG_F_MATCH = 1 << 9
-NLMSG_F_ATOMIC = 1 << 10
-NLMSG_F_DUMP = NLMSG_F_ROOT | NLMSG_F_MATCH
-
-NLMSG_DONE = 3
-
   CMD_GET_STATION = 17
 
   STA_FLAG_AUTHORIZED = 1 << 0
@@ -260,70 +333,19 @@ STA_FLAG_AUTHENTICATED = 1 << 4
   STA_FLAG_TDLS_PEER = 1 << 5
   STA_FLAG_ASSOCIATED = 1 << 6
 
-
-class Netlink(object):
-  _nlmsghdr = StructParser('LHHLL', ('length', 'type', 'flags', 'seq', 'pid'))
-
   def __init__(self):
-    self._sock = socket.socket(socket.AF_NETLINK, socket.SOCK_DGRAM, 16)
+    self._gnl = GenericNetlink()
-    self._sock.bind((0, 0))
+    self._gnl.RegisterMsgType('nl80211', self._nl80211_attr)
 
-  def Send(self, msg):
+  def Send(self, flags, cmd, version, **attrs):
-    self._sock.send(msg)
+    accumulator = Accumulator()
+    self._nl80211_attr.Pack(accumulator, **attrs)
+    self._gnl.Send('nl80211', flags, cmd, version, str(accumulator))
 
   def Recv(self):
-    while True:
+    return self._gnl.Recv()
-      data = self._sock.recv(4096)
-      iterator = Iterator(data)
-      while not iterator.AtEnd():
-        myhdr = self._nlmsghdr.Unpack(iterator)
-        if myhdr['type'] == NLMSG_DONE:
-          return
-        yield (myhdr['type'], iterator.ExtractIterator(myhdr['length'] - self._nlmsghdr.size))
-        if not myhdr['flags'] & NLMSG_F_MULTI:
-          return
 
 
-class GenericNetlink(object):
+nl = NL80211()
-  _genlmsghdr = StructParser('BBH', ('cmd', 'version', 'reserved'))
+nl.Send(Netlink.NLMSG_F_DUMP, nl.CMD_GET_STATION, 0, ifindex=10)
-
+print list(nl.Recv())
-  _parser = {
-    0x10: ctrl_attr,
-  }
-
-  def __init__(self):
-    self._netlink = Netlink()
-
-  def Send(self, msg):
-    self._netlink.Send(msg)
-
-  def Recv(self):
-    for msgtype, iterator in self._netlink.Recv():
-      genlhdr = self._genlmsghdr.Unpack(iterator)
-      parser = self._parser[msgtype]
-      yield (genlhdr['cmd'], parser.Unpack(iterator))
-
-
-#int_genquery = Accumulator()
-#genlmsghdr.Pack(
-#  int_genquery,
-#  cmd=CMD_GET_STATION,
-#  version=0,
-#  reserved=0)
-#nl80211_attr.Pack(
-#  int_genquery,
-#  ifindex=6)
-#genquery = Accumulator()
-#nlmsghdr.Pack(
-#  genquery, 
-#  length=nlmsghdr.size + len(int_genquery),
-#  type=20, # XXX
-#  flags=F_REQUEST | F_ACK | F_DUMP,
-#  seq=random.randint(0, 2 ** 32 - 1),
-#  pid=os.getpid())
-#genquery.Append(str(int_genquery))
-
-query = '\24\0\0\0\20\0\5\3a\6\256T\v\17\0\0\3\1\0\0'
-conn = GenericNetlink()
-conn.Send(query)
-print ['%s\n' % s for cmd, s in conn.Recv()]