Route traffic by linux user

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
@@ -287,6 +287,12 @@ parser.add_argument(
     """
 )
 parser.add_argument(
+    "--user",
+    help="""
+    apply all the rules only to this linux user
+    """
+)
+parser.add_argument(
     "--firewall",
     action="store_true",
     help="""
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'),