From 2fa0cd06fb47ade553a7d36b9efae0225941ffde Mon Sep 17 00:00:00 2001 From: max Date: Fri, 8 Sep 2017 03:17:37 +0200 Subject: [PATCH] Route traffic by linux user --- sshuttle/client.py | 31 +++++++++++++--- sshuttle/cmdline.py | 3 +- sshuttle/firewall.py | 17 +++++---- sshuttle/methods/__init__.py | 7 ++-- sshuttle/methods/ipfw.py | 6 ++-- sshuttle/methods/nat.py | 38 ++++++++++++++++---- sshuttle/methods/pf.py | 4 +-- sshuttle/methods/tproxy.py | 6 ++-- sshuttle/options.py | 6 ++++ sshuttle/tests/client/test_firewall.py | 12 ++++--- sshuttle/tests/client/test_methods_nat.py | 11 +++--- sshuttle/tests/client/test_methods_pf.py | 33 ++++++++++------- sshuttle/tests/client/test_methods_tproxy.py | 10 +++--- 13 files changed, 130 insertions(+), 54 deletions(-) diff --git a/sshuttle/client.py b/sshuttle/client.py index 54022ad..45dd6e2 100644 --- a/sshuttle/client.py +++ b/sshuttle/client.py @@ -15,6 +15,10 @@ from sshuttle.ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper from sshuttle.helpers import log, debug1, debug2, debug3, Fatal, islocal, \ resolvconf_nameservers from sshuttle.methods import get_method, Features +try: + from pwd import getpwnam +except ImportError: + getpwnam = None try: # try getting recvmsg from python @@ -238,7 +242,8 @@ class FirewallClient: self.method.set_firewall(self) def setup(self, subnets_include, subnets_exclude, nslist, - redirectport_v6, redirectport_v4, dnsport_v6, dnsport_v4, udp): + redirectport_v6, redirectport_v4, dnsport_v6, dnsport_v4, udp, + user): self.subnets_include = subnets_include self.subnets_exclude = subnets_exclude self.nslist = nslist @@ -247,6 +252,7 @@ class FirewallClient: self.dnsport_v6 = dnsport_v6 self.dnsport_v4 = dnsport_v4 self.udp = udp + self.user = user def check(self): rv = self.p.poll() @@ -276,8 +282,14 @@ class FirewallClient: udp = 0 if self.udp: udp = 1 + if self.user is None: + user = b'-' + elif isinstance(self.user, str): + user = bytes(self.user, 'utf-8') + else: + user = b'%d' % self.user - self.pfile.write(b'GO %d\n' % udp) + self.pfile.write(b'GO %d %s\n' % (udp, user)) self.pfile.flush() line = self.pfile.readline() @@ -536,7 +548,8 @@ def _main(tcp_listener, udp_listener, fw, ssh_cmd, remotename, def main(listenip_v6, listenip_v4, ssh_cmd, remotename, python, latency_control, dns, nslist, method_name, seed_hosts, auto_hosts, auto_nets, - subnets_include, subnets_exclude, daemon, to_nameserver, pidfile): + subnets_include, subnets_exclude, daemon, to_nameserver, pidfile, + user): if daemon: try: @@ -573,10 +586,19 @@ def main(listenip_v6, listenip_v4, else: listenip_v6 = None + if user is not None: + if getpwnam is None: + raise Fatal("Routing by user not available on this system.") + try: + user = getpwnam(user).pw_uid + except KeyError: + raise Fatal("User %s does not exist." % user) + required.ipv6 = len(subnets_v6) > 0 or listenip_v6 is not None required.ipv4 = len(subnets_v4) > 0 or listenip_v4 is not None required.udp = avail.udp required.dns = len(nslist) > 0 + required.user = False if user is None else True # if IPv6 not supported, ignore IPv6 DNS servers if not required.ipv6: @@ -592,6 +614,7 @@ def main(listenip_v6, listenip_v4, debug1("IPv6 enabled: %r\n" % required.ipv6) debug1("UDP enabled: %r\n" % required.udp) debug1("DNS enabled: %r\n" % required.dns) + debug1("User enabled: %r\n" % required.user) # bind to required ports if listenip_v4 == "auto": @@ -742,7 +765,7 @@ def main(listenip_v6, listenip_v4, # start the firewall fw.setup(subnets_include, subnets_exclude, nslist, redirectport_v6, redirectport_v4, dnsport_v6, dnsport_v4, - required.udp) + required.udp, user) # start the client process try: diff --git a/sshuttle/cmdline.py b/sshuttle/cmdline.py index ecb01e3..e8342ff 100644 --- a/sshuttle/cmdline.py +++ b/sshuttle/cmdline.py @@ -75,7 +75,8 @@ def main(): excludes, opt.daemon, opt.to_ns, - opt.pidfile) + opt.pidfile, + opt.user) if return_code == 0: log('Normal exit code, exiting...') diff --git a/sshuttle/firewall.py b/sshuttle/firewall.py index b275d08..faf3c55 100644 --- a/sshuttle/firewall.py +++ b/sshuttle/firewall.py @@ -188,9 +188,12 @@ def main(method_name, syslog): elif not line.startswith("GO "): raise Fatal('firewall: expected GO but got %r' % line) - _, _, udp = line.partition(" ") + _, _, args = line.partition(" ") + udp, user = args.strip().split(" ", 1) udp = bool(int(udp)) - debug2('firewall manager: Got udp: %r\n' % udp) + if user == '-': + user = None + debug2('firewall manager: Got udp: %r, user: %r\n' % (udp, user)) subnets_v6 = [i for i in subnets if i[0] == socket.AF_INET6] nslist_v6 = [i for i in nslist if i[0] == socket.AF_INET6] @@ -204,13 +207,15 @@ def main(method_name, syslog): debug2('firewall manager: setting up IPv6.\n') method.setup_firewall( port_v6, dnsport_v6, nslist_v6, - socket.AF_INET6, subnets_v6, udp) + socket.AF_INET6, subnets_v6, udp, + user) if len(subnets_v4) > 0 or len(nslist_v4) > 0: debug2('firewall manager: setting up IPv4.\n') method.setup_firewall( port_v4, dnsport_v4, nslist_v4, - socket.AF_INET, subnets_v4, udp) + socket.AF_INET, subnets_v4, udp, + user) stdout.write('STARTED\n') @@ -246,7 +251,7 @@ def main(method_name, syslog): try: if len(subnets_v6) > 0 or len(nslist_v6) > 0: debug2('firewall manager: undoing IPv6 changes.\n') - method.restore_firewall(port_v6, socket.AF_INET6, udp) + method.restore_firewall(port_v6, socket.AF_INET6, udp, user) except: try: debug1("firewall manager: " @@ -259,7 +264,7 @@ def main(method_name, syslog): try: if len(subnets_v4) > 0 or len(nslist_v4) > 0: debug2('firewall manager: undoing IPv4 changes.\n') - method.restore_firewall(port_v4, socket.AF_INET, udp) + method.restore_firewall(port_v4, socket.AF_INET, udp, user) except: try: debug1("firewall manager: " diff --git a/sshuttle/methods/__init__.py b/sshuttle/methods/__init__.py index 8cce0e7..6359d19 100644 --- a/sshuttle/methods/__init__.py +++ b/sshuttle/methods/__init__.py @@ -40,6 +40,7 @@ class BaseMethod(object): result.ipv6 = False result.udp = False result.dns = True + result.user = False return result def get_tcp_dstip(self, sock): @@ -64,16 +65,16 @@ class BaseMethod(object): def assert_features(self, features): avail = self.get_supported_features() - for key in ["udp", "dns", "ipv6"]: + for key in ["udp", "dns", "ipv6", "user"]: if getattr(features, key) and not getattr(avail, key): raise Fatal( "Feature %s not supported with method %s.\n" % (key, self.name)) - def setup_firewall(self, port, dnsport, nslist, family, subnets, udp): + def setup_firewall(self, port, dnsport, nslist, family, subnets, udp, user): raise NotImplementedError() - def restore_firewall(self, port, family, udp): + def restore_firewall(self, port, family, udp, user): raise NotImplementedError() def firewall_command(self, line): diff --git a/sshuttle/methods/ipfw.py b/sshuttle/methods/ipfw.py index 97d2ca6..e696b1f 100644 --- a/sshuttle/methods/ipfw.py +++ b/sshuttle/methods/ipfw.py @@ -193,9 +193,9 @@ class Method(BaseMethod): #if udp_listener.v6 is not None: # udp_listener.v6.setsockopt(SOL_IPV6, IPV6_RECVDSTADDR, 1) - def setup_firewall(self, port, dnsport, nslist, family, subnets, udp): + def setup_firewall(self, port, dnsport, nslist, family, subnets, udp, user): # IPv6 not supported - if family not in [socket.AF_INET ]: + if family not in [socket.AF_INET]: raise Exception( 'Address family "%s" unsupported by ipfw method_name' % family_to_string(family)) @@ -255,7 +255,7 @@ class Method(BaseMethod): else: ipfw('table', '126', 'add', '%s/%s' % (snet, swidth)) - def restore_firewall(self, port, family, udp): + def restore_firewall(self, port, family, udp, user): if family not in [socket.AF_INET]: raise Exception( 'Address family "%s" unsupported by tproxy method' diff --git a/sshuttle/methods/nat.py b/sshuttle/methods/nat.py index a5e7b96..6706003 100644 --- a/sshuttle/methods/nat.py +++ b/sshuttle/methods/nat.py @@ -12,7 +12,7 @@ class Method(BaseMethod): # the multiple copies shouldn't have overlapping subnets, or only the most- # recently-started one will win (because we use "-I OUTPUT 1" instead of # "-A OUTPUT"). - def setup_firewall(self, port, dnsport, nslist, family, subnets, udp): + def setup_firewall(self, port, dnsport, nslist, family, subnets, udp, user): # only ipv4 supported with NAT if family != socket.AF_INET: raise Exception( @@ -29,15 +29,25 @@ class Method(BaseMethod): def _ipt_ttl(*args): return ipt_ttl(family, table, *args) + def _ipm(*args): + return ipt(family, "mangle", *args) + chain = 'sshuttle-%s' % port # basic cleanup/setup of chains - self.restore_firewall(port, family, udp) + self.restore_firewall(port, family, udp, user) _ipt('-N', chain) _ipt('-F', chain) - _ipt('-I', 'OUTPUT', '1', '-j', chain) - _ipt('-I', 'PREROUTING', '1', '-j', chain) + if user is not None: + _ipm('-I', 'OUTPUT', '1', '-m', 'owner', '--uid-owner', str(user), + '-j', 'MARK', '--set-mark', str(port)) + args = '-m', 'mark', '--mark', str(port), '-j', chain + else: + args = '-j', chain + + _ipt('-I', 'OUTPUT', '1', *args) + _ipt('-I', 'PREROUTING', '1', *args) # create new subnet entries. for f, swidth, sexclude, snet, fport, lport \ @@ -62,7 +72,7 @@ class Method(BaseMethod): '--dport', '53', '--to-ports', str(dnsport)) - def restore_firewall(self, port, family, udp): + def restore_firewall(self, port, family, udp, user): # only ipv4 supported with NAT if family != socket.AF_INET: raise Exception( @@ -79,11 +89,25 @@ class Method(BaseMethod): def _ipt_ttl(*args): return ipt_ttl(family, table, *args) + def _ipm(*args): + return ipt(family, "mangle", *args) + chain = 'sshuttle-%s' % port # basic cleanup/setup of chains if ipt_chain_exists(family, table, chain): - nonfatal(_ipt, '-D', 'OUTPUT', '-j', chain) - nonfatal(_ipt, '-D', 'PREROUTING', '-j', chain) + if user is not None: + nonfatal(_ipm, '-D', 'OUTPUT', '-m', 'owner', '--uid-owner', str(user), + '-j', 'MARK', '--set-mark', str(port)) + args = '-m', 'mark', '--mark', str(port), '-j', chain + else: + args = '-j', chain + nonfatal(_ipt, '-D', 'OUTPUT', *args) + nonfatal(_ipt, '-D', 'PREROUTING', *args) nonfatal(_ipt, '-F', chain) _ipt('-X', chain) + + def get_supported_features(self): + result = super(Method, self).get_supported_features() + result.user = True + return result diff --git a/sshuttle/methods/pf.py b/sshuttle/methods/pf.py index 498ba8c..444e0e0 100644 --- a/sshuttle/methods/pf.py +++ b/sshuttle/methods/pf.py @@ -417,7 +417,7 @@ class Method(BaseMethod): return sock.getsockname() - def setup_firewall(self, port, dnsport, nslist, family, subnets, udp): + def setup_firewall(self, port, dnsport, nslist, family, subnets, udp, user): tables = [] translating_rules = [] filtering_rules = [] @@ -446,7 +446,7 @@ class Method(BaseMethod): pf.add_rules(anchor, includes, port, dnsport, nslist, family) pf.enable() - def restore_firewall(self, port, family, udp): + def restore_firewall(self, port, family, udp, user): if family not in [socket.AF_INET, socket.AF_INET6]: raise Exception( 'Address family "%s" unsupported by pf method_name' diff --git a/sshuttle/methods/tproxy.py b/sshuttle/methods/tproxy.py index 44b8fd7..39cd7cd 100644 --- a/sshuttle/methods/tproxy.py +++ b/sshuttle/methods/tproxy.py @@ -150,7 +150,7 @@ class Method(BaseMethod): if udp_listener.v6 is not None: udp_listener.v6.setsockopt(SOL_IPV6, IPV6_RECVORIGDSTADDR, 1) - def setup_firewall(self, port, dnsport, nslist, family, subnets, udp): + def setup_firewall(self, port, dnsport, nslist, family, subnets, udp, user): if family not in [socket.AF_INET, socket.AF_INET6]: raise Exception( 'Address family "%s" unsupported by tproxy method' @@ -174,7 +174,7 @@ class Method(BaseMethod): divert_chain = 'sshuttle-d-%s' % port # basic cleanup/setup of chains - self.restore_firewall(port, family, udp) + self.restore_firewall(port, family, udp, user) _ipt('-N', mark_chain) _ipt('-F', mark_chain) @@ -251,7 +251,7 @@ class Method(BaseMethod): '-m', 'udp', *(udp_ports + ('--on-port', str(port)))) - def restore_firewall(self, port, family, udp): + def restore_firewall(self, port, family, udp, user): if family not in [socket.AF_INET, socket.AF_INET6]: raise Exception( 'Address family "%s" unsupported by tproxy method' diff --git a/sshuttle/options.py b/sshuttle/options.py index f4647b1..9448b02 100644 --- a/sshuttle/options.py +++ b/sshuttle/options.py @@ -286,6 +286,12 @@ parser.add_argument( pidfile name (only if using --daemon) [%(default)s] """ ) +parser.add_argument( + "--user", + help=""" + apply all the rules only to this linux user + """ +) parser.add_argument( "--firewall", action="store_true", diff --git a/sshuttle/tests/client/test_firewall.py b/sshuttle/tests/client/test_firewall.py index 6201601..bc7a8ac 100644 --- a/sshuttle/tests/client/test_firewall.py +++ b/sshuttle/tests/client/test_firewall.py @@ -15,7 +15,7 @@ NSLIST 2,1.2.3.33 10,2404:6800:4004:80c::33 PORTS 1024,1025,1026,1027 -GO 1 +GO 1 - HOST 1.2.3.3,existing """) stdout = Mock() @@ -121,14 +121,16 @@ def test_main(mock_get_method, mock_setup_daemon, mock_rewrite_etc_hosts): 10, [(10, 64, False, u'2404:6800:4004:80c::', 0, 0), (10, 128, True, u'2404:6800:4004:80c::101f', 80, 80)], - True), + True, + None), call().setup_firewall( 1025, 1027, [(2, u'1.2.3.33')], 2, [(2, 24, False, u'1.2.3.0', 8000, 9000), (2, 32, True, u'1.2.3.66', 8080, 8080)], - True), - call().restore_firewall(1024, 10, True), - call().restore_firewall(1025, 2, True), + True, + None), + call().restore_firewall(1024, 10, True, None), + call().restore_firewall(1025, 2, True, None), ] diff --git a/sshuttle/tests/client/test_methods_nat.py b/sshuttle/tests/client/test_methods_nat.py index 4ae571b..84b4f85 100644 --- a/sshuttle/tests/client/test_methods_nat.py +++ b/sshuttle/tests/client/test_methods_nat.py @@ -88,7 +88,8 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): 10, [(10, 64, False, u'2404:6800:4004:80c::', 0, 0), (10, 128, True, u'2404:6800:4004:80c::101f', 80, 80)], - True) + True, + None) assert str(excinfo.value) \ == 'Address family "AF_INET6" unsupported by nat method_name' assert mock_ipt_chain_exists.mock_calls == [] @@ -102,7 +103,8 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): 2, [(2, 24, False, u'1.2.3.0', 8000, 9000), (2, 32, True, u'1.2.3.66', 8080, 8080)], - True) + True, + None) assert str(excinfo.value) == 'UDP not supported by nat method_name' assert mock_ipt_chain_exists.mock_calls == [] assert mock_ipt_ttl.mock_calls == [] @@ -114,7 +116,8 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): 2, [(2, 24, False, u'1.2.3.0', 8000, 9000), (2, 32, True, u'1.2.3.66', 8080, 8080)], - False) + False, + None) assert mock_ipt_chain_exists.mock_calls == [ call(2, 'nat', 'sshuttle-1025') ] @@ -142,7 +145,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): mock_ipt_ttl.reset_mock() mock_ipt.reset_mock() - method.restore_firewall(1025, 2, False) + method.restore_firewall(1025, 2, False, None) assert mock_ipt_chain_exists.mock_calls == [ call(2, 'nat', 'sshuttle-1025') ] diff --git a/sshuttle/tests/client/test_methods_pf.py b/sshuttle/tests/client/test_methods_pf.py index 36b64a9..637f347 100644 --- a/sshuttle/tests/client/test_methods_pf.py +++ b/sshuttle/tests/client/test_methods_pf.py @@ -184,7 +184,8 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl): 10, [(10, 64, False, u'2404:6800:4004:80c::', 8000, 9000), (10, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)], - False) + False, + None) assert mock_ioctl.mock_calls == [ call(mock_pf_get_dev(), 0xC4704433, ANY), call(mock_pf_get_dev(), 0xCC20441A, ANY), @@ -222,7 +223,8 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl): 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - True) + True, + None) assert str(excinfo.value) == 'UDP not supported by pf method_name' assert mock_pf_get_dev.mock_calls == [] assert mock_ioctl.mock_calls == [] @@ -233,7 +235,8 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl): [(2, u'1.2.3.33')], 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - False) + False, + None) assert mock_ioctl.mock_calls == [ call(mock_pf_get_dev(), 0xC4704433, ANY), call(mock_pf_get_dev(), 0xCC20441A, ANY), @@ -262,7 +265,7 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl): mock_ioctl.reset_mock() mock_pfctl.reset_mock() - method.restore_firewall(1025, 2, False) + method.restore_firewall(1025, 2, False, None) assert mock_ioctl.mock_calls == [] assert mock_pfctl.mock_calls == [ call('-a sshuttle-1025 -F all'), @@ -290,7 +293,8 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): 10, [(10, 64, False, u'2404:6800:4004:80c::', 8000, 9000), (10, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)], - False) + False, + None) assert mock_pfctl.mock_calls == [ call('-s all'), @@ -319,7 +323,8 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - True) + True, + None) assert str(excinfo.value) == 'UDP not supported by pf method_name' assert mock_pf_get_dev.mock_calls == [] assert mock_ioctl.mock_calls == [] @@ -330,7 +335,8 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): [(2, u'1.2.3.33')], 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - False) + False, + None) assert mock_ioctl.mock_calls == [ call(mock_pf_get_dev(), 0xC4704433, ANY), call(mock_pf_get_dev(), 0xCBE0441A, ANY), @@ -357,7 +363,7 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): mock_ioctl.reset_mock() mock_pfctl.reset_mock() - method.restore_firewall(1025, 2, False) + method.restore_firewall(1025, 2, False, None) assert mock_ioctl.mock_calls == [] assert mock_pfctl.mock_calls == [ call('-a sshuttle-1025 -F all'), @@ -385,7 +391,8 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): 10, [(10, 64, False, u'2404:6800:4004:80c::', 8000, 9000), (10, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)], - False) + False, + None) assert mock_ioctl.mock_calls == [ call(mock_pf_get_dev(), 0xcd48441a, ANY), @@ -420,7 +427,8 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - True) + True, + None) assert str(excinfo.value) == 'UDP not supported by pf method_name' assert mock_pf_get_dev.mock_calls == [] assert mock_ioctl.mock_calls == [] @@ -432,7 +440,8 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - False) + False, + None) assert mock_ioctl.mock_calls == [ call(mock_pf_get_dev(), 0xcd48441a, ANY), call(mock_pf_get_dev(), 0xcd48441a, ANY), @@ -457,7 +466,7 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl): mock_ioctl.reset_mock() mock_pfctl.reset_mock() - method.restore_firewall(1025, 2, False) + method.restore_firewall(1025, 2, False, None) assert mock_ioctl.mock_calls == [] assert mock_pfctl.mock_calls == [ call('-a sshuttle-1025 -F all'), diff --git a/sshuttle/tests/client/test_methods_tproxy.py b/sshuttle/tests/client/test_methods_tproxy.py index 268e60c..89b47d1 100644 --- a/sshuttle/tests/client/test_methods_tproxy.py +++ b/sshuttle/tests/client/test_methods_tproxy.py @@ -104,7 +104,8 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): 10, [(10, 64, False, u'2404:6800:4004:80c::', 8000, 9000), (10, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)], - True) + True, + None) assert mock_ipt_chain_exists.mock_calls == [ call(10, 'mangle', 'sshuttle-m-1024'), call(10, 'mangle', 'sshuttle-t-1024'), @@ -173,7 +174,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): mock_ipt_ttl.reset_mock() mock_ipt.reset_mock() - method.restore_firewall(1025, 10, True) + method.restore_firewall(1025, 10, True, None) assert mock_ipt_chain_exists.mock_calls == [ call(10, 'mangle', 'sshuttle-m-1025'), call(10, 'mangle', 'sshuttle-t-1025'), @@ -201,7 +202,8 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): [(2, u'1.2.3.33')], 2, [(2, 24, False, u'1.2.3.0', 0, 0), (2, 32, True, u'1.2.3.66', 80, 80)], - True) + True, + None) assert mock_ipt_chain_exists.mock_calls == [ call(2, 'mangle', 'sshuttle-m-1025'), call(2, 'mangle', 'sshuttle-t-1025'), @@ -267,7 +269,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt): mock_ipt_ttl.reset_mock() mock_ipt.reset_mock() - method.restore_firewall(1025, 2, True) + method.restore_firewall(1025, 2, True, None) assert mock_ipt_chain_exists.mock_calls == [ call(2, 'mangle', 'sshuttle-m-1025'), call(2, 'mangle', 'sshuttle-t-1025'),