From 3edeb726b83335f88de2b65cacd88e3a30d972f2 Mon Sep 17 00:00:00 2001
From: Julian Wollrath <jwollrath@web.de>
Date: Fri, 10 Jan 2020 14:09:54 +0100
Subject: [PATCH] Simplify nftables based method

---
 docs/manpage.rst        |  2 +-
 sshuttle/linux.py       | 50 ++---------------------------------------
 sshuttle/methods/nft.py | 29 +++++++++---------------
 3 files changed, 13 insertions(+), 68 deletions(-)

diff --git a/docs/manpage.rst b/docs/manpage.rst
index 34516ed..02a669e 100644
--- a/docs/manpage.rst
+++ b/docs/manpage.rst
@@ -44,7 +44,7 @@ Options
     to during startup will be routed over the VPN. Valid examples are
     example.com, example.com:8000 and example.com:8000-9000.
 
-.. option:: --method [auto|nat|tproxy|pf]
+.. option:: --method [auto|nat|nft|tproxy|pf]
 
    Which firewall method should sshuttle use? For auto, sshuttle attempts to
    guess the appropriate method depending on what it can find in PATH. The
diff --git a/sshuttle/linux.py b/sshuttle/linux.py
index 76ae7e0..0388a3c 100644
--- a/sshuttle/linux.py
+++ b/sshuttle/linux.py
@@ -1,8 +1,6 @@
-import re
 import os
 import socket
 import subprocess as ssubprocess
-
 from sshuttle.helpers import log, debug1, Fatal, family_to_string
 
 
@@ -52,10 +50,8 @@ def ipt(family, table, *args):
 
 
 def nft(family, table, action, *args):
-    if family == socket.AF_INET:
-        argv = ['nft', action, 'ip', table] + list(args)
-    elif family == socket.AF_INET6:
-        argv = ['nft', action, 'ip6', table] + list(args)
+    if family in (socket.AF_INET, socket.AF_INET6):
+        argv = ['nft', action, 'inet', table] + list(args)
     else:
         raise Exception('Unsupported family "%s"' % family_to_string(family))
     debug1('>> %s\n' % ' '.join(argv))
@@ -68,48 +64,6 @@ def nft(family, table, action, *args):
         raise Fatal('%r returned %d' % (argv, rv))
 
 
-def nft_chain_exists(family, table, name):
-    if family == socket.AF_INET:
-        fam = 'ip'
-    elif family == socket.AF_INET6:
-        fam = 'ip6'
-    else:
-        raise Exception('Unsupported family "%s"' % family_to_string(family))
-    argv = ['nft', 'list', 'chain', fam, table, name]
-    debug1('>> %s\n' % ' '.join(argv))
-    env = {
-        'PATH': os.environ['PATH'],
-        'LC_ALL': "C",
-    }
-    try:
-        table_exists = False
-        output = ssubprocess.check_output(argv, env=env,
-                                          stderr=ssubprocess.STDOUT)
-        for line in output.decode('ASCII').split('\n'):
-            if line.startswith('table %s %s ' % (fam, table)):
-                table_exists = True
-            if table_exists and ('chain %s {' % name) in line:
-                return True
-    except ssubprocess.CalledProcessError:
-        return False
-
-
-def nft_get_handle(expression, chain):
-    cmd = 'nft'
-    argv = [cmd, 'list', expression, '-a']
-    env = {
-        'PATH': os.environ['PATH'],
-        'LC_ALL': "C",
-    }
-    try:
-        output = ssubprocess.check_output(argv, env=env)
-        for line in output.decode('utf-8').split('\n'):
-            if ('jump %s' % chain) in line:
-                return re.sub('.*# ', '', line)
-    except ssubprocess.CalledProcessError as e:
-        raise Fatal('%r returned %d' % (argv, e.returncode))
-
-
 _no_ttl_module = False
 
 
diff --git a/sshuttle/methods/nft.py b/sshuttle/methods/nft.py
index dbc0501..0d58abb 100644
--- a/sshuttle/methods/nft.py
+++ b/sshuttle/methods/nft.py
@@ -1,6 +1,6 @@
 import socket
 from sshuttle.firewall import subnet_weight
-from sshuttle.linux import nft, nft_get_handle, nft_chain_exists, nonfatal
+from sshuttle.linux import nft, nonfatal
 from sshuttle.methods import BaseMethod
 
 
@@ -16,22 +16,19 @@ class Method(BaseMethod):
         if udp:
             raise Exception("UDP not supported by nft")
 
-        table = "nat"
+        table = 'sshuttle-%s' % port
 
         def _nft(action, *args):
             return nft(family, table, action, *args)
 
+        chain = table
+
         # basic cleanup/setup of chains
         _nft('add table', '')
-        # prerouting, postrouting, and output chains may already exist
-        for chain in ['prerouting', 'postrouting', 'output']:
-            rules = '{{ type nat hook {} priority -100; policy accept; }}' \
-                    .format(chain)
-            if not nft_chain_exists(family, table, chain):
-                _nft('add chain', chain, rules)
-
-        chain = 'sshuttle-%s' % port
-
+        _nft('add chain', 'prerouting',
+             '{ type nat hook prerouting priority -100; policy accept; }')
+        _nft('add chain', 'output',
+             '{ type nat hook output priority -100; policy accept; }')
         _nft('add chain', chain)
         _nft('flush chain', chain)
         _nft('add rule', 'output jump %s' % chain)
@@ -70,16 +67,10 @@ class Method(BaseMethod):
         if udp:
             raise Exception("UDP not supported by nft method_name")
 
-        table = "nat"
+        table = 'sshuttle-%s' % port
 
         def _nft(action, *args):
             return nft(family, table, action, *args)
 
-        chain = 'sshuttle-%s' % port
-
         # basic cleanup/setup of chains
-        handle = nft_get_handle('chain ip nat output', chain)
-        nonfatal(_nft, 'delete rule', 'output', handle)
-        handle = nft_get_handle('chain ip nat prerouting', chain)
-        nonfatal(_nft, 'delete rule', 'prerouting', handle)
-        nonfatal(_nft, 'delete chain', chain)
+        nonfatal(_nft, 'delete table', '')