import ctypes
import functools
import io
import json
import os
import pathlib
import shlex
import signal
import socket
import subprocess
import tempfile
import unittest
from scapy.all import StreamSocket, sndrcv, Ether, conf, Route, ARP
SLIRPHELPER = os.environ.get("SLIRPHELPER")
LIBC = ctypes.CDLL("libc.so.6")
CLONE_NEWNET = 0x40000000
ORIGINAL_NET_NS = open("/proc/self/ns/net", "rb")
THISDIR = pathlib.Path(__file__).parent.absolute()
DBUS_SESSION_BUS_ADDRESS = os.environ.get("DBUS_SESSION_BUS_ADDRESS")
@functools.lru_cache()
def helper_capabilities():
p = subprocess.run(
[SLIRPHELPER, "--print-capabilities"], stdout=subprocess.PIPE, text=True
)
return json.loads(p.stdout)
def has_cap(cap):
return cap in helper_capabilities()["features"]
class Process:
def __init__(self, argv, close_fds=True, env=None):
self.p = subprocess.Popen(
argv,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
close_fds=close_fds,
env=env,
)
self.rc = None
def stdout_all(self):
return self.p.stdout.read()
def stdout_line(self):
return self.p.stdout.readline()
def stderr_all(self):
return self.p.stderr.read()
def stderr_line(self):
return self.p.stderr.readline()
def close(self, kill=True):
if self.p:
if kill:
self.p.send_signal(signal.SIGINT)
self.p.send_signal(signal.SIGTERM)
self.rc = self.p.wait()
self.p.stderr.close()
self.p.stdout.close()
self.p = None
return self.rc
def graceful_stop(self, wait=True):
self.p.send_signal(signal.SIGINT)
if wait:
self.p.wait()
class TestCase(unittest.TestCase):
has_notify_socket = None
execno = 0
def setUp(self):
if self.has_notify_socket is None:
self.has_notify_socket = has_cap("notify-socket")
self.cleanups = None
prev_net_fd = open("/proc/self/ns/net", "rb")
r = LIBC.unshare(CLONE_NEWNET)
if r != 0:
self.fail('Are you running within "unshare -Ur" ? Need unshare() syscall.')
self.guest_net_fd = open("/proc/self/ns/net", "rb")
self._add_teardown(self.guest_net_fd)
os.system(
"ip link set lo up;"
"ip tuntap add mode tap name tun0;"
"ip link set tun0 mtu 65521;"
"ip link set tun0 up;"
"ip addr add 10.0.2.100/24 dev tun0;"
"ip addr add 2001:2::100/32 dev tun0 nodad;"
"ip route add 0.0.0.0/0 via 10.0.2.2 dev tun0;"
"ip route add ::/0 via 2001:2::2 dev tun0;"
)
w = subprocess.Popen(["/bin/sleep", "1073741824"])
self.guest_ns_pid = w.pid
self._add_teardown(w)
LIBC.setns(prev_net_fd.fileno(), CLONE_NEWNET)
prev_net_fd.close()
self._tmpdir = tempfile.TemporaryDirectory()
self._add_teardown(self._tmpdir)
def tearDown(self):
while self.cleanups:
item = self.cleanups.pop()
if isinstance(item, subprocess.Popen):
item.send_signal(signal.SIGINT)
item.wait()
elif isinstance(item, Process):
item.close()
if getattr(item, "stdout", None):
item.stdout.close()
if getattr(item, "stderr", None):
item.stderr.close()
elif isinstance(item, io.BufferedReader):
item.close()
elif isinstance(item, tempfile.TemporaryDirectory):
item.cleanup()
else:
print("Unknown cleanup type")
print(type(item))
def run_helper(self, argv1=[], wait_ready=True, netns=True):
if isinstance(argv1, str):
argv1 = shlex.split(argv1)
a = [SLIRPHELPER] + argv1
if netns:
a = a + ["--netns", self.net_ns_path(), "--interface", "tun0"]
sn = None
env = None
if self.has_notify_socket and wait_ready:
self.execno += 1
sn = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
path = self.get_tmp_filename("sn-%d" % self.execno)
sn.bind(path)
env = dict(os.environ, NOTIFY_SOCKET="%s" % path)
p = Process(a, close_fds=False, env=env)
if sn:
sn.settimeout(1) try:
self.assertIn("READY=1", sn.recv(4096).decode())
except:
print(p.stderr_all())
sn.close()
self._add_teardown(p)
return p
def skipIfNotCapable(self, cap):
if not has_cap(cap):
self.skipTest("since '%s' capability is missing" % cap)
def start_echo(self, udp=False):
cmd = [THISDIR / "echo.py"]
if udp:
cmd += ["-u"]
p = Process(cmd)
self._add_teardown(p)
return int(p.stdout_line())
def assertTcpEcho(self, ip, port):
data = os.getrandom(16)
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.connect((ip, port))
s.sendall(data)
self.assertEqual(s.recv(len(data)), data)
def assertUdpEcho(self, ip, port):
data = os.getrandom(16)
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
s.sendto(data, (ip, port))
self.assertEqual(s.recv(len(data)), data)
def get_tmp_filename(self, name):
return os.path.join(self._tmpdir.name, name)
def _add_teardown(self, item):
if not self.cleanups:
self.cleanups = []
self.cleanups.append(item)
def net_ns_path(self):
return "/proc/%s/ns/net" % self.guest_ns_pid
def guest_netns(self):
xself = self
class controlled_execution:
def __enter__(self):
self.prev_net_fd = open("/proc/self/ns/net", "rb")
LIBC.setns(xself.guest_net_fd.fileno(), CLONE_NEWNET)
def __exit__(self, type, value, traceback):
LIBC.setns(self.prev_net_fd.fileno(), CLONE_NEWNET)
self.prev_net_fd.close()
return controlled_execution()
class testScapySocket:
def __init__(self, fd):
ss = StreamSocket(fd)
ss.basecls = Ether
self.ss = ss
conf.route = Route() self.e = Ether(src="52:55:0a:00:02:42")
def send(self, x):
self.ss.send(self.e / x)
def recv(self, x):
return self.ss.recv(x)
def fileno(self):
return self.ss.fileno()
def sr1(self, x, checkIPaddr=True, *args, **kwargs):
conf.checkIPaddr = checkIPaddr
kwargs.setdefault("verbose", False)
ans, _ = sndrcv(self.ss, self.e / x, *args, **kwargs)
return ans[0][1]
def sr(self, x, checkIPaddr=True, *args, **kwargs):
conf.checkIPaddr = checkIPaddr
kwargs.setdefault("verbose", False)
return sndrcv(self.ss, self.e / x, *args, **kwargs)
def withScapy():
def decorate(fn):
@functools.wraps(fn)
def maybe(*args, **kw):
sp = socket.socketpair(type=socket.SOCK_DGRAM)
os.set_inheritable(sp[0].fileno(), True)
self = args[0]
arg = kw.pop("parg", "")
p = self.run_helper(arg + " --fd %d" % sp[0].fileno(), netns=False)
s = testScapySocket(sp[1])
s.send(ARP(psrc="10.0.2.100", pdst="10.0.2.100", hwsrc=s.e.src))
kw["s"] = s
ret = fn(*args, **kw)
sp[0].close()
sp[1].close()
return ret
return maybe
return decorate
def isolateHostNetwork():
def decorate(fn):
@functools.wraps(fn)
def maybe(*args, **kw):
prev_net_fd = open("/proc/self/ns/net", "rb")
r = LIBC.unshare(CLONE_NEWNET)
if r != 0:
self.fail(
'Are you running within "unshare -Ur" ? Need unshare() syscall.'
)
os.system(
"ip link set lo up;"
"ip tuntap add mode tun name eth0;"
"ip link set eth0 mtu 65521;"
"ip link set eth0 up;"
"ip addr add 192.168.1.100/24 dev eth0;"
"ip addr add 3ffe::100/16 dev eth0 nodad;"
"ip route add 0.0.0.0/0 via 192.168.1.1 dev eth0;"
"ip route add ::/0 via 3ffe::1 dev eth0;"
)
ret = fn(*args, **kw)
LIBC.setns(prev_net_fd.fileno(), CLONE_NEWNET)
prev_net_fd.close()
return ret
return maybe
return decorate