Split setup_firewall method.

* setup_firewall sets the firewall up.
* restore_firewall restores the firewall to initial state.

sshuttle/firewall.py

@@ -232,7 +232,7 @@ def main(method_name, syslog):
         try:
             if port_v6:
                 debug2('firewall manager: undoing IPv6 changes.\n')
-                method.setup_firewall(port_v6, 0, [], socket.AF_INET6, [], udp)
+                method.restore_firewall(port_v6, socket.AF_INET6, udp)
         except:
             try:
                 debug1("firewall manager: "
@@ -245,7 +245,7 @@ def main(method_name, syslog):
         try:
             if port_v4:
                 debug2('firewall manager: undoing IPv4 changes.\n')
-            method.setup_firewall(port_v4, 0, [], socket.AF_INET, [], udp)
+            method.restore_firewall(port_v4, socket.AF_INET, udp)
         except:
             try:
                 debug1("firewall manager: "

sshuttle/methods/__init__.py

@@ -68,6 +68,9 @@ class BaseMethod(object):
     def setup_firewall(self, port, dnsport, nslist, family, subnets, udp):
         raise NotImplementedError()
 
+    def restore_firewall(self, port, family, udp):
+        raise NotImplementedError()
+
     def firewall_command(self, line):
         return False
 

sshuttle/methods/nat.py

@@ -31,19 +31,13 @@ class Method(BaseMethod):
         chain = 'sshuttle-%s' % port
 
         # basic cleanup/setup of chains
-        if ipt_chain_exists(family, table, chain):
+        self.restore_firewall(port, family, udp)
-            nonfatal(_ipt, '-D', 'OUTPUT', '-j', chain)
-            nonfatal(_ipt, '-D', 'PREROUTING', '-j', chain)
-            nonfatal(_ipt, '-F', chain)
-            _ipt('-X', chain)
 
-        if subnets or dnsport:
         _ipt('-N', chain)
         _ipt('-F', chain)
         _ipt('-I', 'OUTPUT', '1', '-j', chain)
         _ipt('-I', 'PREROUTING', '1', '-j', chain)
 
-        if subnets:
         # create new subnet entries.  Note that we're sorting in a very
         # particular order: we need to go from most-specific (largest
         # swidth) to least-specific, and at any given level of specificity,
@@ -68,3 +62,29 @@ class Method(BaseMethod):
                          '-p', 'udp',
                          '--dport', '53',
                          '--to-ports', str(dnsport))
+
+    def restore_firewall(self, port, family, udp):
+        # only ipv4 supported with NAT
+        if family != socket.AF_INET:
+            raise Exception(
+                'Address family "%s" unsupported by nat method_name'
+                % family_to_string(family))
+        if udp:
+            raise Exception("UDP not supported by nat method_name")
+
+        table = "nat"
+
+        def _ipt(*args):
+            return ipt(family, table, *args)
+
+        def _ipt_ttl(*args):
+            return ipt_ttl(family, table, *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)
+            nonfatal(_ipt, '-F', chain)
+            _ipt('-X', chain)

sshuttle/methods/pf.py

@@ -181,7 +181,6 @@ class Method(BaseMethod):
         if udp:
             raise Exception("UDP not supported by pf method_name")
 
-        if subnets:
         includes = []
         # If a given subnet is both included and excluded, list the
         # exclusion first; the table will ignore the second, opposite
@@ -231,7 +230,15 @@ class Method(BaseMethod):
         elif b'INFO:\nStatus: Disabled' in pf_status:
             pfctl('-e')
             _pf_context['started_by_sshuttle'] = True
-        else:
+
+    def restore_firewall(self, port, family, udp):
+        if family != socket.AF_INET:
+            raise Exception(
+                'Address family "%s" unsupported by pf method_name'
+                % family_to_string(family))
+        if udp:
+            raise Exception("UDP not supported by pf method_name")
+
         pfctl('-a sshuttle -F all')
         if sys.platform == "darwin":
             if _pf_context['Xtoken'] is not None:

sshuttle/methods/tproxy.py

@@ -163,21 +163,8 @@ class Method(BaseMethod):
         divert_chain = 'sshuttle-d-%s' % port
 
         # basic cleanup/setup of chains
-        if ipt_chain_exists(family, table, mark_chain):
+        self.restore_firewall(port, family, udp)
-            _ipt('-D', 'OUTPUT', '-j', mark_chain)
-            _ipt('-F', mark_chain)
-            _ipt('-X', mark_chain)
 
-        if ipt_chain_exists(family, table, tproxy_chain):
-            _ipt('-D', 'PREROUTING', '-j', tproxy_chain)
-            _ipt('-F', tproxy_chain)
-            _ipt('-X', tproxy_chain)
-
-        if ipt_chain_exists(family, table, divert_chain):
-            _ipt('-F', divert_chain)
-            _ipt('-X', divert_chain)
-
-        if subnets or dnsport:
         _ipt('-N', mark_chain)
         _ipt('-F', mark_chain)
         _ipt('-N', divert_chain)
@@ -190,7 +177,8 @@ class Method(BaseMethod):
         _ipt('-A', divert_chain, '-j', 'ACCEPT')
         _ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain,
              '-m', 'tcp', '-p', 'tcp')
