codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2# Copyright (C) 2016  Fabio Falcinelli, Maximilian Hils3#4# This program is free software: you can redistribute it and/or modify5# it under the terms of the GNU Lesser General Public License as published by6# the Free Software Foundation, either version 3 of the License, or7# (at your option) any later version.8#9# This program is distributed in the hope that it will be useful,10# but WITHOUT ANY WARRANTY; without even the implied warranty of11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the12# GNU Lesser General Public License for more details.13#14# You should have received a copy of the GNU Lesser General Public License15# along with this program.  If not, see <http://www.gnu.org/licenses/>.16import time17 18import pytest19from pydivert.consts import Param20from pydivert.windivert import WinDivert21 22from .fixtures import scenario, windivert_handle as w23 24assert scenario, w  # keep fixtures25 26 27def test_open():28    w = WinDivert("false")29    w.open()30    assert w.is_open31    w.close()32    assert not w.is_open33 34    with w:35        # open a second one.36        with WinDivert("false") as w2:37            assert w2.is_open38 39        assert w.is_open40        assert "open" in repr(w)41 42        with pytest.raises(RuntimeError):43            w.open()44 45    assert not w.is_open46    assert "closed" in repr(w)47 48    with pytest.raises(RuntimeError):49        w.recv()50    with pytest.raises(RuntimeError):51        w.close()52 53 54def test_register():55    if WinDivert.is_registered():56        WinDivert.unregister()57    while WinDivert.is_registered():58        time.sleep(0.01)  # pragma: no cover59    assert not WinDivert.is_registered()60    WinDivert.register()61    assert WinDivert.is_registered()62 63 64def test_unregister():65    w = WinDivert("false")66    w.open()67    WinDivert.unregister()68    time.sleep(0.1)69    assert WinDivert.is_registered()70    w.close()71    # may not trigger immediately.72    while WinDivert.is_registered():73        time.sleep(0.01)  # pragma: no cover74 75 76class TestParams(object):77    def test_queue_time_range(self, w):78        """79        Tests setting the minimum value for queue time.80        From docs: 128 < default 512 < 204881        """82        def_range = (128, 512, 2048)83        for value in def_range:84            w.set_param(Param.QUEUE_TIME, value)85            assert value == w.get_param(Param.QUEUE_TIME)86 87    def test_queue_len_range(self, w):88        """89        Tests setting the minimum value for queue length.90        From docs: 1< default 512 <819291        """92        for value in (1, 512, 8192):93            w.set_param(Param.QUEUE_LEN, value)94            assert value == w.get_param(Param.QUEUE_LEN)95 96    def test_invalid_set(self, w):97        with pytest.raises(Exception):98            w.set_param(42, 43)99 100    def test_invalid_get(self, w):101        with pytest.raises(Exception):102            w.get_param(42)103 104 105def test_echo(scenario):106    client_addr, server_addr, w, send = scenario107    w = w  # type: WinDivert108    reply = send(server_addr, b"echo")109 110    for p in w:111        assert p.is_loopback112        assert p.is_outbound113        w.send(p)114        done = (115            p.udp and p.dst_port == client_addr[1]116            or117            p.tcp and p.tcp.fin118        )119        if done:120            break121 122    assert reply.get() == b"ECHO"123 124 125def test_divert(scenario):126    client_addr, server_addr, w, send = scenario127    w = w  # type: WinDivert128    target = (server_addr[0], 80)129    reply = send(target, b"echo")130    for p in w:131        if p.src_port == client_addr[1]:132            p.dst_port = server_addr[1]133        if p.src_port == server_addr[1]:134            p.src_port = target[1]135        w.send(p)136 137        done = (138            p.udp and p.dst_port == client_addr[1]139            or140            p.tcp and p.tcp.fin141        )142        if done:143            break144 145    assert reply.get() == b"ECHO"146 147 148def test_modify_payload(scenario):149    client_addr, server_addr, w, send = scenario150    w = w  # type: WinDivert151    reply = send(server_addr, b"echo")152 153    for p in w:154        p.payload = p.payload.replace(b"echo", b"test").replace(b"TEST", b"ECHO")155        w.send(p)156 157        done = (158            p.udp and p.dst_port == client_addr[1]159            or160            p.tcp and p.tcp.fin161        )162        if done:163            break164    assert reply.get() == b"ECHO"165 166 167def test_packet_cutoff(scenario):168    client_addr, server_addr, w, send = scenario169    w = w  # type: WinDivert170    reply = send(server_addr, b"a" * 1000)171 172    cutoff = None173    while True:174        p = w.recv(500)175        if p.ip.packet_len != len(p.raw):176            assert cutoff is None177            cutoff = p.ip.packet_len - len(p.raw)178            p.ip.packet_len = len(p.raw)  # fix length179            if p.udp:180                p.udp.payload_len = len(p.payload)181        w.send(p)182        done = (183            p.udp and p.dst_port == client_addr[1]184            or185            p.tcp and p.tcp.fin186        )187        if done:188            break189    assert cutoff190    assert reply.get() == b"A" * (1000 - cutoff)191 192def test_check_filter():193 194    res, pos, msg = WinDivert.check_filter('true')195    assert res196    assert pos == 0197    assert msg is not None198    res, pos, msg = WinDivert.check_filter('something wrong here')199    assert not res200    assert pos == 0201    assert msg is not None202    res, pos, msg = WinDivert.check_filter('outbound and something wrong here')203    assert not res204    assert pos == 13205