1#!/usr/bin/python
2#
3# Copyright 2014 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Partial Python implementation of iproute functionality."""
18
19# pylint: disable=g-bad-todo
20
21import errno
22import os
23import socket
24import struct
25import sys
26
27import cstruct
28
29
30# Request constants.
31NLM_F_REQUEST = 1
32NLM_F_ACK = 4
33NLM_F_REPLACE = 0x100
34NLM_F_EXCL = 0x200
35NLM_F_CREATE = 0x400
36NLM_F_DUMP = 0x300
37
38# Message types.
39NLMSG_ERROR = 2
40NLMSG_DONE = 3
41
42# Data structure formats.
43# These aren't constants, they're classes. So, pylint: disable=invalid-name
44NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
45NLMsgErr = cstruct.Struct("NLMsgErr", "=i", "error")
46NLAttr = cstruct.Struct("NLAttr", "=HH", "nla_len nla_type")
47
48# Alignment / padding.
49NLA_ALIGNTO = 4
50
51
52def PaddedLength(length):
53  # TODO: This padding is probably overly simplistic.
54  return NLA_ALIGNTO * ((length / NLA_ALIGNTO) + (length % NLA_ALIGNTO != 0))
55
56
57class NetlinkSocket(object):
58  """A basic netlink socket object."""
59
60  BUFSIZE = 65536
61  DEBUG = False
62  # List of netlink messages to print, e.g., [], ["NEIGH", "ROUTE"], or ["ALL"]
63  NL_DEBUG = []
64
65  def _Debug(self, s):
66    if self.DEBUG:
67      print s
68
69  def _NlAttr(self, nla_type, data):
70    datalen = len(data)
71    # Pad the data if it's not a multiple of NLA_ALIGNTO bytes long.
72    padding = "\x00" * (PaddedLength(datalen) - datalen)
73    nla_len = datalen + len(NLAttr)
74    return NLAttr((nla_len, nla_type)).Pack() + data + padding
75
76  def _NlAttrU32(self, nla_type, value):
77    return self._NlAttr(nla_type, struct.pack("=I", value))
78
79  def _GetConstantName(self, module, value, prefix):
80    thismodule = sys.modules[module]
81    for name in dir(thismodule):
82      if name.startswith("INET_DIAG_BC"):
83        break
84      if (name.startswith(prefix) and
85          not name.startswith(prefix + "F_") and
86          name.isupper() and getattr(thismodule, name) == value):
87          return name
88    return value
89
90  def _Decode(self, command, msg, nla_type, nla_data):
91    """No-op, nonspecific version of decode."""
92    return nla_type, nla_data
93
94  def _ParseAttributes(self, command, family, msg, data):
95    """Parses and decodes netlink attributes.
96
97    Takes a block of NLAttr data structures, decodes them using Decode, and
98    returns the result in a dict keyed by attribute number.
99
100    Args:
101      command: An integer, the rtnetlink command being carried out.
102      family: The address family.
103      msg: A Struct, the type of the data after the netlink header.
104      data: A byte string containing a sequence of NLAttr data structures.
105
106    Returns:
107      A dictionary mapping attribute types (integers) to decoded values.
108
109    Raises:
110      ValueError: There was a duplicate attribute type.
111    """
112    attributes = {}
113    while data:
114      # Read the nlattr header.
115      nla, data = cstruct.Read(data, NLAttr)
116
117      # Read the data.
118      datalen = nla.nla_len - len(nla)
119      padded_len = PaddedLength(nla.nla_len) - len(nla)
120      nla_data, data = data[:datalen], data[padded_len:]
121
122      # If it's an attribute we know about, try to decode it.
123      nla_name, nla_data = self._Decode(command, msg, nla.nla_type, nla_data)
124
125      # We only support unique attributes for now, except for INET_DIAG_NONE,
126      # which can appear more than once but doesn't seem to contain any data.
127      if nla_name in attributes and nla_name != "INET_DIAG_NONE":
128        raise ValueError("Duplicate attribute %s" % nla_name)
129
130      attributes[nla_name] = nla_data
131      self._Debug("      %s" % str((nla_name, nla_data)))
132
133    return attributes
134
135  def __init__(self):
136    # Global sequence number.
137    self.seq = 0
138    self.sock = socket.socket(socket.AF_NETLINK, socket.SOCK_RAW, self.FAMILY)
139    self.sock.connect((0, 0))  # The kernel.
140    self.pid = self.sock.getsockname()[1]
141
142  def _Send(self, msg):
143    # self._Debug(msg.encode("hex"))
144    self.seq += 1
145    self.sock.send(msg)
146
147  def _Recv(self):
148    data = self.sock.recv(self.BUFSIZE)
149    # self._Debug(data.encode("hex"))
150    return data
151
152  def _ExpectDone(self):
153    response = self._Recv()
154    hdr = NLMsgHdr(response)
155    if hdr.type != NLMSG_DONE:
156      raise ValueError("Expected DONE, got type %d" % hdr.type)
157
158  def _ParseAck(self, response):
159    # Find the error code.
160    hdr, data = cstruct.Read(response, NLMsgHdr)
161    if hdr.type == NLMSG_ERROR:
162      error = NLMsgErr(data).error
163      if error:
164        raise IOError(error, os.strerror(-error))
165    else:
166      raise ValueError("Expected ACK, got type %d" % hdr.type)
167
168  def _ExpectAck(self):
169    response = self._Recv()
170    self._ParseAck(response)
171
172  def _SendNlRequest(self, command, data, flags):
173    """Sends a netlink request and expects an ack."""
174    length = len(NLMsgHdr) + len(data)
175    nlmsg = NLMsgHdr((length, command, flags, self.seq, self.pid)).Pack()
176
177    self.MaybeDebugCommand(command, nlmsg + data)
178
179    # Send the message.
180    self._Send(nlmsg + data)
181
182    if flags & NLM_F_ACK:
183      self._ExpectAck()
184
185  def _ParseNLMsg(self, data, msgtype):
186    """Parses a Netlink message into a header and a dictionary of attributes."""
187    nlmsghdr, data = cstruct.Read(data, NLMsgHdr)
188    self._Debug("  %s" % nlmsghdr)
189
190    if nlmsghdr.type == NLMSG_ERROR or nlmsghdr.type == NLMSG_DONE:
191      print "done"
192      return (None, None), data
193
194    nlmsg, data = cstruct.Read(data, msgtype)
195    self._Debug("    %s" % nlmsg)
196
197    # Parse the attributes in the nlmsg.
198    attrlen = nlmsghdr.length - len(nlmsghdr) - len(nlmsg)
199    attributes = self._ParseAttributes(nlmsghdr.type, nlmsg.family,
200                                       nlmsg, data[:attrlen])
201    data = data[attrlen:]
202    return (nlmsg, attributes), data
203
204  def _GetMsg(self, msgtype):
205    data = self._Recv()
206    if NLMsgHdr(data).type == NLMSG_ERROR:
207      self._ParseAck(data)
208    return self._ParseNLMsg(data, msgtype)[0]
209
210  def _GetMsgList(self, msgtype, data, expect_done):
211    out = []
212    while data:
213      msg, data = self._ParseNLMsg(data, msgtype)
214      if msg is None:
215        break
216      out.append(msg)
217    if expect_done:
218      self._ExpectDone()
219    return out
220
221  def _Dump(self, command, msg, msgtype, attrs):
222    """Sends a dump request and returns a list of decoded messages.
223
224    Args:
225      command: An integer, the command to run (e.g., RTM_NEWADDR).
226      msg: A string, the raw bytes of the request (e.g., a packed RTMsg).
227      msgtype: A cstruct.Struct, the data type to parse the dump results as.
228      attrs: A string, the raw bytes of any request attributes to include.
229
230    Returns:
231      A list of (msg, attrs) tuples where msg is of type msgtype and attrs is
232      a dict of attributes.
233    """
234    # Create a netlink dump request containing the msg.
235    flags = NLM_F_DUMP | NLM_F_REQUEST
236    length = len(NLMsgHdr) + len(msg) + len(attrs)
237    nlmsghdr = NLMsgHdr((length, command, flags, self.seq, self.pid))
238
239    # Send the request.
240    self._Send(nlmsghdr.Pack() + msg.Pack() + attrs)
241
242    # Keep reading netlink messages until we get a NLMSG_DONE.
243    out = []
244    while True:
245      data = self._Recv()
246      response_type = NLMsgHdr(data).type
247      if response_type == NLMSG_DONE:
248        break
249      elif response_type == NLMSG_ERROR:
250        # Likely means that the kernel didn't like our dump request.
251        # Parse the error and throw an exception.
252        self._ParseAck(data)
253      out.extend(self._GetMsgList(msgtype, data, False))
254
255    return out
256