From 7c32d4ea2a435484e83aa459831f74ca483d8e3c Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 19 Oct 2016 14:48:42 +1300 Subject: flow.state -> addons.state --- examples/flowbasic | 2 +- mitmproxy/addons/onboarding/app.py | 93 ------------ mitmproxy/addons/state.py | 288 ++++++++++++++++++++++++++++++++++++ mitmproxy/console/master.py | 5 +- mitmproxy/flow/__init__.py | 2 - mitmproxy/flow/state.py | 288 ------------------------------------ mitmproxy/web/master.py | 6 +- test/mitmproxy/addons/test_state.py | 20 +++ test/mitmproxy/test_flow.py | 27 ++-- test/mitmproxy/test_flow_state.py | 19 --- test/mitmproxy/tservers.py | 2 +- 11 files changed, 330 insertions(+), 422 deletions(-) delete mode 100644 mitmproxy/addons/onboarding/app.py create mode 100644 mitmproxy/addons/state.py delete mode 100644 mitmproxy/flow/state.py create mode 100644 test/mitmproxy/addons/test_state.py delete mode 100644 test/mitmproxy/test_flow_state.py diff --git a/examples/flowbasic b/examples/flowbasic index 0eb163a4..bac98916 100755 --- a/examples/flowbasic +++ b/examples/flowbasic @@ -37,7 +37,7 @@ class MyMaster(master.Master): opts = options.Options(cadir="~/.mitmproxy/") config = ProxyConfig(opts) -state = flow.State() +state = state.State() server = ProxyServer(config) m = MyMaster(opts, server, state) m.run() diff --git a/mitmproxy/addons/onboarding/app.py b/mitmproxy/addons/onboarding/app.py deleted file mode 100644 index 20197c09..00000000 --- a/mitmproxy/addons/onboarding/app.py +++ /dev/null @@ -1,93 +0,0 @@ -import os - -import tornado.template -import tornado.web -import tornado.wsgi - -from mitmproxy import utils -from mitmproxy.proxy import config - -loader = tornado.template.Loader(utils.pkg_data.path("addons/onboardingapp/templates")) - - -class Adapter(tornado.wsgi.WSGIAdapter): - # Tornado doesn't make the WSGI environment available to pages, so this - # hideous monkey patch is the easiest way to get to the mitmproxy.master - # variable. - - def __init__(self, application): - self._application = application - - def application(self, request): - request.master = self.environ["mitmproxy.master"] - return self._application(request) - - def __call__(self, environ, start_response): - self.environ = environ - return tornado.wsgi.WSGIAdapter.__call__( - self, - environ, - start_response - ) - - -class Index(tornado.web.RequestHandler): - - def get(self): - t = loader.load("index.html") - self.write(t.generate()) - - -class PEM(tornado.web.RequestHandler): - - @property - def filename(self): - return config.CONF_BASENAME + "-ca-cert.pem" - - def get(self): - p = os.path.join(self.request.master.options.cadir, self.filename) - p = os.path.expanduser(p) - self.set_header("Content-Type", "application/x-x509-ca-cert") - self.set_header( - "Content-Disposition", - "inline; filename={}".format( - self.filename)) - - with open(p, "rb") as f: - self.write(f.read()) - - -class P12(tornado.web.RequestHandler): - - @property - def filename(self): - return config.CONF_BASENAME + "-ca-cert.p12" - - def get(self): - p = os.path.join(self.request.master.options.cadir, self.filename) - p = os.path.expanduser(p) - self.set_header("Content-Type", "application/x-pkcs12") - self.set_header( - "Content-Disposition", - "inline; filename={}".format( - self.filename)) - - with open(p, "rb") as f: - self.write(f.read()) - - -application = tornado.web.Application( - [ - (r"/", Index), - (r"/cert/pem", PEM), - (r"/cert/p12", P12), - ( - r"/static/(.*)", - tornado.web.StaticFileHandler, - { - "path": utils.pkg_data.path("onboarding/static") - } - ), - ], - # debug=True -) diff --git a/mitmproxy/addons/state.py b/mitmproxy/addons/state.py new file mode 100644 index 00000000..bb7460b6 --- /dev/null +++ b/mitmproxy/addons/state.py @@ -0,0 +1,288 @@ +from abc import abstractmethod, ABCMeta + +from typing import List # noqa + +from mitmproxy import flowfilter +from mitmproxy import models # noqa + + +class FlowList(metaclass=ABCMeta): + def __init__(self): + self._list = [] # type: List[models.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) + + # These are compiled filter expressions: + self.intercept = 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 set_intercept(self, txt): + if txt: + flt = flowfilter.parse(txt) + if not flt: + return "Invalid filter expression." + self.intercept = flt + else: + self.intercept = None + + @property + def intercept_txt(self): + return getattr(self.intercept, "pattern", 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/console/master.py b/mitmproxy/console/master.py index 4e49a51a..46dd8254 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -22,6 +22,7 @@ from mitmproxy import master from mitmproxy import flow from mitmproxy import flowfilter from mitmproxy import utils +from mitmproxy.addons import state import mitmproxy.options from mitmproxy.console import flowlist from mitmproxy.console import flowview @@ -39,10 +40,10 @@ from netlib import tcp, strutils EVENTLOG_SIZE = 500 -class ConsoleState(flow.State): +class ConsoleState(state.State): def __init__(self): - flow.State.__init__(self) + state.State.__init__(self) self.focus = None self.follow_focus = None self.default_body_view = contentviews.get("Auto") diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py index 2688966a..69e3ed08 100644 --- a/mitmproxy/flow/__init__.py +++ b/mitmproxy/flow/__init__.py @@ -1,9 +1,7 @@ from mitmproxy.flow import export from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths -from mitmproxy.flow.state import State, FlowView __all__ = [ "export", "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths", - "State", "FlowView", ] diff --git a/mitmproxy/flow/state.py b/mitmproxy/flow/state.py deleted file mode 100644 index bb7460b6..00000000 --- a/mitmproxy/flow/state.py +++ /dev/null @@ -1,288 +0,0 @@ -from abc import abstractmethod, ABCMeta - -from typing import List # noqa - -from mitmproxy import flowfilter -from mitmproxy import models # noqa - - -class FlowList(metaclass=ABCMeta): - def __init__(self): - self._list = [] # type: List[models.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) - - # These are compiled filter expressions: - self.intercept = 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 set_intercept(self, txt): - if txt: - flt = flowfilter.parse(txt) - if not flt: - return "Invalid filter expression." - self.intercept = flt - else: - self.intercept = None - - @property - def intercept_txt(self): - return getattr(self.intercept, "pattern", 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/web/master.py b/mitmproxy/web/master.py index c7976fcb..f07d28c8 100644 --- a/mitmproxy/web/master.py +++ b/mitmproxy/web/master.py @@ -9,7 +9,7 @@ from typing import Optional from mitmproxy import addons from mitmproxy import controller from mitmproxy import exceptions -from mitmproxy import flow +from mitmproxy.addons import state from mitmproxy import options from mitmproxy import master from mitmproxy.web import app @@ -20,7 +20,7 @@ class Stop(Exception): pass -class WebFlowView(flow.FlowView): +class WebFlowView(state.FlowView): def __init__(self, store): super().__init__(store, None) @@ -57,7 +57,7 @@ class WebFlowView(flow.FlowView): ) -class WebState(flow.State): +class WebState(state.State): def __init__(self): super().__init__() diff --git a/test/mitmproxy/addons/test_state.py b/test/mitmproxy/addons/test_state.py new file mode 100644 index 00000000..71c46dcb --- /dev/null +++ b/test/mitmproxy/addons/test_state.py @@ -0,0 +1,20 @@ +from mitmproxy import proxy +from mitmproxy import master +from mitmproxy.addons import state + +from .. import tutils + + +class TestState: + def test_duplicate_flow(self): + s = state.State() + fm = master.Master(None, proxy.DummyServer()) + fm.addons.add(s) + f = tutils.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 2b387f5c..86cd7d16 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -4,6 +4,7 @@ import io import netlib.utils from netlib.http import Headers from mitmproxy import flowfilter, flow, options +from mitmproxy.addons import state from mitmproxy.contrib import tnetstring from mitmproxy.exceptions import FlowReadException, Kill from mitmproxy.models import Error @@ -110,7 +111,7 @@ class TestHTTPFlow: def test_killall(self): srv = DummyServer(None) - s = flow.State() + s = state.State() fm = master.Master(None, srv) fm.addons.add(s) @@ -190,7 +191,7 @@ class TestTCPFlow: class TestState: def test_backup(self): - c = flow.State() + c = state.State() f = tutils.tflow() c.add_flow(f) f.backup() @@ -202,7 +203,7 @@ class TestState: connect -> request -> response """ - c = flow.State() + c = state.State() f = tutils.tflow() c.add_flow(f) assert f @@ -226,13 +227,13 @@ class TestState: assert c.active_flow_count() == 0 def test_err(self): - c = flow.State() + c = state.State() f = tutils.tflow() c.add_flow(f) f.error = Error("message") assert c.update_flow(f) - c = flow.State() + c = state.State() f = tutils.tflow() c.add_flow(f) c.set_view_filter("~e") @@ -242,7 +243,7 @@ class TestState: assert c.view def test_set_view_filter(self): - c = flow.State() + c = state.State() f = tutils.tflow() assert len(c.view) == 0 @@ -270,7 +271,7 @@ class TestState: assert "Invalid" in c.set_view_filter("~") def test_set_intercept(self): - c = flow.State() + c = state.State() assert not c.set_intercept("~q") assert c.intercept_txt == "~q" assert "Invalid" in c.set_intercept("~") @@ -293,7 +294,7 @@ class TestState: state.add_flow(f) def test_clear(self): - c = flow.State() + c = state.State() f = self._add_request(c) f.intercepted = True @@ -301,7 +302,7 @@ class TestState: assert c.flow_count() == 0 def test_dump_flows(self): - c = flow.State() + c = state.State() self._add_request(c) self._add_response(c) self._add_request(c) @@ -317,7 +318,7 @@ class TestState: assert isinstance(c.flows[0], Flow) def test_accept_all(self): - c = flow.State() + c = state.State() self._add_request(c) self._add_response(c) self._add_request(c) @@ -363,7 +364,7 @@ class TestSerialize: def test_load_flows(self): r = self._treader() - s = flow.State() + s = state.State() fm = master.Master(None, DummyServer()) fm.addons.add(s) fm.load_flows(r) @@ -371,7 +372,7 @@ class TestSerialize: def test_load_flows_reverse(self): r = self._treader() - s = flow.State() + s = state.State() opts = options.Options( mode="reverse", upstream_server="https://use-this-domain" @@ -440,7 +441,7 @@ class TestFlowMaster: assert fm.create_request("GET", "http", "example.com", 80, "/") def test_all(self): - s = flow.State() + s = state.State() fm = master.Master(None, DummyServer()) fm.addons.add(s) f = tutils.tflow(req=None) diff --git a/test/mitmproxy/test_flow_state.py b/test/mitmproxy/test_flow_state.py deleted file mode 100644 index 05f4cbb4..00000000 --- a/test/mitmproxy/test_flow_state.py +++ /dev/null @@ -1,19 +0,0 @@ -from mitmproxy import flow -from mitmproxy import proxy -from mitmproxy import master -from . import tutils - - -class TestState: - def test_duplicate_flow(self): - s = flow.State() - fm = master.Master(None, proxy.DummyServer()) - fm.addons.add(s) - f = tutils.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/tservers.py b/test/mitmproxy/tservers.py index 254af2f0..1243bca0 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -7,7 +7,7 @@ import sys from mitmproxy.proxy.config import ProxyConfig from mitmproxy.proxy.server import ProxyServer from mitmproxy import master -from mitmproxy.flow import state +from mitmproxy.addons import state import pathod.test import pathod.pathoc from mitmproxy import controller, options -- cgit v1.2.3