From 08f410caccef68f1c41a1ec197d188732c85698d Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 16 Mar 2012 17:13:11 +1300 Subject: Add a hooks mechanism, based on filter expressions. --- libmproxy/filt.py | 1 + libmproxy/flow.py | 47 ++++++++++++++++++++++++++++++++++++++++++++++- libmproxy/tnetstring.py | 14 +++++++------- test/test_flow.py | 35 +++++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 8 deletions(-) diff --git a/libmproxy/filt.py b/libmproxy/filt.py index b49c856e..0e0a82f1 100644 --- a/libmproxy/filt.py +++ b/libmproxy/filt.py @@ -87,6 +87,7 @@ class _Rex(_Action): except: raise ValueError, "Cannot compile expression." + def _check_content_type(expr, o): val = o.headers["content-type"] if val and re.search(expr, val[0]): diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 1ae2bb8a..4c6f2915 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -26,6 +26,48 @@ import controller, version HDR_FORM_URLENCODED = "application/x-www-form-urlencoded" +class Hooks: + def __init__(self): + self.lst = [] + + def add(self, patt, func): + """ + Add a hook. + + patt: A string specifying a filter pattern. + func: A callable taking the matching flow as argument. + + Returns True if hook was added, False if the pattern could not be + parsed. + """ + cpatt = filt.parse(patt) + if not cpatt: + return False + self.lst.append((patt, func, cpatt)) + return True + + def remove(self, patt, func=None): + """ + Remove a hook. + + patt: A string specifying a filter pattern. + func: Optional callable. If not specified, all hooks matching patt are removed. + """ + for i in range(len(self.lst)-1, -1, -1): + if func and (patt, func) == self.lst[i][:2]: + del self.lst[i] + elif not func and patt == self.lst[i][0]: + del self.lst[i] + + def run(self, f): + for _, func, cpatt in self.lst: + if cpatt(f): + func(f) + + def clear(self): + self.lst = [] + + class ScriptContext: def __init__(self, master): self._master = master @@ -75,7 +117,6 @@ class ODict: def __getitem__(self, k): """ Returns a list of values matching key. - """ ret = [] k = self._kconv(k) @@ -1229,6 +1270,7 @@ class FlowMaster(controller.Master): self.anticache = False self.anticomp = False self.refresh_server_playback = False + self.hooks = Hooks() def add_event(self, e, level="info"): """ @@ -1438,6 +1480,7 @@ class FlowMaster(controller.Master): def handle_error(self, r): f = self.state.add_error(r) + self.hooks.run(f) if f: self.run_script_hook("error", f) if self.client_playback: @@ -1447,12 +1490,14 @@ class FlowMaster(controller.Master): def handle_request(self, r): f = self.state.add_request(r) + self.hooks.run(f) self.run_script_hook("request", f) self.process_new_request(f) return f def handle_response(self, r): f = self.state.add_response(r) + self.hooks.run(f) if f: self.run_script_hook("response", f) if self.client_playback: diff --git a/libmproxy/tnetstring.py b/libmproxy/tnetstring.py index 76b15a02..d9be19bc 100644 --- a/libmproxy/tnetstring.py +++ b/libmproxy/tnetstring.py @@ -57,12 +57,12 @@ functions. They're only here so you can use load() to read precisely one item from a file or socket without consuming any extra data. By default tnetstrings work only with byte strings, not unicode. If you want -unicode strings then pass an optional encoding to the various functions, +unicode strings then pass an optional encoding to the various functions, like so:: >>> print repr(tnetstring.loads("2:\\xce\\xb1,")) '\\xce\\xb1' - >>> + >>> >>> print repr(tnetstring.loads("2:\\xce\\xb1,","utf8")) u'\u03b1' @@ -129,7 +129,7 @@ def _rdumpq(q,size,value,encoding=None): write("5:false!") return size + 8 if isinstance(value,(int,long)): - data = str(value) + data = str(value) ldata = len(data) span = str(ldata) write("#") @@ -142,7 +142,7 @@ def _rdumpq(q,size,value,encoding=None): # It round-trips more accurately. # Probably unnecessary in later python versions that # use David Gay's ftoa routines. - data = repr(value) + data = repr(value) ldata = len(data) span = str(ldata) write("^") @@ -207,13 +207,13 @@ def _gdumps(value,encoding): elif value is False: yield "5:false!" elif isinstance(value,(int,long)): - data = str(value) + data = str(value) yield str(len(data)) yield ":" yield data yield "#" elif isinstance(value,(float,)): - data = repr(value) + data = repr(value) yield str(len(data)) yield ":" yield data @@ -334,7 +334,7 @@ def load(file,encoding=None): d[key] = val return d raise ValueError("unknown type tag") - + def pop(string,encoding=None): diff --git a/test/test_flow.py b/test/test_flow.py index ff35f899..8f7551c7 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -1036,9 +1036,44 @@ class udecoded(libpry.AutoTree): assert r.content == "foo" +class uHooks(libpry.AutoTree): + def test_add_remove(self): + f = lambda(x): None + h = flow.Hooks() + h.add("~q", f) + assert h.lst + + h.remove("~q", f) + assert not h.lst + + h.add("~q", f) + h.add("~s", f) + assert len(h.lst) == 2 + h.remove("~q", f) + assert len(h.lst) == 1 + h.remove("~q") + assert len(h.lst) == 1 + h.remove("~s") + assert len(h.lst) == 0 + + track = [] + def func(x): + track.append(x) + + h.add("~s", func) + + f = tutils.tflow() + h.run(f) + assert not track + + f = tutils.tflow_full() + h.run(f) + assert len(track) == 1 + tests = [ + uHooks(), uStickyCookieState(), uStickyAuthState(), uServerPlaybackState(), -- cgit v1.2.3