Python methods.get_method函数代码示例

def test_get_tcp_dstip():
    sock = Mock()
    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)]

def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl):
    mock_pfctl.side_effect = pfctl

    method = get_method("pf")
    assert method.name == "pf"

    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 pf method_name'
    assert mock_pf_get_dev.mock_calls == []
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.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 pf method_name"
    assert mock_pf_get_dev.mock_calls == []
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.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_ioctl.mock_calls == [call(mock_pf_get_dev(), 0xCD48441A, ANY), call(mock_pf_get_dev(), 0xCD48441A, ANY)]
    assert mock_pfctl.mock_calls == [
        call("-f /dev/stdin", b"match on lo\n"),
        call("-s all"),
        call(
            "-a sshuttle -f /dev/stdin",
            b"table <forward_subnets> {!1.2.3.66/32,1.2.3.0/24}\n"
            b"table <dns_servers> {1.2.3.33}\n"
            b"pass in on lo0 inet proto tcp divert-to 127.0.0.1 port 1025\n"
            b"pass in on lo0 inet proto udp to "
            b"<dns_servers>port 53 rdr-to 127.0.0.1 port 1027\n"
            b"pass out inet proto tcp to "
            b"<forward_subnets> route-to lo0 keep state\n"
            b"pass out inet proto udp to "
            b"<dns_servers> port 53 route-to lo0 keep state\n",
        ),
        call("-e"),
    ]
    mock_pf_get_dev.reset_mock()
    mock_ioctl.reset_mock()
    mock_pfctl.reset_mock()

    method.restore_firewall(1025, 2, False)
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.mock_calls == [call("-a sshuttle -F all"), call("-d")]
    mock_pf_get_dev.reset_mock()
    mock_pfctl.reset_mock()
    mock_ioctl.reset_mock()

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_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_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")

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_setup_firewall(mock_pf_get_dev, mock_ioctl, mock_pfctl):
    method = get_method('pf')
    assert method.name == 'pf'

    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 pf method_name'
    assert mock_pf_get_dev.mock_calls == []
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.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 pf method_name'
    assert mock_pf_get_dev.mock_calls == []
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.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_ioctl.mock_calls == [
        call(mock_pf_get_dev(), 3295691827, ANY),
        call(mock_pf_get_dev(), 3424666650, ANY),
        call(mock_pf_get_dev(), 3424666650, ANY),
        call(mock_pf_get_dev(), 3295691827, ANY),
        call(mock_pf_get_dev(), 3424666650, ANY),
        call(mock_pf_get_dev(), 3424666650, ANY),
    ]
    # FIXME - needs more work
    # print(mock_pfctl.mock_calls)
    # assert mock_pfctl.mock_calls == []
    mock_pf_get_dev.reset_mock()
    mock_ioctl.reset_mock()
    mock_pfctl.reset_mock()

    method.setup_firewall(1025, 0, [], 2, [], False)
    assert mock_ioctl.mock_calls == []
    assert mock_pfctl.mock_calls == [call('-a sshuttle -F all')]
    mock_pf_get_dev.reset_mock()
    mock_pfctl.reset_mock()
    mock_ioctl.reset_mock()

    def __init__(self, method_name):
        self.auto_nets = []
        python_path = os.path.dirname(os.path.dirname(__file__))
        argvbase = ([sys.executable, sys.argv[0]] +
                    ['-v'] * (helpers.verbose or 0) +
                    ['--method', method_name] +
                    ['--firewall'])
        if ssyslog._p:
            argvbase += ['--syslog']
        argv_tries = [
            ['sudo', '-p', '[local sudo] Password: ',
                ('PYTHONPATH=%s' % python_path), '--'] + argvbase,
            argvbase
        ]

        # we can't use stdin/stdout=subprocess.PIPE here, as we normally would,
        # because stupid Linux 'su' requires that stdin be attached to a tty.
        # Instead, attach a *bidirectional* socket to its stdout, and use
        # that for talking in both directions.
        (s1, s2) = socket.socketpair()

        def setup():
            # run in the child process
            s2.close()
        e = None
        if os.getuid() == 0:
            argv_tries = argv_tries[-1:]  # last entry only
        for argv in argv_tries:
            try:
                if argv[0] == 'su':
                    sys.stderr.write('[local su] ')
                self.p = ssubprocess.Popen(argv, stdout=s1, preexec_fn=setup)
                e = None
                break
            except OSError as e:
                pass
        self.argv = argv
        s1.close()
        if sys.version_info < (3, 0):
            # python 2.7
            self.pfile = s2.makefile('wb+')
        else:
            # python 3.5
            self.pfile = s2.makefile('rwb')
        if e:
            log('Spawning firewall manager: %r\n' % self.argv)
            raise Fatal(e)
        line = self.pfile.readline()
        self.check()
        if line[0:5] != b'READY':
            raise Fatal('%r expected READY, got %r' % (self.argv, line))
        method_name = line[6:-1]
        self.method = get_method(method_name.decode("ASCII"))
        self.method.set_firewall(self)

