diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/console.py | 20 | ||||
| -rw-r--r-- | libmproxy/dump.py | 14 | ||||
| -rw-r--r-- | libmproxy/flow.py | 73 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 43 |
4 files changed, 71 insertions, 79 deletions
diff --git a/libmproxy/console.py b/libmproxy/console.py index 895974d2..815eebdf 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -66,7 +66,7 @@ def format_flow(f, focus, extended=False, padding=2): f.request.url(), ), ] - if f.response or f.error or f.is_replay(): + if f.response or f.error or f.request.is_replay(): tsr = f.response or f.error if extended and tsr: ts = ("highlight", utils.format_timestamp(tsr.timestamp) + " ") @@ -77,7 +77,7 @@ def format_flow(f, focus, extended=False, padding=2): txt.append(("text", ts)) txt.append(" "*(padding+2)) met = "" - if f.is_replay(): + if f.request.is_replay(): txt.append(("method", "[replay] ")) elif f.modified(): txt.append(("method", "[edited] ")) @@ -715,17 +715,13 @@ class ConsoleState(flow.State): self.last_script = "" self.last_saveload = "" - def add_browserconnect(self, f): - flow.State.add_browserconnect(self, f) + def add_request(self, req): + f = flow.State.add_request(self, req) if self.focus is None: self.set_focus(0) else: self.set_focus(self.focus + 1) - - def add_request(self, req): - if self.focus is None: - self.set_focus(0) - return flow.State.add_request(self, req) + return f def add_response(self, resp): if self.store is not None: @@ -1305,7 +1301,7 @@ class ConsoleMaster(flow.FlowMaster): def process_flow(self, f, r): if f.match(self.state.beep): urwid.curses_display.curses.beep() - if f.match(self.state.intercept) and not f.is_replay(): + if f.match(self.state.intercept) and not f.request.is_replay(): f.intercept() else: r.ack() @@ -1313,8 +1309,8 @@ class ConsoleMaster(flow.FlowMaster): self.refresh_connection(f) # Handlers - def handle_clientconnection(self, r): - f = flow.FlowMaster.handle_clientconnection(self, r) + def handle_clientconnect(self, r): + f = flow.FlowMaster.handle_clientconnect(self, r) if f: self.sync_list_view() diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 372e6ef6..f6a7ae7e 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -38,14 +38,6 @@ class DumpMaster(flow.FlowMaster): except IOError, v: raise DumpError(v.strerror) - def handle_clientconnection(self, r): - flow.FlowMaster.handle_clientconnection(self, r) - r.ack() - - def handle_error(self, r): - flow.FlowMaster.handle_error(self, r) - r.ack() - def _runscript(self, f, script): try: ret = f.run_script(script) @@ -80,12 +72,12 @@ class DumpMaster(flow.FlowMaster): return sz = utils.pretty_size(len(f.response.content)) if self.o.verbosity == 1: - print >> self.outfile, f.client_conn.address[0], + print >> self.outfile, f.request.client_conn.address[0], print >> self.outfile, f.request.short() print >> self.outfile, " <<", print >> self.outfile, f.response.short(), sz elif self.o.verbosity == 2: - print >> self.outfile, f.client_conn.address[0], + print >> self.outfile, f.request.client_conn.address[0], print >> self.outfile, f.request.short() print >> self.outfile, self.indent(4, f.request.headers) print >> self.outfile @@ -93,7 +85,7 @@ class DumpMaster(flow.FlowMaster): print >> self.outfile, self.indent(4, f.response.headers) print >> self.outfile, "\n" elif self.o.verbosity == 3: - print >> self.outfile, f.client_conn.address[0], + print >> self.outfile, f.request.client_conn.address[0], print >> self.outfile, f.request.short() print >> self.outfile, self.indent(4, f.request.headers) if utils.isBin(f.request.content): diff --git a/libmproxy/flow.py b/libmproxy/flow.py index cea0ca1c..d9df7a1a 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -32,9 +32,9 @@ class ReplayThread(threading.Thread): class Flow: - def __init__(self, client_conn): - self.client_conn = client_conn - self.request, self.response, self.error = None, None, None + def __init__(self, request): + self.request = request + self.response, self.error = None, None self.intercepting = False self._backup = None @@ -90,7 +90,6 @@ class Flow: request = self.request.get_state() if self.request else None, response = self.response.get_state() if self.response else None, error = self.error.get_state() if self.error else None, - client_conn = self.client_conn.get_state() ) if nobackup: d["backup"] = None @@ -99,10 +98,8 @@ class Flow: return d def load_state(self, state): - self.client_conn = proxy.ClientConnection.from_state(state["client_conn"]) self._backup = state["backup"] - if state["request"]: - self.request = proxy.Request.from_state(self.client_conn, state["request"]) + self.request = proxy.Request.from_state(state["request"]) if state["response"]: self.response = proxy.Response.from_state(self.request, state["response"]) if state["error"]: @@ -141,9 +138,6 @@ class Flow: return pattern(self.request) return False - def is_replay(self): - return self.client_conn.is_replay() - def kill(self): if self.request and not self.request.acked: self.request.ack(None) @@ -165,35 +159,43 @@ class Flow: class State: def __init__(self): + self.client_connections = [] self.flow_map = {} self.flow_list = [] + # These are compiled filt expressions: self.limit = None self.intercept = None - def add_browserconnect(self, f): + def clientconnect(self, cc): + if not isinstance(cc, proxy.ClientConnect): + assert False + self.client_connections.append(cc) + + def clientdisconnect(self, dc): """ Start a browser connection. """ - self.flow_list.insert(0, f) - self.flow_map[f.client_conn] = f + self.client_connections.remove(dc.client_conn) def add_request(self, req): """ Add a request to the state. Returns the matching flow. """ - f = self.flow_map.get(req.client_conn) - if not f: - f = Flow(req.client_conn) - self.add_browserconnect(f) - f.request = req + if not isinstance(req, proxy.Request): + assert False + f = Flow(req) + self.flow_list.insert(0, f) + self.flow_map[req] = f return f def add_response(self, resp): """ Add a response to the state. Returns the matching flow. """ - f = self.flow_map.get(resp.request.client_conn) + if not isinstance(resp, proxy.Response): + assert False + f = self.flow_map.get(resp.request) if not f: return False f.response = resp @@ -204,7 +206,7 @@ class State: Add an error response to the state. Returns the matching flow, or None if there isn't one. """ - f = self.flow_map.get(err.client_conn) + f = self.flow_map.get(err.flow.request) if not f: return None f.error = err @@ -213,7 +215,7 @@ class State: def load_flows(self, flows): self.flow_list.extend(flows) for i in flows: - self.flow_map[i.client_conn] = i + self.flow_map[i.request] = i def set_limit(self, limit): """ @@ -229,27 +231,17 @@ class State: return tuple(self.flow_list[:]) def get_client_conn(self, itm): - if isinstance(itm, proxy.ClientConnection): + if isinstance(itm, proxy.ClientConnect): return itm elif hasattr(itm, "client_conn"): return itm.client_conn elif hasattr(itm, "request"): return itm.request.client_conn - def lookup(self, itm): - """ - Checks for matching client_conn, using a Flow, Replay Connection, - ClientConnection, Request, Response or Error object. Returns None - if not found. - """ - client_conn = self.get_client_conn(itm) - return self.flow_map.get(client_conn) - def delete_flow(self, f): if not f.intercepting: - c = self.get_client_conn(f) - if c in self.flow_map: - del self.flow_map[c] + if f.request in self.flow_map: + del self.flow_map[f.request] self.flow_list.remove(f) return True return False @@ -280,7 +272,7 @@ class State: if f.request: f.backup() conn = self.get_client_conn(f) - f.client_conn.set_replay() + f.request.set_replay() if f.request.content: f.request.headers["content-length"] = [str(len(f.request.content))] f.response = None @@ -295,12 +287,13 @@ class FlowMaster(controller.Master): controller.Master.__init__(self, server) self.state = state - # Handlers - def handle_clientconnection(self, r): - f = Flow(r) - self.state.add_browserconnect(f) + def handle_clientconnect(self, r): + self.state.clientconnect(r) + r.ack() + + def handle_clientdisconnect(self, r): + self.state.clientdisconnect(r) r.ack() - return f def handle_error(self, r): f = self.state.add_error(r) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 54ff2ec3..1c4d4d71 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -136,11 +136,21 @@ class Request(controller.Msg): self.close = False controller.Msg.__init__(self) + def set_replay(self): + self.client_conn = None + + def is_replay(self): + if self.client_conn: + return False + else: + return True + def is_cached(self): return False def get_state(self): return dict( + client_conn = self.client_conn.get_state(), host = self.host, port = self.port, scheme = self.scheme, @@ -152,9 +162,9 @@ class Request(controller.Msg): ) @classmethod - def from_state(klass, client_conn, state): + def from_state(klass, state): return klass( - client_conn, + ClientConnect.from_state(state["client_conn"]), state["host"], state["port"], state["scheme"], @@ -165,6 +175,9 @@ class Request(controller.Msg): state["timestamp"] ) + def __hash__(self): + return id(self) + def __eq__(self, other): return self.get_state() == other.get_state() @@ -296,7 +309,13 @@ class Response(controller.Msg): return self.FMT%data -class ClientConnection(controller.Msg): +class ClientDisconnect(controller.Msg): + def __init__(self, client_conn): + controller.Msg.__init__(self) + self.client_conn = client_conn + + +class ClientConnect(controller.Msg): def __init__(self, address): """ address is an (address, port) tuple, or None if this connection has @@ -313,22 +332,13 @@ class ClientConnection(controller.Msg): def from_state(klass, state): return klass(state) - def set_replay(self): - self.address = None - - def is_replay(self): - if self.address: - return False - else: - return True - def copy(self): return copy.copy(self) class Error(controller.Msg): - def __init__(self, client_conn, msg, timestamp=None): - self.client_conn, self.msg = client_conn, msg + def __init__(self, flow, msg, timestamp=None): + self.flow, self.msg = flow, msg self.timestamp = timestamp or time.time() controller.Msg.__init__(self) @@ -453,11 +463,12 @@ class ProxyHandler(SocketServer.StreamRequestHandler): SocketServer.StreamRequestHandler.__init__(self, request, client_address, server) def handle(self): - cc = ClientConnection(self.client_address) + cc = ClientConnect(self.client_address) cc.send(self.mqueue) while not cc.close: self.handle_request(cc) - cc = cc.copy() + cd = ClientDisconnect(cc) + cd.send(self.mqueue) self.finish() def handle_request(self, cc): |