diff options
| -rw-r--r-- | libmproxy/console.py | 31 | ||||
| -rw-r--r-- | libmproxy/flow.py | 36 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 2 |
3 files changed, 49 insertions, 20 deletions
diff --git a/libmproxy/console.py b/libmproxy/console.py index dbeef04a..794d5483 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -779,7 +779,7 @@ VIEW_CONNLIST = 0 VIEW_FLOW = 1 VIEW_HELP = 2 -class ConsoleMaster(controller.Master): +class ConsoleMaster(flow.FlowMaster): palette = [] footer_text_default = [ ('key', "?"), ":help ", @@ -794,10 +794,10 @@ class ConsoleMaster(controller.Master): ('key', "q"), ":back ", ] def __init__(self, server, options): + flow.FlowMaster.__init__(self, server, ConsoleState()) + self.conn_list_view = None self.set_palette() - controller.Master.__init__(self, server) - self.state = ConsoleState() r = self.set_limit(options.limit) if r: @@ -1313,23 +1313,18 @@ class ConsoleMaster(controller.Master): # Handlers def handle_clientconnection(self, r): - f = flow.Flow(r) - self.state.add_browserconnect(f) - r.ack() - self.sync_list_view() + f = flow.FlowMaster.handle_clientconnection(self, r) + if f: + self.sync_list_view() def handle_error(self, r): - f = self.state.add_error(r) - if not f: - r.ack() - else: + f = flow.FlowMaster.handle_error(self, r) + if f: self.process_flow(f, r) def handle_request(self, r): - f = self.state.add_request(r) - if not f: - r.ack() - else: + f = flow.FlowMaster.handle_request(self, r) + if f: if f.match(self.stickycookie): hid = (f.request.host, f.request.port) if f.request.headers.has_key("cookie"): @@ -1339,10 +1334,8 @@ class ConsoleMaster(controller.Master): self.process_flow(f, r) def handle_response(self, r): - f = self.state.add_response(r) - if not f: - r.ack() - else: + f = flow.FlowMaster.handle_request(self, r) + if f: if f.match(self.stickycookie): hid = (f.request.host, f.request.port) if f.response.headers.has_key("set-cookie"): diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 921abe1d..e791055f 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -4,6 +4,7 @@ """ import subprocess, base64, sys, json import proxy, threading, netstring +import controller class RunException(Exception): def __init__(self, msg, returncode, errout): @@ -286,6 +287,41 @@ class State: #end nocover +class FlowMaster(controller.Master): + def __init__(self, server, state): + controller.Master.__init__(self, server) + self.state = state + + # Handlers + def handle_clientconnection(self, r): + f = Flow(r) + self.state.add_browserconnect(f) + r.ack() + return f + + def handle_error(self, r): + f = self.state.add_error(r) + if not f: + r.ack() + return f + + def handle_request(self, r): + f = self.state.add_request(r) + if not f: + r.ack() + return f + + def handle_response(self, r): + f = self.state.add_response(r) + if not f: + r.ack() + else: + if f.match(self.stickycookie): + hid = (f.request.host, f.request.port) + if f.response.headers.has_key("set-cookie"): + self.stickyhosts[hid] = f.response.headers["set-cookie"] + self.process_flow(f, r) + class FlowWriter: def __init__(self, fo): diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 4a6d3614..6093c936 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -307,7 +307,7 @@ class ClientConnection(controller.Msg): controller.Msg.__init__(self) def get_state(self): - return list(self.address) + return list(self.address) if self.address else None @classmethod def from_state(klass, state): |