diff options
| -rw-r--r-- | mitmproxy/addons/clientplayback.py | 21 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 16 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_clientplayback.py | 75 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_script.py | 12 | ||||
| -rw-r--r-- | test/mitmproxy/test_flowfilter.py | 1 |
5 files changed, 80 insertions, 45 deletions
diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index e69cd27a..34c6c9c9 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -1,21 +1,24 @@ from mitmproxy import exceptions from mitmproxy import ctx from mitmproxy import io +from mitmproxy import flow + +import typing class ClientPlayback: def __init__(self): self.flows = None - self.current = None - self.keepserving = None + self.current_thread = None + self.keepserving = False self.has_replayed = False - def count(self): + def count(self) -> int: if self.flows: return len(self.flows) return 0 - def load(self, flows): + def load(self, flows: typing.Sequence[flow.Flow]): self.flows = flows def configure(self, options, updated): @@ -32,11 +35,11 @@ class ClientPlayback: self.keepserving = options.keepserving def tick(self): - if self.current and not self.current.is_alive(): - self.current = None - if self.flows and not self.current: - self.current = ctx.master.replay_request(self.flows.pop(0)) + if self.current_thread and not self.current_thread.is_alive(): + self.current_thread = None + if self.flows and not self.current_thread: + self.current_thread = ctx.master.replay_request(self.flows.pop(0)) self.has_replayed = True if self.has_replayed: - if not self.flows and not self.current and not self.keepserving: + if not self.flows and not self.current_thread and not self.keepserving: ctx.master.shutdown() diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 18395be3..b9c4935c 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -7,7 +7,7 @@ from mitmproxy import stateobject from mitmproxy import connections from mitmproxy import version -from typing import Optional, Dict # noqa +import typing # noqa class Error(stateobject.StateObject): @@ -26,7 +26,7 @@ class Error(stateobject.StateObject): timestamp: Seconds since the epoch """ - def __init__(self, msg, timestamp=None): + def __init__(self, msg: str, timestamp=None) -> None: """ @type msg: str @type timestamp: float @@ -70,20 +70,20 @@ class Flow(stateobject.StateObject): type: str, client_conn: connections.ClientConnection, server_conn: connections.ServerConnection, - live=None - ): + live: bool=None + ) -> None: self.type = type self.id = str(uuid.uuid4()) self.client_conn = client_conn self.server_conn = server_conn self.live = live - self.error = None # type: Optional[Error] + self.error = None # type: typing.Optional[Error] self.intercepted = False # type: bool - self._backup = None # type: Optional[Flow] - self.reply = None # type: Optional[controller.Reply] + self._backup = None # type: typing.Optional[Flow] + self.reply = None # type: typing.Optional[controller.Reply] self.marked = False # type: bool - self.metadata = dict() # type: Dict[str, str] + self.metadata = dict() # type: typing.Dict[str, str] _stateobject_attributes = dict( id=str, diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py index 03fcfe47..32eca536 100644 --- a/test/mitmproxy/addons/test_clientplayback.py +++ b/test/mitmproxy/addons/test_clientplayback.py @@ -1,38 +1,65 @@ -from mitmproxy.test import tflow +import os import mock +from mitmproxy.test import tflow +from mitmproxy.test import tutils +from mitmproxy import io +from mitmproxy import exceptions + from mitmproxy.addons import clientplayback -from mitmproxy import options +from mitmproxy.test import taddons + + +def tdump(path, flows): + w = io.FlowWriter(open(path, "wb")) + for i in flows: + w.add(i) -from .. import mastertest + +class MockThread(): + def is_alive(self): + return False class TestClientPlayback: def test_playback(self): cp = clientplayback.ClientPlayback() - cp.configure(options.Options(), []) - assert cp.count() == 0 - f = tflow.tflow(resp=True) - cp.load([f]) - assert cp.count() == 1 - RP = "mitmproxy.proxy.protocol.http_replay.RequestReplayThread" - with mock.patch(RP) as rp: - assert not cp.current - with mastertest.mockctx(): + with taddons.context(): + assert cp.count() == 0 + f = tflow.tflow(resp=True) + cp.load([f]) + assert cp.count() == 1 + RP = "mitmproxy.proxy.protocol.http_replay.RequestReplayThread" + with mock.patch(RP) as rp: + assert not cp.current_thread + cp.tick() + rp.assert_called() + assert cp.current_thread + + cp.keepserving = False + cp.flows = None + cp.current_thread = None + with mock.patch("mitmproxy.master.Master.shutdown") as sd: cp.tick() - rp.assert_called() - assert cp.current - - cp.keepserving = False - cp.flows = None - cp.current = None - with mock.patch("mitmproxy.master.Master.shutdown") as sd: - with mastertest.mockctx(): + sd.assert_called() + + cp.current_thread = MockThread() + with mock.patch("mitmproxy.master.Master.shutdown") as sd: cp.tick() - sd.assert_called() + assert cp.current_thread is None def test_configure(self): cp = clientplayback.ClientPlayback() - cp.configure( - options.Options(), [] - ) + with taddons.context() as tctx: + with tutils.tmpdir() as td: + path = os.path.join(td, "flows") + tdump(path, [tflow.tflow()]) + tctx.configure(cp, client_replay=[path]) + tctx.configure(cp, client_replay=[]) + tctx.configure(cp) + tutils.raises( + exceptions.OptionsError, + tctx.configure, + cp, + client_replay=["nonexistent"] + ) diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index c8d7318f..c72dac40 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -34,9 +34,15 @@ class TestParseCommand: def test_parse_args(self): with tutils.chdir(tutils.test_data.dirname): - assert script.parse_command("mitmproxy/data/addonscripts/recorder.py") == ("mitmproxy/data/addonscripts/recorder.py", []) - assert script.parse_command("mitmproxy/data/addonscripts/recorder.py foo bar") == ("mitmproxy/data/addonscripts/recorder.py", ["foo", "bar"]) - assert script.parse_command("mitmproxy/data/addonscripts/recorder.py 'foo bar'") == ("mitmproxy/data/addonscripts/recorder.py", ["foo bar"]) + assert script.parse_command( + "mitmproxy/data/addonscripts/recorder.py" + ) == ("mitmproxy/data/addonscripts/recorder.py", []) + assert script.parse_command( + "mitmproxy/data/addonscripts/recorder.py foo bar" + ) == ("mitmproxy/data/addonscripts/recorder.py", ["foo", "bar"]) + assert script.parse_command( + "mitmproxy/data/addonscripts/recorder.py 'foo bar'" + ) == ("mitmproxy/data/addonscripts/recorder.py", ["foo bar"]) @ttutils.skip_not_windows def test_parse_windows(self): diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py index a1b77c05..2d409994 100644 --- a/test/mitmproxy/test_flowfilter.py +++ b/test/mitmproxy/test_flowfilter.py @@ -3,7 +3,6 @@ from mitmproxy.test import tflow from mock import patch from mitmproxy import flowfilter -from mitmproxy.test import tutils from . import tutils as ttutils |