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
19from socket import *  # pylint: disable=wildcard-import
20import struct
21import unittest
22
23from scapy import all as scapy
24
25SOL_IPV6 = 41
26IP_RECVERR = 11
27IPV6_RECVERR = 25
28IP_TRANSPARENT = 19
29IPV6_TRANSPARENT = 75
30IPV6_TCLASS = 67
31SO_BINDTODEVICE = 25
32SO_MARK = 36
33IPV6_FLOWLABEL_MGR = 32
34IPV6_FLOWINFO_SEND = 33
35
36ETH_P_IP = 0x0800
37ETH_P_IPV6 = 0x86dd
38
39IPPROTO_GRE = 47
40
41SIOCSIFHWADDR = 0x8924
42
43IPV6_FL_A_GET = 0
44IPV6_FL_A_PUT = 1
45IPV6_FL_A_RENEW = 1
46
47IPV6_FL_F_CREATE = 1
48IPV6_FL_F_EXCL = 2
49
50IPV6_FL_S_NONE = 0
51IPV6_FL_S_EXCL = 1
52IPV6_FL_S_ANY = 255
53
54IPV4_PING = "\x08\x00\x00\x00\x0a\xce\x00\x03"
55IPV6_PING = "\x80\x00\x00\x00\x0a\xce\x00\x03"
56
57IPV4_ADDR = "8.8.8.8"
58IPV6_ADDR = "2001:4860:4860::8888"
59
60IPV6_SEQ_DGRAM_HEADER = ("  sl  "
61                         "local_address                         "
62                         "remote_address                        "
63                         "st tx_queue rx_queue tr tm->when retrnsmt"
64                         "   uid  timeout inode ref pointer drops\n")
65
66# Unix group to use if we want to open sockets as non-root.
67AID_INET = 3003
68
69
70def LinuxVersion():
71  # Example: "3.4.67-00753-gb7a556f".
72  # Get the part before the dash.
73  version = os.uname()[2].split("-")[0]
74  # Convert it into a tuple such as (3, 4, 67). That allows comparing versions
75  # using < and >, since tuples are compared lexicographically.
76  version = tuple(int(i) for i in version.split("."))
77  return version
78
79
80LINUX_VERSION = LinuxVersion()
81
82
83def SetSocketTimeout(sock, ms):
84  s = ms / 1000
85  us = (ms % 1000) * 1000
86  sock.setsockopt(SOL_SOCKET, SO_RCVTIMEO, struct.pack("LL", s, us))
87
88
89def SetSocketTos(s, tos):
90  level = {AF_INET: SOL_IP, AF_INET6: SOL_IPV6}[s.family]
91  option = {AF_INET: IP_TOS, AF_INET6: IPV6_TCLASS}[s.family]
92  s.setsockopt(level, option, tos)
93
94
95def SetNonBlocking(fd):
96  flags = fcntl.fcntl(fd, fcntl.F_GETFL, 0)
97  fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
98
99
100# Convenience functions to create sockets.
101def Socket(family, sock_type, protocol):
102  s = socket(family, sock_type, protocol)
103  SetSocketTimeout(s, 1000)
104  return s
105
106
107def PingSocket(family):
108  proto = {AF_INET: IPPROTO_ICMP, AF_INET6: IPPROTO_ICMPV6}[family]
109  return Socket(family, SOCK_DGRAM, proto)
110
111
112def IPv4PingSocket():
113  return PingSocket(AF_INET)
114
115
116def IPv6PingSocket():
117  return PingSocket(AF_INET6)
118
119
120def TCPSocket(family):
121  s = Socket(family, SOCK_STREAM, IPPROTO_TCP)
122  SetNonBlocking(s.fileno())
123  return s
124
125
126def IPv4TCPSocket():
127  return TCPSocket(AF_INET)
128
129
130def IPv6TCPSocket():
131  return TCPSocket(AF_INET6)
132
133
134def UDPSocket(family):
135  return Socket(family, SOCK_DGRAM, IPPROTO_UDP)
136
137
138def RawGRESocket(family):
139  s = Socket(family, SOCK_RAW, IPPROTO_GRE)
140  return s
141
142
143def GetInterfaceIndex(ifname):
144  s = IPv4PingSocket()
145  ifr = struct.pack("16si", ifname, 0)
146  ifr = fcntl.ioctl(s, scapy.SIOCGIFINDEX, ifr)
147  return struct.unpack("16si", ifr)[1]
148
149
150def SetInterfaceHWAddr(ifname, hwaddr):
151  s = IPv4PingSocket()
152  hwaddr = hwaddr.replace(":", "")
153  hwaddr = hwaddr.decode("hex")
154  if len(hwaddr) != 6:
155    raise ValueError("Unknown hardware address length %d" % len(hwaddr))
156  ifr = struct.pack("16sH6s", ifname, scapy.ARPHDR_ETHER, hwaddr)
157  fcntl.ioctl(s, SIOCSIFHWADDR, ifr)
158
159
160def SetInterfaceState(ifname, up):
161  s = IPv4PingSocket()
162  ifr = struct.pack("16sH", ifname, 0)
163  ifr = fcntl.ioctl(s, scapy.SIOCGIFFLAGS, ifr)
164  _, flags = struct.unpack("16sH", ifr)
165  if up:
166    flags |= scapy.IFF_UP
167  else:
168    flags &= ~scapy.IFF_UP
169  ifr = struct.pack("16sH", ifname, flags)
170  ifr = fcntl.ioctl(s, scapy.SIOCSIFFLAGS, ifr)
171
172
173def SetInterfaceUp(ifname):
174  return SetInterfaceState(ifname, True)
175
176
177def SetInterfaceDown(ifname):
178  return SetInterfaceState(ifname, False)
179
180
181def FormatProcAddress(unformatted):
182  groups = []
183  for i in xrange(0, len(unformatted), 4):
184    groups.append(unformatted[i:i+4])
185  formatted = ":".join(groups)
186  # Compress the address.
187  address = inet_ntop(AF_INET6, inet_pton(AF_INET6, formatted))
188  return address
189
190
191def FormatSockStatAddress(address):
192  if ":" in address:
193    family = AF_INET6
194  else:
195    family = AF_INET
196  binary = inet_pton(family, address)
197  out = ""
198  for i in xrange(0, len(binary), 4):
199    out += "%08X" % struct.unpack("=L", binary[i:i+4])
200  return out
201
202
203def GetLinkAddress(ifname, linklocal):
204  addresses = open("/proc/net/if_inet6").readlines()
205  for address in addresses:
206    address = [s for s in address.strip().split(" ") if s]
207    if address[5] == ifname:
208      if (linklocal and address[0].startswith("fe80")
209          or not linklocal and not address[0].startswith("fe80")):
210        # Convert the address from raw hex to something with colons in it.
211        return FormatProcAddress(address[0])
212  return None
213
214
215def GetDefaultRoute(version=6):
216  if version == 6:
217    routes = open("/proc/net/ipv6_route").readlines()
218    for route in routes:
219      route = [s for s in route.strip().split(" ") if s]
220      if (route[0] == "00000000000000000000000000000000" and route[1] == "00"
221          # Routes in non-default tables end up in /proc/net/ipv6_route!!!
222          and route[9] != "lo" and not route[9].startswith("nettest")):
223        return FormatProcAddress(route[4]), route[9]
224    raise ValueError("No IPv6 default route found")
225  elif version == 4:
226    routes = open("/proc/net/route").readlines()
227    for route in routes:
228      route = [s for s in route.strip().split("\t") if s]
229      if route[1] == "00000000" and route[7] == "00000000":
230        gw, iface = route[2], route[0]
231        gw = inet_ntop(AF_INET, gw.decode("hex")[::-1])
232        return gw, iface
233    raise ValueError("No IPv4 default route found")
234  else:
235    raise ValueError("Don't know about IPv%s" % version)
236
237
238def GetDefaultRouteInterface():
239  unused_gw, iface = GetDefaultRoute()
240  return iface
241
242
243def MakeFlowLabelOption(addr, label):
244  # struct in6_flowlabel_req {
245  #         struct in6_addr flr_dst;
246  #         __be32  flr_label;
247  #         __u8    flr_action;
248  #         __u8    flr_share;
249  #         __u16   flr_flags;
250  #         __u16   flr_expires;
251  #         __u16   flr_linger;
252  #         __u32   __flr_pad;
253  #         /* Options in format of IPV6_PKTOPTIONS */
254  # };
255  fmt = "16sIBBHHH4s"
256  assert struct.calcsize(fmt) == 32
257  addr = inet_pton(AF_INET6, addr)
258  assert len(addr) == 16
259  label = htonl(label & 0xfffff)
260  action = IPV6_FL_A_GET
261  share = IPV6_FL_S_ANY
262  flags = IPV6_FL_F_CREATE
263  pad = "\x00" * 4
264  return struct.pack(fmt, addr, label, action, share, flags, 0, 0, pad)
265
266
267def SetFlowLabel(s, addr, label):
268  opt = MakeFlowLabelOption(addr, label)
269  s.setsockopt(SOL_IPV6, IPV6_FLOWLABEL_MGR, opt)
270  # Caller also needs to do s.setsockopt(SOL_IPV6, IPV6_FLOWINFO_SEND, 1).
271
272
273# Determine network configuration.
274try:
275  GetDefaultRoute(version=4)
276  HAVE_IPV4 = True
277except ValueError:
278  HAVE_IPV4 = False
279
280try:
281  GetDefaultRoute(version=6)
282  HAVE_IPV6 = True
283except ValueError:
284  HAVE_IPV6 = False
285
286
287class RunAsUid(object):
288
289  """Context guard to run a code block as a given UID."""
290
291  def __init__(self, uid):
292    self.uid = uid
293
294  def __enter__(self):
295    if self.uid:
296      self.saved_uid = os.geteuid()
297      self.saved_groups = os.getgroups()
298      if self.uid:
299        os.setgroups(self.saved_groups + [AID_INET])
300        os.seteuid(self.uid)
301
302  def __exit__(self, unused_type, unused_value, unused_traceback):
303    if self.uid:
304      os.seteuid(self.saved_uid)
305      os.setgroups(self.saved_groups)
306
307
308class NetworkTest(unittest.TestCase):
309
310  def assertRaisesErrno(self, err_num, f, *args):
311    msg = os.strerror(err_num)
312    self.assertRaisesRegexp(EnvironmentError, msg, f, *args)
313
314
315if __name__ == "__main__":
316  unittest.main()
317