nl80211 named, pythonic queries

nl80211.py

@@ -173,103 +173,38 @@ u32 = SingleStructParser('L')
 u64 = SingleStructParser('Q')
 
 
-rate_info = Attributes({
-  1: ('bitrate', u16),
-  2: ('mcs', u8),
-  4: ('short_gi', flag),
-  5: ('bitrate32', u32),
-  9: ('80p80_mhz_width', u32),
-  10: ('160_mhz_width', u32),
-})
-
-bss_param = Attributes({
-  2: ('short_preamble', flag),
-  3: ('short_slot_time', flag),
-  4: ('dtim_period', u8),
-  5: ('beacon_interval', u16),
-})
-
-sta_info = Attributes({
-  1: ('inactive_time', u32),
-  2: ('rx_bytes', u32),
-  3: ('tx_bytes', u32),
-  7: ('signal', u8),
-  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),
-  15: ('bss_param', bss_param),
-  16: ('connected_time', u32),
-  17: ('sta_flags', StructParser('LL', ('mask', 'values'))),
-  18: ('beacon_loss', u32),
-  23: ('rx_bytes_64', u64),
-  24: ('tx_bytes_64', u64),
-})
-
-nl80211_attr = Attributes({
-  3: ('ifindex', u32),
-  6: ('mac', string),
-  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
-STA_FLAG_SHORT_PREAMBLE = 1 << 1
-STA_FLAG_WME = 1 << 2
-STA_FLAG_MFP = 1 << 3
-STA_FLAG_AUTHENTICATED = 1 << 4
-STA_FLAG_TDLS_PEER = 1 << 5
-STA_FLAG_ASSOCIATED = 1 << 6
-
-
 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, msg):
-    self._sock.send(msg)
+  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:
@@ -277,53 +212,140 @@ class Netlink(object):
       iterator = Iterator(data)
       while not iterator.AtEnd():
         myhdr = self._nlmsghdr.Unpack(iterator)
-        if myhdr['type'] == NLMSG_DONE:
+        if myhdr['type'] == self._NLMSG_DONE:
           return
         yield (myhdr['type'], iterator.ExtractIterator(myhdr['length'] - self._nlmsghdr.size))
-        if not myhdr['flags'] & NLMSG_F_MULTI:
+        if not myhdr['flags'] & self._NLMSG_F_MULTI:
           return
 
 
 class GenericNetlink(object):
   _genlmsghdr = StructParser('BBH', ('cmd', 'version', 'reserved'))
 
-  _parser = {
-    0x10: ctrl_attr,
+  _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 Send(self, msg):
-    self._netlink.Send(msg)
+  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 = self._parser[msgtype]
+      parser = [v[0] for v in self._msgtypes.itervalues() if v[1] == msgtype][0]
       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))
+class NL80211(object):
+  _rate_info = Attributes({
+    1: ('bitrate', u16),
+    2: ('mcs', u8),
+    4: ('short_gi', flag),
+    5: ('bitrate32', u32),
+    9: ('80p80_mhz_width', u32),
+    10: ('160_mhz_width', u32),
+  })
 
-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()]
+  _bss_param = Attributes({
+    2: ('short_preamble', flag),
+    3: ('short_slot_time', flag),
+    4: ('dtim_period', u8),
+    5: ('beacon_interval', u16),
+  })
+
+  _sta_info = Attributes({
+    1: ('inactive_time', u32),
+    2: ('rx_bytes', u32),
+    3: ('tx_bytes', u32),
+    7: ('signal', u8),
+    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),
+    15: ('bss_param', _bss_param),
+    16: ('connected_time', u32),
+    17: ('sta_flags', StructParser('LL', ('mask', 'values'))),
+    18: ('beacon_loss', u32),
+    23: ('rx_bytes_64', u64),
+    24: ('tx_bytes_64', u64),
+  })
+
+  _nl80211_attr = Attributes({
+    3: ('ifindex', u32),
+    6: ('mac', string),
+    21: ('sta_info', _sta_info),
+    46: ('generation', u32),
+  })
+
+  CMD_GET_STATION = 17
+
+  STA_FLAG_AUTHORIZED = 1 << 0
+  STA_FLAG_SHORT_PREAMBLE = 1 << 1
+  STA_FLAG_WME = 1 << 2
+  STA_FLAG_MFP = 1 << 3
+  STA_FLAG_AUTHENTICATED = 1 << 4
+  STA_FLAG_TDLS_PEER = 1 << 5
+  STA_FLAG_ASSOCIATED = 1 << 6
+
+  def __init__(self):
+    self._gnl = GenericNetlink()
+    self._gnl.RegisterMsgType('nl80211', self._nl80211_attr)
+
+  def Send(self, flags, cmd, version, **attrs):
+    accumulator = Accumulator()
+    self._nl80211_attr.Pack(accumulator, **attrs)
+    self._gnl.Send('nl80211', flags, cmd, version, str(accumulator))
+
+  def Recv(self):
+    return self._gnl.Recv()
+
+
+nl = NL80211()
+nl.Send(Netlink.NLMSG_F_DUMP, nl.CMD_GET_STATION, 0, ifindex=10)
+print list(nl.Recv())