From 262a42055344b9ece7b94dc8ed762d8a9a8882de Mon Sep 17 00:00:00 2001
From: Shadab Zafar <dufferzafar0@gmail.com>
Date: Fri, 15 Jul 2016 14:35:18 +0530
Subject: [PATCH] Use ~bq / ~bs for messages sent from client / server

---
 mitmproxy/filt.py           | 34 ++++++++++++++++++++--------------
 test/mitmproxy/test_filt.py | 16 +++++++++++-----
 2 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/mitmproxy/filt.py b/mitmproxy/filt.py
index 009712189..e0e5606b8 100644
--- a/mitmproxy/filt.py
+++ b/mitmproxy/filt.py
@@ -219,18 +219,14 @@ class FBod(_Rex):
 
     @only(HTTPFlow, TCPFlow)
     def __call__(self, f):
-
-        # HTTPFlow
-        if hasattr(f, 'request'):
+        if isinstance(f, HTTPFlow):
             if f.request and f.request.content:
                 if self.re.search(f.request.get_decoded_content()):
                     return True
             if f.response and f.response.content:
                 if self.re.search(f.response.get_decoded_content()):
                     return True
-
-        # TCPFlow
-        elif hasattr(f, 'messages'):
+        elif isinstance(f, TCPFlow):
             for msg in f.messages:
                 if self.re.search(msg.content):
                     return True
@@ -242,22 +238,32 @@ class FBodRequest(_Rex):
     code = "bq"
     help = "Request body"
 
-    @only(HTTPFlow)
+    @only(HTTPFlow, TCPFlow)
     def __call__(self, f):
-        if f.request and f.request.content:
-            if self.re.search(f.request.get_decoded_content()):
-                return True
+        if isinstance(f, HTTPFlow):
+            if f.request and f.request.content:
+                if self.re.search(f.request.get_decoded_content()):
+                    return True
+        elif isinstance(f, TCPFlow):
+            for msg in f.messages:
+                if msg.from_client and self.re.search(msg.content):
+                    return True
 
 
 class FBodResponse(_Rex):
     code = "bs"
     help = "Response body"
 
-    @only(HTTPFlow)
+    @only(HTTPFlow, TCPFlow)
     def __call__(self, f):
-        if f.response and f.response.content:
-            if self.re.search(f.response.get_decoded_content()):
-                return True
+        if isinstance(f, HTTPFlow):
+            if f.response and f.response.content:
+                if self.re.search(f.response.get_decoded_content()):
+                    return True
+        elif isinstance(f, TCPFlow):
+            for msg in f.messages:
+                if not msg.from_client and self.re.search(msg.content):
+                    return True
 
 
 class FMethod(_Rex):
diff --git a/test/mitmproxy/test_filt.py b/test/mitmproxy/test_filt.py
index 2d2b87e6f..8330b68c5 100644
--- a/test/mitmproxy/test_filt.py
+++ b/test/mitmproxy/test_filt.py
@@ -264,15 +264,21 @@ class TestMatchingTCPFlow:
 
     def test_body(self):
         f = self.flow()
-        assert not self.q("~b nonexistent", f)
+
+        # Messages sent by client or server
         assert self.q("~b hello", f)
         assert self.q("~b me", f)
+        assert not self.q("~b nonexistent", f)
 
-        # Request Body
-        assert not self.q("~bq whatever", f)
+        # Messages sent by client
+        assert self.q("~bq hello", f)
+        assert not self.q("~bq me", f)
+        assert not self.q("~bq nonexistent", f)
 
-        # Response Body
-        assert not self.q("~bs whatever", f)
+        # Messages sent by server
+        assert self.q("~bs me", f)
+        assert not self.q("~bs hello", f)
+        assert not self.q("~bs nonexistent", f)
 
     def test_src(self):
         f = self.flow()