def test_assert_features():
    method = get_method("nat")
    features = method.get_supported_features()
    method.assert_features(features)

    features.udp = True
    with pytest.raises(Fatal):
        method.assert_features(features)

    features.ipv6 = True
    with pytest.raises(Fatal):
        method.assert_features(features)

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_get_tcp_dstip():
    sock = Mock()
    sock.getpeername.return_value = ("127.0.0.1", 1024)
    sock.getsockname.return_value = ("127.0.0.2", 1025)
    sock.family = socket.AF_INET

    firewall = Mock()
    firewall.pfile.readline.return_value = b"QUERY_PF_NAT_SUCCESS 127.0.0.3,1026\n"

    method = get_method("pf")
    method.set_firewall(firewall)
    assert method.get_tcp_dstip(sock) == ("127.0.0.3", 1026)

    assert sock.mock_calls == [call.getpeername(), call.getsockname()]
    assert firewall.mock_calls == [
        call.pfile.write(b"QUERY_PF_NAT 2,6,127.0.0.1,1024,127.0.0.2,1025\n"),
        call.pfile.flush(),
        call.pfile.readline(),
    ]

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,
            [(AF_INET6, u'2404:6800:4004:80c::33')],
            AF_INET6,
            [(AF_INET6, 64, False, u'2404:6800:4004:80c::', 0, 0),
                (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 80, 80)],
            True,
            None)
    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,
            [(AF_INET, u'1.2.3.33')],
            AF_INET,
            [(AF_INET, 24, False, u'1.2.3.0', 8000, 9000),
                (AF_INET, 32, True, u'1.2.3.66', 8080, 8080)],
            True,
            None)
    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,
        [(AF_INET, u'1.2.3.33')],
        AF_INET,
        [(AF_INET, 24, False, u'1.2.3.0', 8000, 9000),
            (AF_INET, 32, True, u'1.2.3.66', 8080, 8080)],
        False,
        None)
    assert mock_ipt_chain_exists.mock_calls == [
        call(AF_INET, 'nat', 'sshuttle-1025')
    ]
    assert mock_ipt_ttl.mock_calls == [
        call(AF_INET, 'nat', '-A', 'sshuttle-1025', '-j', 'REDIRECT',
             '--dest', u'1.2.3.0/24', '-p', 'tcp', '--dport', '8000:9000',
             '--to-ports', '1025'),
        call(AF_INET, '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(AF_INET, 'nat', '-D', 'OUTPUT', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-D', 'PREROUTING', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-F', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-X', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-N', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-F', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-I', 'OUTPUT', '1', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-I', 'PREROUTING', '1', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-A', 'sshuttle-1025', '-j', 'RETURN',
             '--dest', u'1.2.3.66/32', '-p', 'tcp', '--dport', '8080:8080')
    ]
    mock_ipt_chain_exists.reset_mock()
    mock_ipt_ttl.reset_mock()
    mock_ipt.reset_mock()

    method.restore_firewall(1025, AF_INET, False, None)
    assert mock_ipt_chain_exists.mock_calls == [
        call(AF_INET, 'nat', 'sshuttle-1025')
    ]
    assert mock_ipt_ttl.mock_calls == []
    assert mock_ipt.mock_calls == [
        call(AF_INET, 'nat', '-D', 'OUTPUT', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-D', 'PREROUTING', '-j', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-F', 'sshuttle-1025'),
        call(AF_INET, 'nat', '-X', 'sshuttle-1025')
    ]
    mock_ipt_chain_exists.reset_mock()
    mock_ipt_ttl.reset_mock()
    mock_ipt.reset_mock()

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_get_supported_features():
    method = get_method('tproxy')
    features = method.get_supported_features()
    assert features.ipv6
    assert features.udp

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.restore_firewall(1025, 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()