From 262a42055344b9ece7b94dc8ed762d8a9a8882de Mon Sep 17 00:00:00 2001 From: Shadab Zafar Date: Fri, 15 Jul 2016 14:35:18 +0530 Subject: 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 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() -- cgit v1.2.3