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