IPv6 support in nft method.

This works for me but needs testing by others. Remember to specify a
::0/0 subnet or similar to route IPv6 through sshuttle.

I'm adding this to nft before nat since it is not sshuttle's default
method on Linux. Documentation updates may be required too.

This patch uses the ipaddress module, but that appears to be included
since Python 3.3.

3 files changed, 64 insertions, 17 deletions

diff --git a/sshuttle/methods/__init__.py b/sshuttle/methods/__init__.py
index 7a1d493..3968c9d 100644
--- a/sshuttle/methods/__init__.py
+++ b/sshuttle/methods/__init__.py
@@ -3,24 +3,34 @@ import importlib
 import socket
 import struct
 import errno
+import ipaddress
 from sshuttle.helpers import Fatal, debug3
 
 
 def original_dst(sock):
+    ip = "0.0.0.0"
+    port = -1
     try:
+        family = sock.family
         SO_ORIGINAL_DST = 80
-        SOCKADDR_MIN = 16
-        sockaddr_in = sock.getsockopt(socket.SOL_IP,
-                                      SO_ORIGINAL_DST, SOCKADDR_MIN)
-        (proto, port, a, b, c, d) = struct.unpack('!HHBBBB', sockaddr_in[:8])
-        # FIXME: decoding is IPv4 only.
-        assert(socket.htons(proto) == socket.AF_INET)
-        ip = '%d.%d.%d.%d' % (a, b, c, d)
-        return (ip, port)
+
+        if family == socket.AF_INET:
+            SOCKADDR_MIN = 16
+            sockaddr_in = sock.getsockopt(socket.SOL_IP,
+                                          SO_ORIGINAL_DST, SOCKADDR_MIN)
+            port, raw_ip = struct.unpack_from('!2xH4s', sockaddr_in[:8])
+            ip = str(ipaddress.IPv4Address(raw_ip))
+        elif family == socket.AF_INET6:
+            sockaddr_in = sock.getsockopt(41, SO_ORIGINAL_DST, 64)
+            port, raw_ip = struct.unpack_from("!2xH4x16s", sockaddr_in)
+            ip = str(ipaddress.IPv6Address(raw_ip))
+        else:
+            raise Fatal("fw: Unknown family type.")
     except socket.error as e:
         if e.args[0] == errno.ENOPROTOOPT:
             return sock.getsockname()
         raise
+    return (ip, port)
 
 
 class Features(object):
diff --git a/sshuttle/methods/nft.py b/sshuttle/methods/nft.py
index 3ec47e0..058baa9 100644
--- a/sshuttle/methods/nft.py
+++ b/sshuttle/methods/nft.py
@@ -16,7 +16,10 @@ class Method(BaseMethod):
         if udp:
             raise Exception("UDP not supported by nft")
 
-        table = 'sshuttle-%s' % port
+        if family == socket.AF_INET:
+            table = 'sshuttle-ipv4-%s' % port
+        if family == socket.AF_INET6:
+            table = 'sshuttle-ipv6-%s' % port
 
         def _nft(action, *args):
             return nft(family, table, action, *args)
@@ -37,7 +40,10 @@ class Method(BaseMethod):
         # This TTL hack allows the client and server to run on the
         # same host. The connections the sshuttle server makes will
         # have TTL set to 63.
-        _nft('add rule', chain, 'ip ttl == 63 return')
+        if family == socket.AF_INET:
+            _nft('add rule', chain, 'ip ttl == 63 return')
+        elif family == socket.AF_INET6:
+            _nft('add rule', chain, 'ip6 hoplimit == 63 return')
 
         # Redirect DNS traffic as requested. This includes routing traffic
         # to localhost DNS servers through sshuttle.
@@ -57,7 +63,11 @@ class Method(BaseMethod):
         # create new subnet entries.
         for _, swidth, sexclude, snet, fport, lport \
                 in sorted(subnets, key=subnet_weight, reverse=True):
-            tcp_ports = ('ip', 'protocol', 'tcp')
+            if family == socket.AF_INET:
+                tcp_ports = ('ip', 'protocol', 'tcp')
+            elif family == socket.AF_INET6:
+                tcp_ports = ('ip6', 'nexthdr', 'tcp')
+
             if fport and fport != lport:
                 tcp_ports = \
                     tcp_ports + \
@@ -66,21 +76,38 @@ class Method(BaseMethod):
                 tcp_ports = tcp_ports + ('tcp', 'dport', '%d' % (fport))
 
             if sexclude:
-                _nft('add rule', chain, *(tcp_ports + (
-                     'ip daddr %s/%s' % (snet, swidth), 'return')))
+                if family == socket.AF_INET:
+                    _nft('add rule', chain, *(tcp_ports + (
+                        'ip daddr %s/%s' % (snet, swidth), 'return')))
+                elif family == socket.AF_INET6:
+                    _nft('add rule', chain, *(tcp_ports + (
+                        'ip6 daddr %s/%s' % (snet, swidth), 'return')))
             else:
-                _nft('add rule', chain, *(tcp_ports + (
-                     'ip daddr %s/%s' % (snet, swidth),
-                     ('redirect to :' + str(port)))))
+                if family == socket.AF_INET:
+                    _nft('add rule', chain, *(tcp_ports + (
+                        'ip daddr %s/%s' % (snet, swidth),
+                        ('redirect to :' + str(port)))))
+                elif family == socket.AF_INET6:
+                    _nft('add rule', chain, *(tcp_ports + (
+                        'ip6 daddr %s/%s' % (snet, swidth),
+                        ('redirect to :' + str(port)))))
 
     def restore_firewall(self, port, family, udp, user):
         if udp:
             raise Exception("UDP not supported by nft method_name")
 
-        table = 'sshuttle-%s' % port
+        if family == socket.AF_INET:
+            table = 'sshuttle-ipv4-%s' % port
+        if family == socket.AF_INET6:
+            table = 'sshuttle-ipv6-%s' % port
 
         def _nft(action, *args):
             return nft(family, table, action, *args)
 
         # basic cleanup/setup of chains
         nonfatal(_nft, 'delete table', '')
+
+    def get_supported_features(self):
+        result = super(Method, self).get_supported_features()
+        result.ipv6 = True
+        return result
diff --git a/tests/client/test_methods_nat.py b/tests/client/test_methods_nat.py
index 11a901b..a9d2a25 100644
--- a/tests/client/test_methods_nat.py
+++ b/tests/client/test_methods_nat.py
@@ -18,12 +18,22 @@ def test_get_supported_features():
 
 def test_get_tcp_dstip():
     sock = Mock()
+    sock.family = AF_INET
     sock.getsockopt.return_value = struct.pack(
         '!HHBBBB', socket.ntohs(AF_INET), 1024, 127, 0, 0, 1)
     method = get_method('nat')
     assert method.get_tcp_dstip(sock) == ('127.0.0.1', 1024)
     assert sock.mock_calls == [call.getsockopt(0, 80, 16)]
 
+    sock = Mock()
+    sock.family = AF_INET6
+    sock.getsockopt.return_value = struct.pack(
+        '!HH4xBBBBBBBBBBBBBBBB', socket.ntohs(AF_INET6),
+        1024, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1)
+    method = get_method('nft')
+    assert method.get_tcp_dstip(sock) == ('::1', 1024)
+    assert sock.mock_calls == [call.getsockopt(41, 80, 64)]
+
 
 def test_recv_udp():
     sock = Mock()