Add more methods tests.

Fix bug in tproxy recv_udp() method.

3 files changed, 423 insertions, 1 deletions

diff --git a/sshuttle/methods/tproxy.py b/sshuttle/methods/tproxy.py
index 4c02286..5b4f408 100644
--- a/sshuttle/methods/tproxy.py
+++ b/sshuttle/methods/tproxy.py
@@ -118,9 +118,14 @@ class Method(BaseMethod):
                 "-- ignored UDP from %r: "
                 "couldn't determine destination IP address\n" % (srcip,))
             return None
-        return None
+        return srcip, dstip, data
 
     def send_udp(self, sock, srcip, dstip, data):
+        if not srcip:
+            debug1(
+                "-- ignored UDP to %r: "
+                "couldn't determine source IP address\n" % (dstip,))
+            return
         sender = socket.socket(sock.family, socket.SOCK_DGRAM)
         sender.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         sender.setsockopt(socket.SOL_IP, IP_TRANSPARENT, 1)
diff --git a/sshuttle/tests/test_methods_nat.py b/sshuttle/tests/test_methods_nat.py
new file mode 100644
index 0000000..433c709
--- /dev/null
+++ b/sshuttle/tests/test_methods_nat.py
@@ -0,0 +1,145 @@
+import pytest
+from mock import Mock, patch, call
+import socket
+import struct
+
+from sshuttle.methods import get_method
+
+
+def test_get_supported_features():
+    method = get_method('nat')
+    features = method.get_supported_features()
+    assert not features.ipv6
+    assert not features.udp
+
+
+def test_get_tcp_dstip():
+    sock = Mock()
+    sock.getsockopt.return_value = struct.pack(
+        '!HHBBBB', socket.ntohs(socket.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)]
+
+
+def test_recv_udp():
+    sock = Mock()
+    sock.recvfrom.return_value = "11111", "127.0.0.1"
+    method = get_method('nat')
+    result = method.recv_udp(sock, 1024)
+    assert sock.mock_calls == [call.recvfrom(1024)]
+    assert result == ("127.0.0.1", None, "11111")
+
+
+def test_send_udp():
+    sock = Mock()
+    method = get_method('nat')
+    method.send_udp(sock, None, "127.0.0.1", "22222")
+    assert sock.mock_calls == [call.sendto("22222", "127.0.0.1")]
+
+
+def test_setup_tcp_listener():
+    listener = Mock()
+    method = get_method('nat')
+    method.setup_tcp_listener(listener)
+    assert listener.mock_calls == []
+
+
+def test_setup_udp_listener():
+    listener = Mock()
+    method = get_method('nat')
+    method.setup_udp_listener(listener)
+    assert listener.mock_calls == []
+
+
+def test_check_settings():
+    method = get_method('nat')
+    method.check_settings(True, True)
+    method.check_settings(False, True)
+
+
+def test_firewall_command():
+    method = get_method('nat')
+    assert not method.firewall_command("somthing")
+
+
+@patch('sshuttle.methods.nat.ipt')
+@patch('sshuttle.methods.nat.ipt_ttl')
+@patch('sshuttle.methods.nat.ipt_chain_exists')
+def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt):
+    mock_ipt_chain_exists.return_value = True
+    method = get_method('nat')
+    assert method.name == 'nat'
+
+    with pytest.raises(Exception) as excinfo:
+        method.setup_firewall(
+            1024, 1026,
+            [(10, u'2404:6800:4004:80c::33')],
+            10,
+            [(10, 64, False, u'2404:6800:4004:80c::'),
+                (10, 128, True, u'2404:6800:4004:80c::101f')],
+            True)
+    assert str(excinfo.value) \
+        == 'Address family "AF_INET6" unsupported by nat method_name'
+    assert mock_ipt_chain_exists.mock_calls == []
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == []
+
+    with pytest.raises(Exception) as excinfo:
+        method.setup_firewall(
+            1025, 1027,
+            [(2, u'1.2.3.33')],
+            2,
+            [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')],
+            True)
+    assert str(excinfo.value) == 'UDP not supported by nat method_name'
+    assert mock_ipt_chain_exists.mock_calls == []
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == []
+
+    method.setup_firewall(
+        1025, 1027,
+        [(2, u'1.2.3.33')],
+        2,
+        [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')],
+        False)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(2, 'nat', 'sshuttle-1025')
+    ]
+    assert mock_ipt_ttl.mock_calls == [
+        call(2, 'nat', '-A', 'sshuttle-1025', '-j', 'REDIRECT',
+             '--dest', u'1.2.3.0/24', '-p', 'tcp', '--to-ports', '1025'),
+        call(2, 'nat', '-A', 'sshuttle-1025', '-j', 'REDIRECT',
+             '--dest', u'1.2.3.33/32', '-p', 'udp',
+             '--dport', '53', '--to-ports', '1027')
+    ]
+    assert mock_ipt.mock_calls == [
+        call(2, 'nat', '-D', 'OUTPUT', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-D', 'PREROUTING', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-F', 'sshuttle-1025'),
+        call(2, 'nat', '-X', 'sshuttle-1025'),
+        call(2, 'nat', '-N', 'sshuttle-1025'),
+        call(2, 'nat', '-F', 'sshuttle-1025'),
+        call(2, 'nat', '-I', 'OUTPUT', '1', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-I', 'PREROUTING', '1', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-A', 'sshuttle-1025', '-j', 'RETURN',
+             '--dest', u'1.2.3.66/32', '-p', 'tcp')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()
+
+    method.setup_firewall(1025, 0, [], 2, [], False)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(2, 'nat', 'sshuttle-1025')
+    ]
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == [
+        call(2, 'nat', '-D', 'OUTPUT', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-D', 'PREROUTING', '-j', 'sshuttle-1025'),
+        call(2, 'nat', '-F', 'sshuttle-1025'),
+        call(2, 'nat', '-X', 'sshuttle-1025')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()
diff --git a/sshuttle/tests/test_methods_tproxy.py b/sshuttle/tests/test_methods_tproxy.py
new file mode 100644
index 0000000..82efed5
--- /dev/null
+++ b/sshuttle/tests/test_methods_tproxy.py
@@ -0,0 +1,272 @@
+from mock import Mock, patch, call
+
+from sshuttle.methods import get_method
+
+
+def test_get_supported_features():
+    method = get_method('tproxy')
+    features = method.get_supported_features()
+    assert features.ipv6
+    assert features.udp
+
+
+def test_get_tcp_dstip():
+    sock = Mock()
+    sock.getsockname.return_value = ('127.0.0.1', 1024)
+    method = get_method('tproxy')
+    assert method.get_tcp_dstip(sock) == ('127.0.0.1', 1024)
+    assert sock.mock_calls == [call.getsockname()]
+
+
+@patch("sshuttle.methods.tproxy.recv_udp")
+def test_recv_udp(mock_recv_udp):
+    mock_recv_udp.return_value = ("127.0.0.1", "127.0.0.2", "11111")
+
+    sock = Mock()
+    method = get_method('tproxy')
+    result = method.recv_udp(sock, 1024)
+    assert sock.mock_calls == []
+    assert mock_recv_udp.mock_calls == [call(sock, 1024)]
+    assert result == ("127.0.0.1", "127.0.0.2", "11111")
+
+
+@patch("sshuttle.methods.socket.socket")
+def test_send_udp(mock_socket):
+    sock = Mock()
+    method = get_method('tproxy')
+    method.send_udp(sock, "127.0.0.2", "127.0.0.1", "2222222")
+    assert sock.mock_calls == []
+    assert mock_socket.mock_calls == [
+        call(sock.family, 2),
+        call().setsockopt(1, 2, 1),
+        call().setsockopt(0, 19, 1),
+        call().bind('127.0.0.2'),
+        call().sendto("2222222", '127.0.0.1'),
+        call().close()
+    ]
+
+
+def test_setup_tcp_listener():
+    listener = Mock()
+    method = get_method('tproxy')
+    method.setup_tcp_listener(listener)
+    assert listener.mock_calls == [
+        call.setsockopt(0, 19, 1)
+    ]
+
+
+def test_setup_udp_listener():
+    listener = Mock()
+    method = get_method('tproxy')
+    method.setup_udp_listener(listener)
+    assert listener.mock_calls == [
+        call.setsockopt(0, 19, 1),
+        call.v4.setsockopt(0, 20, 1),
+        call.v6.setsockopt(41, 74, 1)
+    ]
+
+
+def test_check_settings():
+    method = get_method('tproxy')
+    method.check_settings(True, True)
+    method.check_settings(False, True)
+
+
+def test_firewall_command():
+    method = get_method('tproxy')
+    assert not method.firewall_command("somthing")
+
+
+@patch('sshuttle.methods.tproxy.ipt')
+@patch('sshuttle.methods.tproxy.ipt_ttl')
+@patch('sshuttle.methods.tproxy.ipt_chain_exists')
+def test_setup_firewall(mock_ipt_chain_exists, mock_ipt_ttl, mock_ipt):
+    mock_ipt_chain_exists.return_value = True
+    method = get_method('tproxy')
+    assert method.name == 'tproxy'
+
+    # IPV6
+
+    method.setup_firewall(
+        1024, 1026,
+        [(10, u'2404:6800:4004:80c::33')],
+        10,
+        [(10, 64, False, u'2404:6800:4004:80c::'),
+            (10, 128, True, u'2404:6800:4004:80c::101f')],
+        True)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(10, 'mangle', 'sshuttle-m-1024'),
+        call(10, 'mangle', 'sshuttle-t-1024'),
+        call(10, 'mangle', 'sshuttle-d-1024')
+    ]
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == [
+        call(10, 'mangle', '-D', 'OUTPUT', '-j', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-X', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-D', 'PREROUTING', '-j', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-X', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-d-1024'),
+        call(10, 'mangle', '-X', 'sshuttle-d-1024'),
+        call(10, 'mangle', '-N', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-N', 'sshuttle-d-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-d-1024'),
+        call(10, 'mangle', '-N', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-F', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-I', 'OUTPUT', '1', '-j', 'sshuttle-m-1024'),
+        call(10, 'mangle', '-I', 'PREROUTING', '1', '-j', 'sshuttle-t-1024'),
+        call(10, 'mangle', '-A', 'sshuttle-d-1024', '-j', 'MARK',
+             '--set-mark', '1'),
+        call(10, 'mangle', '-A', 'sshuttle-d-1024', '-j', 'ACCEPT'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-m', 'socket',
+             '-j', 'sshuttle-d-1024', '-m', 'tcp', '-p', 'tcp'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-m', 'socket',
+             '-j', 'sshuttle-d-1024', '-m', 'udp', '-p', 'udp'),
+        call(10, 'mangle', '-A', 'sshuttle-m-1024', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'2404:6800:4004:80c::33/32',
+             '-m', 'udp', '-p', 'udp', '--dport', '53'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1',
+             '--dest', u'2404:6800:4004:80c::33/32',
+             '-m', 'udp', '-p', 'udp', '--dport', '53', '--on-port', '1026'),
+        call(10, 'mangle', '-A', 'sshuttle-m-1024', '-j', 'RETURN',
+             '--dest', u'2404:6800:4004:80c::101f/128',
+             '-m', 'tcp', '-p', 'tcp'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-j', 'RETURN',
+             '--dest', u'2404:6800:4004:80c::101f/128',
+             '-m', 'tcp', '-p', 'tcp'),
+        call(10, 'mangle', '-A', 'sshuttle-m-1024', '-j', 'RETURN',
+             '--dest', u'2404:6800:4004:80c::101f/128',
+             '-m', 'udp', '-p', 'udp'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-j', 'RETURN',
+             '--dest', u'2404:6800:4004:80c::101f/128',
+             '-m', 'udp', '-p', 'udp'),
+        call(10, 'mangle', '-A', 'sshuttle-m-1024', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'2404:6800:4004:80c::/64',
+             '-m', 'tcp', '-p', 'tcp'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1', '--dest', u'2404:6800:4004:80c::/64',
+             '-m', 'tcp', '-p', 'tcp', '--on-port', '1024'),
+        call(10, 'mangle', '-A', 'sshuttle-m-1024', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'2404:6800:4004:80c::/64',
+             '-m', 'udp', '-p', 'udp'),
+        call(10, 'mangle', '-A', 'sshuttle-t-1024', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1', '--dest', u'2404:6800:4004:80c::/64',
+             '-m', 'udp', '-p', 'udp', '--on-port', '1024')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()
+
+    method.setup_firewall(1025, 0, [], 10, [], True)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(10, 'mangle', 'sshuttle-m-1025'),
+        call(10, 'mangle', 'sshuttle-t-1025'),
+        call(10, 'mangle', 'sshuttle-d-1025')
+    ]
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == [
+        call(10, 'mangle', '-D', 'OUTPUT', '-j', 'sshuttle-m-1025'),
+        call(10, 'mangle', '-F', 'sshuttle-m-1025'),
+        call(10, 'mangle', '-X', 'sshuttle-m-1025'),
+        call(10, 'mangle', '-D', 'PREROUTING', '-j', 'sshuttle-t-1025'),
+        call(10, 'mangle', '-F', 'sshuttle-t-1025'),
+        call(10, 'mangle', '-X', 'sshuttle-t-1025'),
+        call(10, 'mangle', '-F', 'sshuttle-d-1025'),
+        call(10, 'mangle', '-X', 'sshuttle-d-1025')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()
+
+    # IPV4
+
+    method.setup_firewall(
+        1025, 1027,
+        [(2, u'1.2.3.33')],
+        2,
+        [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')],
+        True)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(2, 'mangle', 'sshuttle-m-1025'),
+        call(2, 'mangle', 'sshuttle-t-1025'),
+        call(2, 'mangle', 'sshuttle-d-1025')
+    ]
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == [
+        call(2, 'mangle', '-D', 'OUTPUT', '-j', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-D', 'PREROUTING', '-j', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-d-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-d-1025'),
+        call(2, 'mangle', '-N', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-N', 'sshuttle-d-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-d-1025'),
+        call(2, 'mangle', '-N', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-I', 'OUTPUT', '1', '-j', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-I', 'PREROUTING', '1', '-j', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-A', 'sshuttle-d-1025',
+             '-j', 'MARK', '--set-mark', '1'),
+        call(2, 'mangle', '-A', 'sshuttle-d-1025', '-j', 'ACCEPT'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-m', 'socket',
+             '-j', 'sshuttle-d-1025', '-m', 'tcp', '-p', 'tcp'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-m', 'socket',
+             '-j', 'sshuttle-d-1025', '-m', 'udp', '-p', 'udp'),
+        call(2, 'mangle', '-A', 'sshuttle-m-1025', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'1.2.3.33/32',
+             '-m', 'udp', '-p', 'udp', '--dport', '53'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1', '--dest', u'1.2.3.33/32',
+             '-m', 'udp', '-p', 'udp', '--dport', '53', '--on-port', '1027'),
+        call(2, 'mangle', '-A', 'sshuttle-m-1025', '-j', 'RETURN',
+             '--dest', u'1.2.3.66/32', '-m', 'tcp', '-p', 'tcp'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-j', 'RETURN',
+             '--dest', u'1.2.3.66/32', '-m', 'tcp', '-p', 'tcp'),
+        call(2, 'mangle', '-A', 'sshuttle-m-1025', '-j', 'RETURN',
+             '--dest', u'1.2.3.66/32', '-m', 'udp', '-p', 'udp'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-j', 'RETURN',
+             '--dest', u'1.2.3.66/32', '-m', 'udp', '-p', 'udp'),
+        call(2, 'mangle', '-A', 'sshuttle-m-1025', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'1.2.3.0/24',
+             '-m', 'tcp', '-p', 'tcp'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1', '--dest', u'1.2.3.0/24',
+             '-m', 'tcp', '-p', 'tcp', '--on-port', '1025'),
+        call(2, 'mangle', '-A', 'sshuttle-m-1025', '-j', 'MARK',
+             '--set-mark', '1', '--dest', u'1.2.3.0/24',
+             '-m', 'udp', '-p', 'udp'),
+        call(2, 'mangle', '-A', 'sshuttle-t-1025', '-j', 'TPROXY',
+             '--tproxy-mark', '0x1/0x1', '--dest', u'1.2.3.0/24',
+             '-m', 'udp', '-p', 'udp', '--on-port', '1025')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()
+
+    method.setup_firewall(1025, 0, [], 2, [], True)
+    assert mock_ipt_chain_exists.mock_calls == [
+        call(2, 'mangle', 'sshuttle-m-1025'),
+        call(2, 'mangle', 'sshuttle-t-1025'),
+        call(2, 'mangle', 'sshuttle-d-1025')
+    ]
+    assert mock_ipt_ttl.mock_calls == []
+    assert mock_ipt.mock_calls == [
+        call(2, 'mangle', '-D', 'OUTPUT', '-j', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-m-1025'),
+        call(2, 'mangle', '-D', 'PREROUTING', '-j', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-t-1025'),
+        call(2, 'mangle', '-F', 'sshuttle-d-1025'),
+        call(2, 'mangle', '-X', 'sshuttle-d-1025')
+    ]
+    mock_ipt_chain_exists.reset_mock()
+    mock_ipt_ttl.reset_mock()
+    mock_ipt.reset_mock()