Refactor filter matching.

2 files changed, 51 insertions, 81 deletions

diff --git a/libmproxy/filt.py b/libmproxy/filt.py
index 228a2283..b9d6f83a 100644
--- a/libmproxy/filt.py
+++ b/libmproxy/filt.py
@@ -60,22 +60,22 @@ class _Action(_Token):
 class FErr(_Action):
     code = "e"
     help = "Match error"
-    def __call__(self, conn):
-        return isinstance(conn, flow.Error)
+    def __call__(self, f):
+        return True if f.error else False
 
 
 class FReq(_Action):
     code = "q"
     help = "Match request"
-    def __call__(self, conn):
-        return isinstance(conn, flow.Request)
+    def __call__(self, f):
+        return True if f.request else False
 
 
 class FResp(_Action):
     code = "s"
     help = "Match response"
-    def __call__(self, conn):
-        return isinstance(conn, flow.Response)
+    def __call__(self, f):
+        return True if f.response else False
 
 
 class _Rex(_Action):
@@ -96,72 +96,64 @@ def _check_content_type(expr, o):
 class FContentType(_Rex):
     code = "t"
     help = "Content-type header"
-    def __call__(self, o):
-        if _check_content_type(self.expr, o):
+    def __call__(self, f):
+        if _check_content_type(self.expr, f.request):
             return True
-        elif isinstance(o, flow.Response) and _check_content_type(self.expr, o.request):
+        elif f.response and _check_content_type(self.expr, f.response):
             return True
-        else:
-            return False
+        return False
 
 
 class FRequestContentType(_Rex):
     code = "tq"
     help = "Request Content-Type header"
-    def __call__(self, o):
-        if isinstance(o, flow.Response):
-            return _check_content_type(self.expr, o.request)
-        else:
-            return _check_content_type(self.expr, o)
+    def __call__(self, f):
+        return _check_content_type(self.expr, f.request)
 
 
 class FResponseContentType(_Rex):
     code = "ts"
     help = "Request Content-Type header"
-    def __call__(self, o):
-        if isinstance(o, flow.Response):
-            return _check_content_type(self.expr, o)
-        else:
-            return False
+    def __call__(self, f):
+        if f.response:
+            return _check_content_type(self.expr, f.response)
+        return False
 
 
 class FHead(_Rex):
     code = "h"
     help = "Header"
-    def __call__(self, o):
-        val = o.headers.match_re(self.expr)
-        if not val and isinstance(o, flow.Response):
-            val = o.request.headers.match_re(self.expr)
-        return val
+    def __call__(self, f):
+        if f.request.headers.match_re(self.expr):
+            return True
+        elif f.response and f.response.headers.match_re(self.expr):
+            return True
+        return False
 
 
 class FHeadRequest(_Rex):
     code = "hq"
     help = "Request header"
-    def __call__(self, o):
-        if isinstance(o, flow.Response):
-            h = o.request.headers
-        else:
-            h = o.headers
-        return h.match_re(self.expr)
+    def __call__(self, f):
+        if f.request.headers.match_re(self.expr):
+            return True
 
 
 class FHeadResponse(_Rex):
     code = "hs"
     help = "Response header"
-    def __call__(self, o):
-        if not isinstance(o, flow.Response):
-            return False
-        return o.headers.match_re(self.expr)
+    def __call__(self, f):
+        if f.response and f.response.headers.match_re(self.expr):
+            return True
 
 
 class FBod(_Rex):
     code = "b"
     help = "Body"
-    def __call__(self, o):
-        if o.content and re.search(self.expr, o.content):
+    def __call__(self, f):
+        if f.request.content and re.search(self.expr, f.request.content):
             return True
-        elif isinstance(o, flow.Response) and o.request.content and re.search(self.expr, o.request.content):
+        elif f.response and f.response.content and re.search(self.expr, f.response.content):
             return True
         return False
 
@@ -169,32 +161,24 @@ class FBod(_Rex):
 class FBodRequest(_Rex):
     code = "bq"
     help = "Request body"
-    def __call__(self, o):
-        if isinstance(o, flow.Response) and o.request.content and re.search(self.expr, o.request.content):
+    def __call__(self, f):
+        if f.request.content and re.search(self.expr, f.request.content):
             return True
-        elif not isinstance(o, flow.Response) and o.content and re.search(self.expr, o.content):
-            return True
-        return False
 
 
 class FBodResponse(_Rex):
     code = "bs"
     help = "Response body"
-    def __call__(self, o):
-        if not isinstance(o, flow.Response):
-            return False
-        elif o.content and re.search(self.expr, o.content):
+    def __call__(self, f):
+        if f.response and f.response.content and re.search(self.expr, f.response.content):
             return True
-        return False
 
 
 class FMethod(_Rex):
     code = "m"
     help = "Method"
-    def __call__(self, o):
-        if isinstance(o, flow.Request):
-            return re.search(self.expr, o.method, re.IGNORECASE)
-        return False
+    def __call__(self, f):
+        return bool(re.search(self.expr, f.request.method, re.IGNORECASE))
 
 
 class FUrl(_Rex):
@@ -207,14 +191,8 @@ class FUrl(_Rex):
             toks = toks[1:]
         return klass(*toks)
 
-    def __call__(self, o):
-        if isinstance(o, flow.Response):
-            c = o.request
-        elif isinstance(o, flow.Request):
-            c = o
-        else:
-            return False
-        return re.search(self.expr, c.get_url())
+    def __call__(self, f):
+        return re.search(self.expr, f.request.get_url())
 
 
 class _Int(_Action):
@@ -225,10 +203,9 @@ class _Int(_Action):
 class FCode(_Int):
     code = "c"
     help = "HTTP response code"
-    def __call__(self, o):
-        if isinstance(o, flow.Response):
-            return o.code == self.num
-        return False
+    def __call__(self, f):
+        if f.response and f.response.code == self.num:
+            return True
 
 
 class FAnd(_Token):
@@ -240,8 +217,8 @@ class FAnd(_Token):
         for i in self.lst:
             i.dump(indent+1, fp)
 
-    def __call__(self, o):
-        return all(i(o) for i in self.lst)
+    def __call__(self, f):
+        return all(i(f) for i in self.lst)
 
 
 class FOr(_Token):
@@ -253,8 +230,8 @@ class FOr(_Token):
         for i in self.lst:
             i.dump(indent+1, fp)
 
-    def __call__(self, o):
-        return any(i(o) for i in self.lst)
+    def __call__(self, f):
+        return any(i(f) for i in self.lst)
 
 
 class FNot(_Token):
@@ -265,8 +242,9 @@ class FNot(_Token):
         print >> fp, "\t"*indent, self.__class__.__name__
         self.itm.dump(indent + 1, fp)
 
-    def __call__(self, o):
-        return not self.itm(o)
+    def __call__(self, f):
+        return not self.itm(f)
+
 
 filt_unary = [
     FReq,
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index b40daefd..e010883a 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -979,16 +979,8 @@ class Flow:
             if matched, False if not.
         """
         if f:
-            if self.response and f(self.response):
-                return True
-            elif self.request and f(self.request):
-                return True
-            elif self.error and f(self.error):
-                return True
-            else:
-                return False
-        else:
-            return True
+            return f(self)
+        return True
 
     def kill(self, master):
         """