diff options
| -rw-r--r-- | mitmproxy/addons/state.py | 272 | ||||
| -rw-r--r-- | mitmproxy/addons/view.py | 32 | ||||
| -rw-r--r-- | mitmproxy/tools/web/app.py | 35 | ||||
| -rw-r--r-- | mitmproxy/tools/web/master.py | 93 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_state.py | 19 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 152 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 24 | ||||
| -rw-r--r-- | test/mitmproxy/test_web_master.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 20 | ||||
| -rw-r--r-- | tox.ini | 2 |
10 files changed, 109 insertions, 542 deletions
diff --git a/mitmproxy/addons/state.py b/mitmproxy/addons/state.py deleted file mode 100644 index b6f575d1..00000000 --- a/mitmproxy/addons/state.py +++ /dev/null @@ -1,272 +0,0 @@ -from abc import abstractmethod, ABCMeta - -from typing import List # noqa - -from mitmproxy import flowfilter -from mitmproxy import flow # noqa - - -class FlowList(metaclass=ABCMeta): - def __init__(self): - self._list = [] # type: List[flow.Flow] - - def __iter__(self): - return iter(self._list) - - def __contains__(self, item): - return item in self._list - - def __getitem__(self, item): - return self._list[item] - - def __bool__(self): - return bool(self._list) - - def __len__(self): - return len(self._list) - - def index(self, f): - return self._list.index(f) - - @abstractmethod - def _add(self, f): - return - - @abstractmethod - def _update(self, f): - return - - @abstractmethod - def _remove(self, f): - return - - -def _pos(*args): - return True - - -class FlowView(FlowList): - def __init__(self, store, flt=None): - super().__init__() - if not flt: - flt = _pos - self._build(store, flt) - - self.store = store - self.store.views.append(self) - - def _close(self): - self.store.views.remove(self) - - def _build(self, flows, flt=None): - if flt: - self.filter = flt - self._list = list(filter(self.filter, flows)) - - def _add(self, f): - if self.filter(f): - self._list.append(f) - - def _update(self, f): - if f not in self._list: - self._add(f) - elif not self.filter(f): - self._remove(f) - - def _remove(self, f): - if f in self._list: - self._list.remove(f) - - def _recalculate(self, flows): - self._build(flows) - - -class FlowStore(FlowList): - """ - Responsible for handling flows in the state: - Keeps a list of all flows and provides views on them. - """ - - def __init__(self): - super().__init__() - self._set = set() # Used for O(1) lookups - self.views = [] - self._recalculate_views() - - def get(self, flow_id): - for f in self._list: - if f.id == flow_id: - return f - - def __contains__(self, f): - return f in self._set - - def _add(self, f): - """ - Adds a flow to the state. - The flow to add must not be present in the state. - """ - self._list.append(f) - self._set.add(f) - for view in self.views: - view._add(f) - - def _update(self, f): - """ - Notifies the state that a flow has been updated. - The flow must be present in the state. - """ - if f in self: - for view in self.views: - view._update(f) - - def _remove(self, f): - """ - Deletes a flow from the state. - The flow must be present in the state. - """ - self._list.remove(f) - self._set.remove(f) - for view in self.views: - view._remove(f) - - # Expensive bulk operations - - def _extend(self, flows): - """ - Adds a list of flows to the state. - The list of flows to add must not contain flows that are already in the state. - """ - self._list.extend(flows) - self._set.update(flows) - self._recalculate_views() - - def _clear(self): - self._list = [] - self._set = set() - self._recalculate_views() - - def _recalculate_views(self): - """ - Expensive operation: Recalculate all the views after a bulk change. - """ - for view in self.views: - view._recalculate(self) - - # Utility functions. - # There are some common cases where we need to argue about all flows - # irrespective of filters on the view etc (i.e. on shutdown). - - def active_count(self): - c = 0 - for i in self._list: - if not i.response and not i.error: - c += 1 - return c - - # TODO: Should accept_all operate on views or on all flows? - def accept_all(self, master): - for f in self._list: - f.resume(master) - - def kill_all(self, master): - for f in self._list: - if f.killable: - f.kill(master) - - -class State: - def __init__(self): - self.flows = FlowStore() - self.view = FlowView(self.flows, None) - - @property - def filter_txt(self): - return getattr(self.view.filter, "pattern", None) - - def flow_count(self): - return len(self.flows) - - # TODO: All functions regarding flows that don't cause side-effects should - # be moved into FlowStore. - def index(self, f): - return self.flows.index(f) - - def active_flow_count(self): - return self.flows.active_count() - - def add_flow(self, f): - """ - Add a request to the state. - """ - self.flows._add(f) - return f - - def update_flow(self, f): - """ - Add a response to the state. - """ - self.flows._update(f) - return f - - def delete_flow(self, f): - self.flows._remove(f) - - def load_flows(self, flows): - self.flows._extend(flows) - - def set_view_filter(self, txt): - if txt == self.filter_txt: - return - if txt: - flt = flowfilter.parse(txt) - if not flt: - return "Invalid filter expression." - self.view._close() - self.view = FlowView(self.flows, flt) - else: - self.view._close() - self.view = FlowView(self.flows, None) - - def clear(self): - self.flows._clear() - - def accept_all(self, master): - self.flows.accept_all(master) - - def backup(self, f): - f.backup() - self.update_flow(f) - - def revert(self, f): - f.revert() - self.update_flow(f) - - def killall(self, master): - self.flows.kill_all(master) - - def duplicate_flow(self, f): - """ - Duplicate flow, and insert it into state without triggering any of - the normal flow events. - """ - f2 = f.copy() - self.add_flow(f2) - return f2 - - # Event handlers - def intercept(self, f): - self.update_flow(f) - - def resume(self, f): - self.update_flow(f) - - def error(self, f): - self.update_flow(f) - - def request(self, f): - if f not in self.flows: # don't add again on replay - self.add_flow(f) - - def response(self, f): - self.update_flow(f) diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index c09ff454..8fa5dc50 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -66,17 +66,17 @@ class OrderRequestStart(_OrderKey): class OrderRequestMethod(_OrderKey): - def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + def generate(self, f: mitmproxy.flow.Flow) -> str: return f.request.method class OrderRequestURL(_OrderKey): - def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + def generate(self, f: mitmproxy.flow.Flow) -> str: return f.request.url class OrderKeySize(_OrderKey): - def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + def generate(self, f: mitmproxy.flow.Flow) -> int: s = 0 if f.request.raw_content: s += len(f.request.raw_content) @@ -166,8 +166,8 @@ class View(collections.Sequence): return v return self._rev(v - 1) + 1 - def index(self, f: mitmproxy.flow.Flow) -> int: - return self._rev(self._view.index(f)) + def index(self, f: mitmproxy.flow.Flow, start: int = 0, stop: typing.Optional[int] = None) -> int: + return self._rev(self._view.index(f, start, stop)) def __contains__(self, f: mitmproxy.flow.Flow) -> bool: return self._view.__contains__(f) @@ -318,7 +318,7 @@ class Focus: """ def __init__(self, v: View) -> None: self.view = v - self._flow = None + self._flow = None # type: mitmproxy.flow.Flow self.sig_change = blinker.Signal() if len(self.view): self.flow = self.view[0] @@ -374,29 +374,29 @@ class Focus: class Settings(collections.Mapping): def __init__(self, view: View) -> None: self.view = view - self.values = {} + self._values = {} # type: typing.MutableMapping[str, mitmproxy.flow.Flow] view.sig_remove.connect(self._sig_remove) view.sig_refresh.connect(self._sig_refresh) def clear(self): - self.values.clear() + self._values.clear() - def __iter__(self) -> typing.Iterable: - return iter(self.values) + def __iter__(self) -> typing.Iterator: + return iter(self._values) def __len__(self) -> int: - return len(self.values) + return len(self._values) def __getitem__(self, f: mitmproxy.flow.Flow) -> dict: if f.id not in self.view._store: raise KeyError - return self.values.setdefault(f.id, {}) + return self._values.setdefault(f.id, {}) def _sig_remove(self, view, flow): - if flow.id in self.values: - del self.values[flow.id] + if flow.id in self._values: + del self._values[flow.id] def _sig_refresh(self, view): - for fid in list(self.values.keys()): + for fid in list(self._values.keys()): if fid not in view._store: - del self.values[fid] + del self._values[fid] diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py index 41701d36..bbdc2bf9 100644 --- a/mitmproxy/tools/web/app.py +++ b/mitmproxy/tools/web/app.py @@ -116,8 +116,8 @@ class RequestHandler(BasicAuth, tornado.web.RequestHandler): return json.loads(self.request.body.decode()) @property - def state(self): - return self.application.master.state + def view(self): + return self.application.master.view @property def master(self) -> "mitmproxy.tools.web.master.WebMaster": @@ -126,7 +126,8 @@ class RequestHandler(BasicAuth, tornado.web.RequestHandler): @property def flow(self): flow_id = str(self.path_kwargs["flow_id"]) - flow = self.state.flows.get(flow_id) + # FIXME: Add a facility to addon.view to safely access the store + flow = self.view._store.get(flow_id) if flow: return flow else: @@ -184,7 +185,7 @@ class Flows(RequestHandler): def get(self): self.write(dict( - data=[convert_flow_to_json_dict(f) for f in self.state.flows] + data=[convert_flow_to_json_dict(f) for f in self.view] )) @@ -195,31 +196,31 @@ class DumpFlows(RequestHandler): bio = BytesIO() fw = io.FlowWriter(bio) - for f in self.state.flows: + for f in self.view: fw.add(f) self.write(bio.getvalue()) bio.close() def post(self): - self.state.clear() + self.view.clear() content = self.request.files.values()[0][0].body bio = BytesIO(content) - self.state.load_flows(io.FlowReader(bio).stream()) + self.view.load_flows(io.FlowReader(bio).stream()) bio.close() class ClearAll(RequestHandler): def post(self): - self.state.clear() + self.view.clear() class AcceptFlows(RequestHandler): def post(self): - self.state.flows.accept_all(self.master) + self.view.flows.accept_all(self.master) class AcceptFlow(RequestHandler): @@ -233,7 +234,7 @@ class FlowHandler(RequestHandler): def delete(self, flow_id): if self.flow.killable: self.flow.kill(self.master) - self.state.delete_flow(self.flow) + self.view.delete_flow(self.flow) def put(self, flow_id): flow = self.flow @@ -270,19 +271,19 @@ class FlowHandler(RequestHandler): print("Warning: Unknown update {}.{}: {}".format(a, k, v)) else: print("Warning: Unknown update {}: {}".format(a, b)) - self.state.update_flow(flow) + self.view.update_flow(flow) class DuplicateFlow(RequestHandler): def post(self, flow_id): - self.master.state.duplicate_flow(self.flow) + self.master.view.duplicate_flow(self.flow) class RevertFlow(RequestHandler): def post(self, flow_id): - self.state.revert(self.flow) + self.flow.revert() class ReplayFlow(RequestHandler): @@ -290,7 +291,7 @@ class ReplayFlow(RequestHandler): def post(self, flow_id): self.flow.backup() self.flow.response = None - self.state.update_flow(self.flow) + self.view.update_flow(self.flow) r = self.master.replay_request(self.flow) if r: @@ -303,7 +304,7 @@ class FlowContent(RequestHandler): self.flow.backup() message = getattr(self.flow, message) message.content = self.request.files.values()[0][0].body - self.state.update_flow(self.flow) + self.view.update_flow(self.flow) def get(self, flow_id, message): message = getattr(self.flow, message) @@ -340,7 +341,7 @@ class FlowContentView(RequestHandler): message = getattr(self.flow, message) description, lines, error = contentviews.get_message_content_view( - contentviews.get(content_view.replace('_', ' ')), message + contentviews.get(content_view.replace('_', ' ')).name, message ) # if error: # add event log @@ -355,7 +356,7 @@ class Events(RequestHandler): def get(self): self.write(dict( - data=list(self.state.events) + data=list([]) )) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index cf0bfd73..d283e5d5 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -1,5 +1,4 @@ import sys -import collections import tornado.httpserver import tornado.ioloop @@ -8,7 +7,7 @@ from typing import Optional from mitmproxy import addons from mitmproxy import exceptions -from mitmproxy.addons import state +from mitmproxy.addons import view from mitmproxy.addons import intercept from mitmproxy import options from mitmproxy import master @@ -20,53 +19,7 @@ class Stop(Exception): pass -class WebFlowView(state.FlowView): - - def __init__(self, store): - super().__init__(store, None) - - def _add(self, f): - super()._add(f) - app.ClientConnection.broadcast( - type="UPDATE_FLOWS", - cmd="add", - data=app.convert_flow_to_json_dict(f) - ) - - def _update(self, f): - super()._update(f) - app.ClientConnection.broadcast( - type="UPDATE_FLOWS", - cmd="update", - data=app.convert_flow_to_json_dict(f) - ) - - def _remove(self, f): - super()._remove(f) - app.ClientConnection.broadcast( - type="UPDATE_FLOWS", - cmd="remove", - data=dict(id=f.id) - ) - - def _recalculate(self, flows): - super()._recalculate(flows) - app.ClientConnection.broadcast( - type="UPDATE_FLOWS", - cmd="reset" - ) - - -class WebState(state.State): - - def __init__(self): - super().__init__() - self.view._close() - self.view = WebFlowView(self.flows) - - self._last_event_id = 0 - self.events = collections.deque(maxlen=1000) - +class _WebState(): def add_log(self, e, level): self._last_event_id += 1 entry = { @@ -136,9 +89,14 @@ class WebMaster(master.Master): def __init__(self, options, server): super().__init__(options, server) - self.state = WebState() + self.view = view.View() + self.view.sig_add.connect(self._sig_add) + self.view.sig_remove.connect(self._sig_remove) + self.view.sig_update.connect(self._sig_update) + self.view.sig_refresh.connect(self._sig_refresh) + self.addons.add(*addons.default_addons()) - self.addons.add(self.state, intercept.Intercept()) + self.addons.add(self.view, intercept.Intercept()) self.app = app.Application( self, self.options.wdebug, self.options.wauthenticator ) @@ -162,6 +120,33 @@ class WebMaster(master.Master): print("Stream file error: {}".format(err), file=sys.stderr) sys.exit(1) + def _sig_add(self, view, flow): + app.ClientConnection.broadcast( + type="UPDATE_FLOWS", + cmd="add", + data=app.convert_flow_to_json_dict(flow) + ) + + def _sig_update(self, view, flow): + app.ClientConnection.broadcast( + type="UPDATE_FLOWS", + cmd="update", + data=app.convert_flow_to_json_dict(flow) + ) + + def _sig_remove(self, view, flow): + app.ClientConnection.broadcast( + type="UPDATE_FLOWS", + cmd="remove", + data=dict(id=flow.id) + ) + + def _sig_refresh(self, view): + app.ClientConnection.broadcast( + type="UPDATE_FLOWS", + cmd="reset" + ) + def run(self): # pragma: no cover iol = tornado.ioloop.IOLoop.instance() @@ -178,6 +163,6 @@ class WebMaster(master.Master): except (Stop, KeyboardInterrupt): self.shutdown() - def add_log(self, e, level="info"): - super().add_log(e, level) - return self.state.add_log(e, level) + # def add_log(self, e, level="info"): + # super().add_log(e, level) + # return self.state.add_log(e, level) diff --git a/test/mitmproxy/addons/test_state.py b/test/mitmproxy/addons/test_state.py deleted file mode 100644 index 16e79e0d..00000000 --- a/test/mitmproxy/addons/test_state.py +++ /dev/null @@ -1,19 +0,0 @@ -from mitmproxy.test import tflow -from mitmproxy import proxy -from mitmproxy import master -from mitmproxy.addons import state - - -class TestState: - def test_duplicate_flow(self): - s = state.State() - fm = master.Master(None, proxy.DummyServer()) - fm.addons.add(s) - f = tflow.tflow(resp=True) - fm.load_flow(f) - assert s.flow_count() == 1 - - f2 = s.duplicate_flow(f) - assert f2.response - assert s.flow_count() == 2 - assert s.index(f2) == 1 diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 7a9d0901..295e2bea 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -6,7 +6,6 @@ import mitmproxy.test.tutils from mitmproxy.net.http import Headers import mitmproxy.io from mitmproxy import flowfilter, options -from mitmproxy.addons import state from mitmproxy.contrib import tnetstring from mitmproxy.exceptions import FlowReadException, Kill from mitmproxy import flow @@ -15,7 +14,7 @@ from mitmproxy import connections from mitmproxy.proxy import ProxyConfig from mitmproxy.proxy.server import DummyServer from mitmproxy import master -from . import tutils +from . import tutils, tservers class TestHTTPFlow: @@ -107,20 +106,6 @@ class TestHTTPFlow: assert not f.killable assert f.reply.value == Kill - def test_killall(self): - srv = DummyServer(None) - s = state.State() - fm = master.Master(None, srv) - fm.addons.add(s) - - f = tflow.tflow() - f.reply.handle() - f.intercept(fm) - - s.killall(fm) - for i in s.view: - assert "killed" in str(i.error) - def test_resume(self): f = tflow.tflow() f.reply.handle() @@ -186,135 +171,6 @@ class TestTCPFlow: tutils.raises(ValueError, flowfilter.match, "~", f) -class TestState: - - def test_backup(self): - c = state.State() - f = tflow.tflow() - c.add_flow(f) - f.backup() - c.revert(f) - - def test_flow(self): - """ - normal flow: - - connect -> request -> response - """ - c = state.State() - f = tflow.tflow() - c.add_flow(f) - assert f - assert c.flow_count() == 1 - assert c.active_flow_count() == 1 - - newf = tflow.tflow() - assert c.add_flow(newf) - assert c.active_flow_count() == 2 - - f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - assert c.update_flow(f) - assert c.flow_count() == 2 - assert c.active_flow_count() == 1 - - assert not c.update_flow(None) - assert c.active_flow_count() == 1 - - newf.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - assert c.update_flow(newf) - assert c.active_flow_count() == 0 - - def test_err(self): - c = state.State() - f = tflow.tflow() - c.add_flow(f) - f.error = flow.Error("message") - assert c.update_flow(f) - - c = state.State() - f = tflow.tflow() - c.add_flow(f) - c.set_view_filter("~e") - assert not c.view - f.error = tflow.terr() - assert c.update_flow(f) - assert c.view - - def test_set_view_filter(self): - c = state.State() - - f = tflow.tflow() - assert len(c.view) == 0 - - c.add_flow(f) - assert len(c.view) == 1 - - c.set_view_filter("~s") - assert c.filter_txt == "~s" - assert len(c.view) == 0 - f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - c.update_flow(f) - assert len(c.view) == 1 - c.set_view_filter(None) - assert len(c.view) == 1 - - f = tflow.tflow() - c.add_flow(f) - assert len(c.view) == 2 - c.set_view_filter("~q") - assert len(c.view) == 1 - c.set_view_filter("~s") - assert len(c.view) == 1 - - assert "Invalid" in c.set_view_filter("~") - - def _add_request(self, state): - f = tflow.tflow() - state.add_flow(f) - return f - - def _add_response(self, state): - f = tflow.tflow() - state.add_flow(f) - f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - state.update_flow(f) - - def _add_error(self, state): - f = tflow.tflow(err=True) - state.add_flow(f) - - def test_clear(self): - c = state.State() - f = self._add_request(c) - f.intercepted = True - - c.clear() - assert c.flow_count() == 0 - - def test_dump_flows(self): - c = state.State() - self._add_request(c) - self._add_response(c) - self._add_request(c) - self._add_response(c) - self._add_request(c) - self._add_response(c) - self._add_error(c) - - flows = c.view[:] - c.clear() - - c.load_flows(flows) - assert isinstance(c.flows[0], flow.Flow) - - def test_accept_all(self): - c = state.State() - self._add_request(c) - self._add_response(c) - self._add_request(c) - c.accept_all(mock.Mock()) - - class TestSerialize: def _treader(self): @@ -354,7 +210,7 @@ class TestSerialize: def test_load_flows(self): r = self._treader() - s = state.State() + s = tservers.TestState() fm = master.Master(None, DummyServer()) fm.addons.add(s) fm.load_flows(r) @@ -362,7 +218,7 @@ class TestSerialize: def test_load_flows_reverse(self): r = self._treader() - s = state.State() + s = tservers.TestState() opts = options.Options( mode="reverse", upstream_server="https://use-this-domain" @@ -431,7 +287,7 @@ class TestFlowMaster: assert fm.create_request("GET", "http", "example.com", 80, "/") def test_all(self): - s = state.State() + s = tservers.TestState() fm = master.Master(None, DummyServer()) fm.addons.add(s) f = tflow.tflow(req=None) diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 7dd738f4..74cb5ae8 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -50,10 +50,10 @@ class CommonMixin: def test_replay(self): assert self.pathod("304").status_code == 304 if isinstance(self, tservers.HTTPUpstreamProxyTest) and self.ssl: - assert len(self.master.state.view) == 2 + assert len(self.master.state.flows) == 2 else: - assert len(self.master.state.view) == 1 - l = self.master.state.view[-1] + assert len(self.master.state.flows) == 1 + l = self.master.state.flows[-1] assert l.response.status_code == 304 l.request.path = "/p/305" self.wait_until_not_live(l) @@ -87,7 +87,7 @@ class CommonMixin: # In Upstream mode with SSL, we may already have a previous CONNECT # request. - l = self.master.state.view[-1] + l = self.master.state.flows[-1] assert l.client_conn.address assert "host" in l.request.headers assert l.response.status_code == 304 @@ -589,7 +589,7 @@ class TestProxy(tservers.HTTPProxyTest): f = self.pathod("304") assert f.status_code == 304 - f = self.master.state.view[0] + f = self.master.state.flows[0] assert f.client_conn.address assert "host" in f.request.headers assert f.response.status_code == 304 @@ -601,7 +601,7 @@ class TestProxy(tservers.HTTPProxyTest): f = self.pathod("304:b@1k:p50,1") assert f.status_code == 304 - response = self.master.state.view[0].response + response = self.master.state.flows[0].response # timestamp_start might fire a bit late, so we play safe and only require 300ms. assert 0.3 <= response.timestamp_end - response.timestamp_start @@ -620,8 +620,8 @@ class TestProxy(tservers.HTTPProxyTest): connection.recv(50000) connection.close() - request, response = self.master.state.view[ - 0].request, self.master.state.view[0].response + request, response = self.master.state.flows[ + 0].request, self.master.state.flows[0].response assert response.status_code == 304 # sanity test for our low level request # timestamp_start might fire a bit late, so we play safe and only require 300ms. assert 0.3 <= request.timestamp_end - request.timestamp_start @@ -647,8 +647,8 @@ class TestProxy(tservers.HTTPProxyTest): recvd += len(connection.recv(5000)) connection.close() - first_flow = self.master.state.view[0] - second_flow = self.master.state.view[1] + first_flow = self.master.state.flows[0] + second_flow = self.master.state.flows[1] assert first_flow.server_conn.timestamp_tcp_setup assert first_flow.server_conn.timestamp_ssl_setup is None assert second_flow.server_conn.timestamp_tcp_setup @@ -657,7 +657,7 @@ class TestProxy(tservers.HTTPProxyTest): def test_request_ip(self): f = self.pathod("200:b@100") assert f.status_code == 200 - f = self.master.state.view[0] + f = self.master.state.flows[0] assert f.server_conn.address == ("127.0.0.1", self.server.port) @@ -668,7 +668,7 @@ class TestProxySSL(tservers.HTTPProxyTest): # tests that the ssl timestamp is present when ssl is used f = self.pathod("304:b@10k") assert f.status_code == 304 - first_flow = self.master.state.view[0] + first_flow = self.master.state.flows[0] assert first_flow.server_conn.timestamp_ssl_setup diff --git a/test/mitmproxy/test_web_master.py b/test/mitmproxy/test_web_master.py index 77280644..298b14eb 100644 --- a/test/mitmproxy/test_web_master.py +++ b/test/mitmproxy/test_web_master.py @@ -12,4 +12,4 @@ class TestWebMaster(mastertest.MasterTest): m = self.mkmaster() for i in (1, 2, 3): self.dummy_cycle(m, 1, b"") - assert len(m.state.flows) == i + assert len(m.view) == i diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index e07102e1..f9dfde30 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -7,7 +7,6 @@ import sys from mitmproxy.proxy.config import ProxyConfig from mitmproxy.proxy.server import ProxyServer from mitmproxy import master -from mitmproxy.addons import state import pathod.test import pathod.pathoc from mitmproxy import controller @@ -15,6 +14,23 @@ from mitmproxy import options from mitmproxy import exceptions +class TestState: + def __init__(self): + self.flows = [] + + def request(self, f): + if f not in self.flows: + self.flows.append(f) + + def response(self, f): + if f not in self.flows: + self.flows.append(f) + + # FIXME: compat with old state - remove in favor of len(state.flows) + def flow_count(self): + return len(self.flows) + + class TestMaster(master.Master): def __init__(self, opts, config): @@ -23,7 +39,7 @@ class TestMaster(master.Master): def clear_addons(self, addons): self.addons.clear() - self.state = state.State() + self.state = TestState() self.addons.add(self.state) self.addons.add(*addons) @@ -21,4 +21,4 @@ commands = sphinx-build -W -b html -d {envtmpdir}/doctrees . {envtmpdir}/html commands = flake8 --jobs 8 --count mitmproxy pathod examples test rstcheck README.rst - mypy --silent-imports mitmproxy/addonmanager.py mitmproxy/proxy/protocol/ + mypy --silent-imports mitmproxy/addons mitmproxy/addonmanager.py mitmproxy/proxy/protocol/ |