diff --git a/libmproxy/console.py b/libmproxy/console.py index 4c51e786b..12f6dec93 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -44,6 +44,55 @@ def format_keyvals(lst, key="key", val="text", space=5, indent=0): return ret +def format_flow(flow, focus, padding=3): + if not flow.request and not flow.response: + txt = [ + ("title", " Connection from %s..."%(flow.connection.address)), + ] + else: + txt = [ + ("ack", "!") if flow.intercepting and not flow.request.acked else " ", + ("method", flow.request.method), + " ", + ( + "text" if (flow.response or flow.error) else "title", + flow.request.url(), + ), + ] + if flow.response or flow.error or flow.is_replay(): + txt.append("\n" + " "*(padding+2)) + if flow.is_replay(): + txt.append(("method", "[replay] ")) + if not (flow.response or flow.error): + txt.append(("text", "waiting for response...")) + + if flow.response: + txt.append( + ("ack", "!") if flow.intercepting and not flow.response.acked else " " + ) + txt.append("<- ") + if flow.response.code in [200, 304]: + txt.append(("goodcode", str(flow.response.code))) + else: + txt.append(("error", str(flow.response.code))) + t = flow.response.headers.get("content-type") + if t: + t = t[0].split(";")[0] + txt.append(("text", " %s"%t)) + if flow.response.content: + txt.append(", %s"%utils.pretty_size(len(flow.response.content))) + elif flow.error: + txt.append( + ("error", flow.error.msg) + ) + if focus: + txt.insert(0, ("focus", ">>" + " "*(padding-2))) + else: + txt.insert(0, " "*padding) + return txt + + + #begin nocover def int_version(v): @@ -67,8 +116,9 @@ class WWrap(urwid.WidgetWrap): class ConnectionItem(WWrap): - def __init__(self, master, state, flow): + def __init__(self, master, state, flow, focus): self.master, self.state, self.flow = master, state, flow + self.focus = focus w = self.get_text() WWrap.__init__(self, w) @@ -77,7 +127,7 @@ class ConnectionItem(WWrap): self.w = self.get_text() def get_text(self): - return urwid.Text(self.flow.get_text()) + return urwid.Text(format_flow(self.flow, self.focus)) def selectable(self): return True @@ -123,7 +173,7 @@ class ConnectionListView(urwid.ListWalker): def get_focus(self): f, i = self.state.get_focus() - f = ConnectionItem(self.master, self.state, f) if f else None + f = ConnectionItem(self.master, self.state, f, True) if f else None return f, i def set_focus(self, focus): @@ -133,23 +183,24 @@ class ConnectionListView(urwid.ListWalker): def get_next(self, pos): f, i = self.state.get_next(pos) - f = ConnectionItem(self.master, self.state, f) if f else None + f = ConnectionItem(self.master, self.state, f, False) if f else None return f, i def get_prev(self, pos): f, i = self.state.get_prev(pos) - f = ConnectionItem(self.master, self.state, f) if f else None + f = ConnectionItem(self.master, self.state, f, False) if f else None return f, i class ConnectionViewHeader(WWrap): - def __init__(self, flow): - self.flow = flow - self.w = urwid.Text(flow.get_text(nofocus=True, padding=0)) + def __init__(self, master, flow): + self.master, self.flow = master, flow + self.w = urwid.Text(format_flow(flow, False, padding=0)) - def refresh_connection(self, f): + def refresh_connection(self, flow): if f == self.flow: - self.w = urwid.Text(f.get_text(nofocus=True, padding=0)) + self.w = urwid.Text(format_flow(flow, False, padding=0)) + VIEW_BODY_RAW = 0 VIEW_BODY_BINARY = 1 @@ -602,59 +653,6 @@ class StatusBar(WWrap): #end nocover -class ConsoleFlow(flow.Flow): - def __init__(self, connection): - flow.Flow.__init__(self, connection) - self.focus = False - - def get_text(self, nofocus=False, padding=3): - if not self.request and not self.response: - txt = [ - ("title", " Connection from %s..."%(self.connection.address)), - ] - else: - txt = [ - ("ack", "!") if self.intercepting and not self.request.acked else " ", - ("method", self.request.method), - " ", - ( - "text" if (self.response or self.error) else "title", - self.request.url(), - ), - ] - if self.response or self.error or self.is_replay(): - txt.append("\n" + " "*(padding+2)) - if self.is_replay(): - txt.append(("method", "[replay] ")) - if not (self.response or self.error): - txt.append(("text", "waiting for response...")) - - if self.response: - txt.append( - ("ack", "!") if self.intercepting and not self.response.acked else " " - ) - txt.append("<- ") - if self.response.code in [200, 304]: - txt.append(("goodcode", str(self.response.code))) - else: - txt.append(("error", str(self.response.code))) - t = self.response.headers.get("content-type") - if t: - t = t[0].split(";")[0] - txt.append(("text", " %s"%t)) - if self.response.content: - txt.append(", %s"%utils.pretty_size(len(self.response.content))) - elif self.error: - txt.append( - ("error", self.error.msg) - ) - if self.focus and not nofocus: - txt.insert(0, ("focus", ">>" + " "*(padding-2))) - else: - txt.insert(0, " "*padding) - return txt - - class ConsoleState(flow.State): def __init__(self): flow.State.__init__(self) @@ -694,13 +692,10 @@ class ConsoleState(flow.State): def set_focus(self, idx): if self.view: - for i in self.view: - i.focus = False if idx >= len(self.view): idx = len(self.view) - 1 elif idx < 0: idx = 0 - self.view[idx].focus = True self.focus = idx def get_from_pos(self, pos): @@ -817,7 +812,6 @@ class ConsoleMaster(controller.Master): sys.stderr.flush() self.shutdown() - def make_view(self): self.view = urwid.Frame( self.body, @@ -850,7 +844,7 @@ class ConsoleMaster(controller.Master): def view_flow(self, flow): self.statusbar = StatusBar(self, self.footer_text_connview) self.body = ConnectionView(self, self.state, flow) - self.header = ConnectionViewHeader(flow) + self.header = ConnectionViewHeader(self, flow) self.viewstate = VIEW_FLOW self.currentflow = flow self.make_view() @@ -903,7 +897,7 @@ class ConsoleMaster(controller.Master): f.close() except IOError, v: return v.strerror - self.state.load_flows(data, ConsoleFlow) + self.state.load_flows(data) if self.conn_list_view: self.conn_list_view.set_focus(0) self.sync_list_view() @@ -1209,7 +1203,7 @@ class ConsoleMaster(controller.Master): # Handlers def handle_browserconnection(self, r): - f = ConsoleFlow(r) + f = flow.Flow(r) self.state.add_browserconnect(f) r.ack() self.sync_list_view() diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 7f4394b1f..a014f8cb5 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -33,6 +33,12 @@ class Flow: self.intercepting = False self._backup = None + def run_script(self): + """ + Run a script on a flow, returning the modified flow. + """ + pass + def dump(self): data = dict( flows = [self.get_state()] @@ -160,9 +166,9 @@ class State: ) return bson.dumps(data) - def load_flows(self, js, klass): + def load_flows(self, js): data = bson.loads(js) - data = [klass.from_state(i) for i in data["flows"]] + data = [Flow.from_state(i) for i in data["flows"]] self.flow_list.extend(data) def set_limit(self, limit): diff --git a/test/test_console.py b/test/test_console.py index cfafed964..c2791bb71 100644 --- a/test/test_console.py +++ b/test/test_console.py @@ -19,14 +19,14 @@ def tresp(req=None): def tflow(): bc = proxy.BrowserConnection("address", 22) - return console.ConsoleFlow(bc) + return flow.Flow(bc) class uState(libpry.AutoTree): def test_backup(self): bc = proxy.BrowserConnection("address", 22) c = console.ConsoleState() - f = console.ConsoleFlow(bc) + f = flow.Flow(bc) c.add_browserconnect(f) f.backup() @@ -40,7 +40,7 @@ class uState(libpry.AutoTree): """ bc = proxy.BrowserConnection("address", 22) c = console.ConsoleState() - f = console.ConsoleFlow(bc) + f = flow.Flow(bc) c.add_browserconnect(f) assert c.lookup(bc) assert c.get_focus() == (f, 0) @@ -66,7 +66,7 @@ class uState(libpry.AutoTree): def test_err(self): bc = proxy.BrowserConnection("address", 22) c = console.ConsoleState() - f = console.ConsoleFlow(bc) + f = flow.Flow(bc) c.add_browserconnect(f) e = proxy.Error(bc, "message") assert c.add_error(e) @@ -104,7 +104,7 @@ class uState(libpry.AutoTree): c = console.ConsoleState() bc = proxy.BrowserConnection("address", 22) - f = console.ConsoleFlow(bc) + f = flow.Flow(bc) c.add_browserconnect(f) assert c.get_focus() == (f, 0) assert c.get_from_pos(0) == (f, 0) @@ -112,7 +112,7 @@ class uState(libpry.AutoTree): assert c.get_next(0) == (None, None) bc2 = proxy.BrowserConnection("address", 22) - f2 = console.ConsoleFlow(bc2) + f2 = flow.Flow(bc2) c.add_browserconnect(f2) assert c.get_focus() == (f, 1) assert c.get_next(0) == (f, 1) @@ -193,8 +193,8 @@ class uState(libpry.AutoTree): dump = c.dump_flows() c.clear() - c.load_flows(dump, console.ConsoleFlow) - assert isinstance(c.flow_list[0], console.ConsoleFlow) + c.load_flows(dump) + assert isinstance(c.flow_list[0], flow.Flow) class uFlow(libpry.AutoTree): @@ -212,36 +212,44 @@ class uFlow(libpry.AutoTree): def test_getset_state(self): f = tflow() state = f.get_state() - assert f == console.ConsoleFlow.from_state(state) + assert f == flow.Flow.from_state(state) f.response = tresp() f.request = f.response.request state = f.get_state() - assert f == console.ConsoleFlow.from_state(state) + assert f == flow.Flow.from_state(state) def test_simple(self): f = tflow() - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.request = treq() - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.response = tresp() f.response.headers["content-type"] = ["text/html"] - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.response.code = 404 - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.focus = True - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.connection = flow.ReplayConnection() - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.response = None - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) f.error = proxy.Error(200, "test") - assert f.get_text() + assert console.format_flow(f, True) + assert console.format_flow(f, False) def test_kill(self): f = tflow() @@ -274,7 +282,7 @@ class uFlow(libpry.AutoTree): assert f.response.acked def test_serialization(self): - f = console.ConsoleFlow(None) + f = flow.Flow(None) f.request = treq()