diff options
| -rw-r--r-- | mitmproxy/filt.py | 34 | ||||
| -rw-r--r-- | test/mitmproxy/test_filt.py | 16 | 
2 files changed, 31 insertions, 19 deletions
| diff --git a/mitmproxy/filt.py b/mitmproxy/filt.py index 00971218..e0e5606b 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 2d2b87e6..8330b68c 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() | 
