Added nft_chain_exists() and fixed nft to use that

This commit is contained in:
Haw Loeung 2020-02-03 10:37:02 +11:00 committed by Brian May
parent 84076f29fa
commit 13db89916a
2 changed files with 28 additions and 4 deletions

View File

@ -68,6 +68,32 @@ 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']

View File

@ -1,7 +1,7 @@
import socket
from sshuttle.firewall import subnet_weight
from sshuttle.helpers import Fatal, log
from sshuttle.linux import nft, nft_get_handle, nonfatal
from sshuttle.linux import nft, nft_get_handle, nft_chain_exists, nonfatal
from sshuttle.methods import BaseMethod
@ -28,10 +28,8 @@ class Method(BaseMethod):
for chain in ['prerouting', 'postrouting', 'output']:
rules = '{{ type nat hook {} priority -100; policy accept; }}' \
.format(chain)
try:
if not nft_chain_exists(family, table, chain):
_nft('add chain', chain, rules)
except Fatal:
log('Chain {} already exists, ignoring\n'.format(chain))
chain = 'sshuttle-%s' % port