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
17import fcntl
18import os
19import random
20import re
21from socket import *  # pylint: disable=wildcard-import
22import struct
23import unittest
24
25from scapy import all as scapy
26
27SOL_IPV6 = 41
28IP_RECVERR = 11
29IPV6_RECVERR = 25
30IP_TRANSPARENT = 19
31IPV6_TRANSPARENT = 75
32IPV6_TCLASS = 67
33IPV6_FLOWLABEL_MGR = 32
34IPV6_FLOWINFO_SEND = 33
35
36SO_BINDTODEVICE = 25
37SO_MARK = 36
38SO_PROTOCOL = 38
39SO_DOMAIN = 39
40
41ETH_P_IP = 0x0800
42ETH_P_IPV6 = 0x86dd
43
44IPPROTO_GRE = 47
45
46SIOCSIFHWADDR = 0x8924
47
48IPV6_FL_A_GET = 0
49IPV6_FL_A_PUT = 1
50IPV6_FL_A_RENEW = 1
51
52IPV6_FL_F_CREATE = 1
53IPV6_FL_F_EXCL = 2
54
55IPV6_FL_S_NONE = 0
56IPV6_FL_S_EXCL = 1
57IPV6_FL_S_ANY = 255
58
59IFNAMSIZ = 16
60
61IPV4_PING = "\x08\x00\x00\x00\x0a\xce\x00\x03"
62IPV6_PING = "\x80\x00\x00\x00\x0a\xce\x00\x03"
63
64IPV4_ADDR = "8.8.8.8"
65IPV6_ADDR = "2001:4860:4860::8888"
66
67IPV6_SEQ_DGRAM_HEADER = ("  sl  "
68                         "local_address                         "
69                         "remote_address                        "
70                         "st tx_queue rx_queue tr tm->when retrnsmt"
71                         "   uid  timeout inode ref pointer drops\n")
72
73# Arbitrary packet payload.
74UDP_PAYLOAD = str(scapy.DNS(rd=1,
75                            id=random.randint(0, 65535),
76                            qd=scapy.DNSQR(qname="wWW.GoOGle.CoM",
77                                           qtype="AAAA")))
78
79# Unix group to use if we want to open sockets as non-root.
80AID_INET = 3003
81
82
83def LinuxVersion():
84  # Example: "3.4.67-00753-gb7a556f".
85  # Get the part before the dash.
86  version = os.uname()[2].split("-")[0]
87  # Convert it into a tuple such as (3, 4, 67). That allows comparing versions
88  # using < and >, since tuples are compared lexicographically.
89  version = tuple(int(i) for i in version.split("."))
90  return version
91
92
93LINUX_VERSION = LinuxVersion()
94
95
96def SetSocketTimeout(sock, ms):
97  s = ms / 1000
98  us = (ms % 1000) * 1000
99  sock.setsockopt(SOL_SOCKET, SO_RCVTIMEO, struct.pack("LL", s, us))
100
101
102def SetSocketTos(s, tos):
103  level = {AF_INET: SOL_IP, AF_INET6: SOL_IPV6}[s.family]
104  option = {AF_INET: IP_TOS, AF_INET6: IPV6_TCLASS}[s.family]
105  s.setsockopt(level, option, tos)
106
107
108def SetNonBlocking(fd):
109  flags = fcntl.fcntl(fd, fcntl.F_GETFL, 0)
110  fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
111
112
113# Convenience functions to create sockets.
114def Socket(family, sock_type, protocol):
115  s = socket(family, sock_type, protocol)
116  SetSocketTimeout(s, 1000)
117  return s
118
119
120def PingSocket(family):
121  proto = {AF_INET: IPPROTO_ICMP, AF_INET6: IPPROTO_ICMPV6}[family]
122  return Socket(family, SOCK_DGRAM, proto)
123
124
125def IPv4PingSocket():
126  return PingSocket(AF_INET)
127
128
129def IPv6PingSocket():
130  return PingSocket(AF_INET6)
131
132
133def TCPSocket(family):
134  s = Socket(family, SOCK_STREAM, IPPROTO_TCP)
135  SetNonBlocking(s.fileno())
136  return s
137
138
139def IPv4TCPSocket():
140  return TCPSocket(AF_INET)
141
142
143def IPv6TCPSocket():
144  return TCPSocket(AF_INET6)
145
146
147def UDPSocket(family):
148  return Socket(family, SOCK_DGRAM, IPPROTO_UDP)
149
150
151def RawGRESocket(family):
152  s = Socket(family, SOCK_RAW, IPPROTO_GRE)
153  return s
154
155
156def DisableLinger(sock):
157  sock.setsockopt(SOL_SOCKET, SO_LINGER, struct.pack("ii", 1, 0))
158
159
160def CreateSocketPair(family, socktype, addr):
161  clientsock = socket(family, socktype, 0)
162  listensock = socket(family, socktype, 0)
163  listensock.bind((addr, 0))
164  addr = listensock.getsockname()
165  listensock.listen(1)
166  clientsock.connect(addr)
167  acceptedsock, _ = listensock.accept()
168  DisableLinger(clientsock)
169  DisableLinger(acceptedsock)
170  listensock.close()
171  return clientsock, acceptedsock
172
173
174def GetInterfaceIndex(ifname):
175  s = IPv4PingSocket()
176  ifr = struct.pack("%dsi" % IFNAMSIZ, ifname, 0)
177  ifr = fcntl.ioctl(s, scapy.SIOCGIFINDEX, ifr)
178  return struct.unpack("%dsi" % IFNAMSIZ, ifr)[1]
179
180
181def SetInterfaceHWAddr(ifname, hwaddr):
182  s = IPv4PingSocket()
183  hwaddr = hwaddr.replace(":", "")
184  hwaddr = hwaddr.decode("hex")
185  if len(hwaddr) != 6:
186    raise ValueError("Unknown hardware address length %d" % len(hwaddr))
187  ifr = struct.pack("%dsH6s" % IFNAMSIZ, ifname, scapy.ARPHDR_ETHER, hwaddr)
188  fcntl.ioctl(s, SIOCSIFHWADDR, ifr)
189
190
191def SetInterfaceState(ifname, up):
192  s = IPv4PingSocket()
193  ifr = struct.pack("%dsH" % IFNAMSIZ, ifname, 0)
194  ifr = fcntl.ioctl(s, scapy.SIOCGIFFLAGS, ifr)
195  _, flags = struct.unpack("%dsH" % IFNAMSIZ, ifr)
196  if up:
197    flags |= scapy.IFF_UP
198  else:
199    flags &= ~scapy.IFF_UP
200  ifr = struct.pack("%dsH" % IFNAMSIZ, ifname, flags)
201  ifr = fcntl.ioctl(s, scapy.SIOCSIFFLAGS, ifr)
202
203
204def SetInterfaceUp(ifname):
205  return SetInterfaceState(ifname, True)
206
207
208def SetInterfaceDown(ifname):
209  return SetInterfaceState(ifname, False)
210
211
212def FormatProcAddress(unformatted):
213  groups = []
214  for i in xrange(0, len(unformatted), 4):
215    groups.append(unformatted[i:i+4])
216  formatted = ":".join(groups)
217  # Compress the address.
218  address = inet_ntop(AF_INET6, inet_pton(AF_INET6, formatted))
219  return address
220
221
222def FormatSockStatAddress(address):
223  if ":" in address:
224    family = AF_INET6
225  else:
226    family = AF_INET
227  binary = inet_pton(family, address)
228  out = ""
229  for i in xrange(0, len(binary), 4):
230    out += "%08X" % struct.unpack("=L", binary[i:i+4])
231  return out
232
233
234def GetLinkAddress(ifname, linklocal):
235  addresses = open("/proc/net/if_inet6").readlines()
236  for address in addresses:
237    address = [s for s in address.strip().split(" ") if s]
238    if address[5] == ifname:
239      if (linklocal and address[0].startswith("fe80")
240          or not linklocal and not address[0].startswith("fe80")):
241        # Convert the address from raw hex to something with colons in it.
242        return FormatProcAddress(address[0])
243  return None
244
245
246def GetDefaultRoute(version=6):
247  if version == 6:
248    routes = open("/proc/net/ipv6_route").readlines()
249    for route in routes:
250      route = [s for s in route.strip().split(" ") if s]
251      if (route[0] == "00000000000000000000000000000000" and route[1] == "00"
252          # Routes in non-default tables end up in /proc/net/ipv6_route!!!
253          and route[9] != "lo" and not route[9].startswith("nettest")):
254        return FormatProcAddress(route[4]), route[9]
255    raise ValueError("No IPv6 default route found")
256  elif version == 4:
257    routes = open("/proc/net/route").readlines()
258    for route in routes:
259      route = [s for s in route.strip().split("\t") if s]
260      if route[1] == "00000000" and route[7] == "00000000":
261        gw, iface = route[2], route[0]
262        gw = inet_ntop(AF_INET, gw.decode("hex")[::-1])
263        return gw, iface
264    raise ValueError("No IPv4 default route found")
265  else:
266    raise ValueError("Don't know about IPv%s" % version)
267
268
269def GetDefaultRouteInterface():
270  unused_gw, iface = GetDefaultRoute()
271  return iface
272
273
274def MakeFlowLabelOption(addr, label):
275  # struct in6_flowlabel_req {
276  #         struct in6_addr flr_dst;
277  #         __be32  flr_label;
278  #         __u8    flr_action;
279  #         __u8    flr_share;
280  #         __u16   flr_flags;
281  #         __u16   flr_expires;
282  #         __u16   flr_linger;
283  #         __u32   __flr_pad;
284  #         /* Options in format of IPV6_PKTOPTIONS */
285  # };
286  fmt = "16sIBBHHH4s"
287  assert struct.calcsize(fmt) == 32
288  addr = inet_pton(AF_INET6, addr)
289  assert len(addr) == 16
290  label = htonl(label & 0xfffff)
291  action = IPV6_FL_A_GET
292  share = IPV6_FL_S_ANY
293  flags = IPV6_FL_F_CREATE
294  pad = "\x00" * 4
295  return struct.pack(fmt, addr, label, action, share, flags, 0, 0, pad)
296
297
298def SetFlowLabel(s, addr, label):
299  opt = MakeFlowLabelOption(addr, label)
300  s.setsockopt(SOL_IPV6, IPV6_FLOWLABEL_MGR, opt)
301  # Caller also needs to do s.setsockopt(SOL_IPV6, IPV6_FLOWINFO_SEND, 1).
302
303
304# Determine network configuration.
305try:
306  GetDefaultRoute(version=4)
307  HAVE_IPV4 = True
308except ValueError:
309  HAVE_IPV4 = False
310
311try:
312  GetDefaultRoute(version=6)
313  HAVE_IPV6 = True
314except ValueError:
315  HAVE_IPV6 = False
316
317
318class RunAsUid(object):
319  """Context guard to run a code block as a given UID."""
320
321  def __init__(self, uid):
322    self.uid = uid
323
324  def __enter__(self):
325    if self.uid:
326      self.saved_uid = os.geteuid()
327      self.saved_groups = os.getgroups()
328      if self.uid:
329        os.setgroups(self.saved_groups + [AID_INET])
330        os.seteuid(self.uid)
331
332  def __exit__(self, unused_type, unused_value, unused_traceback):
333    if self.uid:
334      os.seteuid(self.saved_uid)
335      os.setgroups(self.saved_groups)
336
337
338class NetworkTest(unittest.TestCase):
339
340  def assertRaisesErrno(self, err_num, f, *args):
341    msg = os.strerror(err_num)
342    self.assertRaisesRegexp(EnvironmentError, msg, f, *args)
343
344  def ReadProcNetSocket(self, protocol):
345    # Read file.
346    filename = "/proc/net/%s" % protocol
347    lines = open(filename).readlines()
348
349    # Possibly check, and strip, header.
350    if protocol in ["icmp6", "raw6", "udp6"]:
351      self.assertEqual(IPV6_SEQ_DGRAM_HEADER, lines[0])
352    lines = lines[1:]
353
354    # Check contents.
355    if protocol.endswith("6"):
356      addrlen = 32
357    else:
358      addrlen = 8
359
360    if protocol.startswith("tcp"):
361      # Real sockets have 5 extra numbers, timewait sockets have none.
362      end_regexp = "(| +[0-9]+ [0-9]+ [0-9]+ [0-9]+ -?[0-9]+|)$"
363    elif re.match("icmp|udp|raw", protocol):
364      # Drops.
365      end_regexp = " +([0-9]+) *$"
366    else:
367      raise ValueError("Don't know how to parse %s" % filename)
368
369    regexp = re.compile(r" *(\d+): "                    # bucket
370                        "([0-9A-F]{%d}:[0-9A-F]{4}) "   # srcaddr, port
371                        "([0-9A-F]{%d}:[0-9A-F]{4}) "   # dstaddr, port
372                        "([0-9A-F][0-9A-F]) "           # state
373                        "([0-9A-F]{8}:[0-9A-F]{8}) "    # mem
374                        "([0-9A-F]{2}:[0-9A-F]{8}) "    # ?
375                        "([0-9A-F]{8}) +"               # ?
376                        "([0-9]+) +"                    # uid
377                        "([0-9]+) +"                    # timeout
378                        "([0-9]+) +"                    # inode
379                        "([0-9]+) +"                    # refcnt
380                        "([0-9a-f]+)"                   # sp
381                        "%s"                            # icmp has spaces
382                        % (addrlen, addrlen, end_regexp))
383    # Return a list of lists with only source / dest addresses for now.
384    # TODO: consider returning a dict or namedtuple instead.
385    out = []
386    for line in lines:
387      (_, src, dst, state, mem,
388       _, _, uid, _, _, refcnt, _, extra) = regexp.match(line).groups()
389      out.append([src, dst, state, mem, uid, refcnt, extra])
390    return out
391
392
393if __name__ == "__main__":
394  unittest.main()
395