-        if subnets and udp:
+
+        if udp:
             _ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain,
                  '-m', 'udp', '-p', 'udp')
 
@@ -205,7 +193,6 @@ class Method(BaseMethod):
                      '-m', 'udp', '-p', 'udp', '--dport', '53',
                      '--on-port', str(dnsport))
 
-        if subnets:
         for f, swidth, sexclude, snet \
                 in sorted(subnets, key=lambda s: s[1], reverse=True):
             if sexclude:
@@ -225,14 +212,15 @@ class Method(BaseMethod):
                      '-m', 'tcp', '-p', 'tcp',
                      '--on-port', str(port))
 
-                if sexclude and udp:
+            if udp:
+                if sexclude:
                     _ipt('-A', mark_chain, '-j', 'RETURN',
                          '--dest', '%s/%s' % (snet, swidth),
                          '-m', 'udp', '-p', 'udp')
                     _ipt('-A', tproxy_chain, '-j', 'RETURN',
                          '--dest', '%s/%s' % (snet, swidth),
                          '-m', 'udp', '-p', 'udp')
-                elif udp:
+                else:
                     _ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1',
                          '--dest', '%s/%s' % (snet, swidth),
                          '-m', 'udp', '-p', 'udp')
@@ -242,6 +230,39 @@ class Method(BaseMethod):
                          '-m', 'udp', '-p', 'udp',
                          '--on-port', str(port))
 
+    def restore_firewall(self, port, family, udp):
+        if family not in [socket.AF_INET, socket.AF_INET6]:
+            raise Exception(
+                'Address family "%s" unsupported by tproxy method'
+                % family_to_string(family))
+
+        table = "mangle"
+
+        def _ipt(*args):
+            return ipt(family, table, *args)
+
+        def _ipt_ttl(*args):
+            return ipt_ttl(family, table, *args)
+
+        mark_chain = 'sshuttle-m-%s' % port
+        tproxy_chain = 'sshuttle-t-%s' % port
+        divert_chain = 'sshuttle-d-%s' % port
+
+        # basic cleanup/setup of chains
+        if ipt_chain_exists(family, table, mark_chain):
+            _ipt('-D', 'OUTPUT', '-j', mark_chain)
+            _ipt('-F', mark_chain)
+            _ipt('-X', mark_chain)
+
+        if ipt_chain_exists(family, table, tproxy_chain):
+            _ipt('-D', 'PREROUTING', '-j', tproxy_chain)
+            _ipt('-F', tproxy_chain)
+            _ipt('-X', tproxy_chain)
+
+        if ipt_chain_exists(family, table, divert_chain):
+            _ipt('-F', divert_chain)
+            _ipt('-X', divert_chain)
+
     def check_settings(self, udp, dns):
         if udp and recvmsg is None:
             Fatal("tproxy UDP support requires recvmsg function.\n")

sshuttle/tests/test_firewall.py

@@ -97,6 +97,6 @@ def test_main(mock_get_method, mock_setup_daemon, mock_rewrite_etc_hosts):
             2,
             [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')],
             True),
-        call().setup_firewall(1024, 0, [], 10, [], True),
+        call().restore_firewall(1024, 10, True),
-        call().setup_firewall(1025, 0, [], 2, [], True),
+        call().restore_firewall(1025, 2, True),
     ]

sshuttle/tests/test_methods_nat.py

@@ -129,7 +129,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt):
     mock_ipt_ttl.reset_mock()
     mock_ipt.reset_mock()
 
-    method.setup_firewall(1025, 0, [], 2, [], False)
+    method.restore_firewall(1025, 2, False)
     assert mock_ipt_chain_exists.mock_calls == [
         call(2, 'nat', 'sshuttle-1025')
     ]

sshuttle/tests/test_methods_pf.py

@@ -173,7 +173,7 @@ def test_setup_firewall(mock_pf_get_dev, mock_ioctl, mock_pfctl):
     mock_ioctl.reset_mock()
     mock_pfctl.reset_mock()
 
-    method.setup_firewall(1025, 0, [], 2, [], False)
+    method.restore_firewall(1025, 2, False)
     assert mock_ioctl.mock_calls == []
     assert mock_pfctl.mock_calls == [
         call('-a sshuttle -F all'),

sshuttle/tests/test_methods_tproxy.py

@@ -160,7 +160,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt):
     mock_ipt_ttl.reset_mock()
     mock_ipt.reset_mock()
 
-    method.setup_firewall(1025, 0, [], 10, [], True)
+    method.restore_firewall(1025, 10, True)
     assert mock_ipt_chain_exists.mock_calls == [
         call(10, 'mangle', 'sshuttle-m-1025'),
         call(10, 'mangle', 'sshuttle-t-1025'),
@@ -250,7 +250,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt):
     mock_ipt_ttl.reset_mock()
     mock_ipt.reset_mock()
 
-    method.setup_firewall(1025, 0, [], 2, [], True)
+    method.restore_firewall(1025, 2, True)
     assert mock_ipt_chain_exists.mock_calls == [
         call(2, 'mangle', 'sshuttle-m-1025'),
         call(2, 'mangle', 'sshuttle-t-1025'),