From b0cfeff06d9dd99a16dfae19c5df3c73c5864fb9 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 3 Sep 2014 16:57:56 +0200 Subject: fix #341 - work on flows instead of request/response internally. --- libmproxy/flow.py | 106 +++++++++++++++++++++++++----------------------------- 1 file changed, 48 insertions(+), 58 deletions(-) (limited to 'libmproxy/flow.py') diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 2540435e..eb183d9f 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -34,11 +34,11 @@ class AppRegistry: """ Returns an WSGIAdaptor instance if request matches an app, or None. """ - if (request.get_host(), request.get_port()) in self.apps: - return self.apps[(request.get_host(), request.get_port())] + if (request.host, request.port) in self.apps: + return self.apps[(request.host, request.port)] if "host" in request.headers: host = request.headers["host"][0] - return self.apps.get((host, request.get_port()), None) + return self.apps.get((host, request.port), None) class ReplaceHooks: @@ -185,11 +185,11 @@ class ClientPlaybackState: n = self.flows.pop(0) n.request.reply = controller.DummyReply() n.client_conn = None - self.current = master.handle_request(n.request) + self.current = master.handle_request(n) if not testing and not self.current.response: - master.replay_request(self.current) # pragma: no cover + master.replay_request(self.current) # pragma: no cover elif self.current.response: - master.handle_response(self.current.response) + master.handle_response(self.current) class ServerPlaybackState: @@ -260,8 +260,8 @@ class StickyCookieState: Returns a (domain, port, path) tuple. """ return ( - m["domain"] or f.request.get_host(), - f.request.get_port(), + m["domain"] or f.request.get_host(False, f), + f.request.get_port(f), m["path"] or "/" ) @@ -279,7 +279,7 @@ class StickyCookieState: c = Cookie.SimpleCookie(str(i)) m = c.values()[0] k = self.ckey(m, f) - if self.domain_match(f.request.get_host(), k[0]): + if self.domain_match(f.request.get_host(False, f), k[0]): self.jar[self.ckey(m, f)] = m def handle_request(self, f): @@ -287,8 +287,8 @@ class StickyCookieState: if f.match(self.flt): for i in self.jar.keys(): match = [ - self.domain_match(f.request.get_host(), i[0]), - f.request.get_port() == i[1], + self.domain_match(f.request.get_host(False, f), i[0]), + f.request.get_port(f) == i[1], f.request.path.startswith(i[2]) ] if all(match): @@ -307,7 +307,7 @@ class StickyAuthState: self.hosts = {} def handle_request(self, f): - host = f.request.get_host() + host = f.request.get_host(False, f) if "authorization" in f.request.headers: self.hosts[host] = f.request.headers["authorization"] elif f.match(self.flt): @@ -342,33 +342,30 @@ class State(object): c += 1 return c - def add_request(self, req): + def add_request(self, flow): """ Add a request to the state. Returns the matching flow. """ - f = req.flow - self._flow_list.append(f) - if f.match(self._limit): - self.view.append(f) - return f + self._flow_list.append(flow) + if flow.match(self._limit): + self.view.append(flow) + return flow - def add_response(self, resp): + def add_response(self, f): """ Add a response to the state. Returns the matching flow. """ - f = resp.flow if not f: return False if f.match(self._limit) and not f in self.view: self.view.append(f) return f - def add_error(self, err): + def add_error(self, f): """ Add an error response to the state. Returns the matching flow, or None if there isn't one. """ - f = err.flow if not f: return None if f.match(self._limit) and not f in self.view: @@ -586,7 +583,7 @@ class FlowMaster(controller.Master): response.is_replay = True if self.refresh_server_playback: response.refresh() - flow.request.reply(response) + flow.reply(response) if self.server_playback.count() == 0: self.stop_server_playback() return True @@ -612,16 +609,14 @@ class FlowMaster(controller.Master): """ Loads a flow, and returns a new flow object. """ + f.reply = controller.DummyReply() if f.request: - f.request.reply = controller.DummyReply() - fr = self.handle_request(f.request) + self.handle_request(f) if f.response: - f.response.reply = controller.DummyReply() - self.handle_response(f.response) + self.handle_response(f) if f.error: - f.error.reply = controller.DummyReply() - self.handle_error(f.error) - return fr + self.handle_error(f) + return f def load_flows(self, fr): """ @@ -647,7 +642,7 @@ class FlowMaster(controller.Master): if self.kill_nonreplay: f.kill(self) else: - f.request.reply() + f.reply() def process_new_response(self, f): if self.stickycookie_state: @@ -694,54 +689,49 @@ class FlowMaster(controller.Master): self.run_script_hook("serverconnect", sc) sc.reply() - def handle_error(self, r): - f = self.state.add_error(r) - if f: - self.run_script_hook("error", f) + def handle_error(self, f): + self.state.add_error(f) + self.run_script_hook("error", f) if self.client_playback: self.client_playback.clear(f) - r.reply() + f.reply() return f - def handle_request(self, r): - if r.flow.live: - app = self.apps.get(r) + def handle_request(self, f): + if f.live: + app = self.apps.get(f.request) if app: - err = app.serve(r, r.flow.client_conn.wfile, **{"mitmproxy.master": self}) + err = app.serve(f, f.client_conn.wfile, **{"mitmproxy.master": self}) if err: self.add_event("Error in wsgi app. %s"%err, "error") - r.reply(protocol.KILL) + f.reply(protocol.KILL) return - f = self.state.add_request(r) + self.state.add_request(f) self.replacehooks.run(f) self.setheaders.run(f) self.run_script_hook("request", f) self.process_new_request(f) return f - def handle_responseheaders(self, resp): - f = resp.flow + def handle_responseheaders(self, f): self.run_script_hook("responseheaders", f) if self.stream_large_bodies: self.stream_large_bodies.run(f, False) - resp.reply() + f.reply() return f - def handle_response(self, r): - f = self.state.add_response(r) - if f: - self.replacehooks.run(f) - self.setheaders.run(f) - self.run_script_hook("response", f) - if self.client_playback: - self.client_playback.clear(f) - self.process_new_response(f) - if self.stream: - self.stream.add(f) - else: - r.reply() + def handle_response(self, f): + self.state.add_response(f) + self.replacehooks.run(f) + self.setheaders.run(f) + self.run_script_hook("response", f) + if self.client_playback: + self.client_playback.clear(f) + self.process_new_response(f) + if self.stream: + self.stream.add(f) return f def shutdown(self): -- cgit v1.2.3 From 2f44b26b4cd014e03dd62a125d79af9b81663a93 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 3 Sep 2014 23:44:54 +0200 Subject: improve HTTPRequest syntax --- libmproxy/flow.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'libmproxy/flow.py') diff --git a/libmproxy/flow.py b/libmproxy/flow.py index eb183d9f..9115ec9d 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -260,8 +260,8 @@ class StickyCookieState: Returns a (domain, port, path) tuple. """ return ( - m["domain"] or f.request.get_host(False, f), - f.request.get_port(f), + m["domain"] or f.request.host, + f.request.port, m["path"] or "/" ) @@ -279,7 +279,7 @@ class StickyCookieState: c = Cookie.SimpleCookie(str(i)) m = c.values()[0] k = self.ckey(m, f) - if self.domain_match(f.request.get_host(False, f), k[0]): + if self.domain_match(f.request.host, k[0]): self.jar[self.ckey(m, f)] = m def handle_request(self, f): @@ -287,8 +287,8 @@ class StickyCookieState: if f.match(self.flt): for i in self.jar.keys(): match = [ - self.domain_match(f.request.get_host(False, f), i[0]), - f.request.get_port(f) == i[1], + self.domain_match(f.request.host, i[0]), + f.request.port == i[1], f.request.path.startswith(i[2]) ] if all(match): @@ -307,7 +307,7 @@ class StickyAuthState: self.hosts = {} def handle_request(self, f): - host = f.request.get_host(False, f) + host = f.request.host if "authorization" in f.request.headers: self.hosts[host] = f.request.headers["authorization"] elif f.match(self.flt): -- cgit v1.2.3 From 649e63ff3c868397f493e1dabdc1c63d572aedd8 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 4 Sep 2014 00:10:01 +0200 Subject: fix some leftovers --- libmproxy/flow.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) (limited to 'libmproxy/flow.py') diff --git a/libmproxy/flow.py b/libmproxy/flow.py index eb183d9f..df72878f 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -183,8 +183,7 @@ class ClientPlaybackState: """ if self.flows and not self.current: n = self.flows.pop(0) - n.request.reply = controller.DummyReply() - n.client_conn = None + n.reply = controller.DummyReply() self.current = master.handle_request(n) if not testing and not self.current.response: master.replay_request(self.current) # pragma: no cover -- cgit v1.2.3 From b23a1aa4a4dd9f09fc199d03f546a8fafc8b27b8 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 4 Sep 2014 19:08:54 +0200 Subject: much tests. so tcp. very wow. --- libmproxy/flow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'libmproxy/flow.py') diff --git a/libmproxy/flow.py b/libmproxy/flow.py index eeb53e81..086710bc 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -612,6 +612,7 @@ class FlowMaster(controller.Master): if f.request: self.handle_request(f) if f.response: + self.handle_responseheaders(f) self.handle_response(f) if f.error: self.handle_error(f) @@ -668,7 +669,7 @@ class FlowMaster(controller.Master): self.masterq, self.should_exit ) - rt.start() # pragma: no cover + rt.start() # pragma: no cover if block: rt.join() -- cgit v1.2.3