Merge remote-tracking branch 'duffer/tcpflow-filter'

5 files changed, 342 insertions, 32 deletions

diff --git a/mitmproxy/filt.py b/mitmproxy/filt.py
index a42988f1..2774f8d1 100644
--- a/mitmproxy/filt.py
+++ b/mitmproxy/filt.py
@@ -35,11 +35,26 @@ from __future__ import absolute_import, print_function, division
 
 import re
 import sys
+import functools
+
+from mitmproxy.models.http import HTTPFlow
+from mitmproxy.models.tcp import TCPFlow
 from netlib import strutils
 
 import pyparsing as pp
 
 
+def only(*types):
+    def decorator(fn):
+        @functools.wraps(fn)
+        def filter_types(self, flow):
+            if isinstance(flow, types):
+                return fn(self, flow)
+            return False
+        return filter_types
+    return decorator
+
+
 class _Token(object):
 
     def dump(self, indent=0, fp=sys.stdout):
@@ -65,10 +80,29 @@ class FErr(_Action):
         return True if f.error else False
 
 
+class FHTTP(_Action):
+    code = "http"
+    help = "Match HTTP flows"
+
+    @only(HTTPFlow)
+    def __call__(self, f):
+        return True
+
+
+class FTCP(_Action):
+    code = "tcp"
+    help = "Match TCP flows"
+
+    @only(TCPFlow)
+    def __call__(self, f):
+        return True
+
+
 class FReq(_Action):
     code = "q"
     help = "Match request with no response"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if not f.response:
             return True
@@ -78,6 +112,7 @@ class FResp(_Action):
     code = "s"
     help = "Match response"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         return bool(f.response)
 
@@ -117,6 +152,7 @@ class FAsset(_Action):
     ]
     ASSET_TYPES = [re.compile(x) for x in ASSET_TYPES]
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.response:
             for i in self.ASSET_TYPES:
@@ -129,6 +165,7 @@ class FContentType(_Rex):
     code = "t"
     help = "Content-type header"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if _check_content_type(self.re, f.request):
             return True
@@ -137,18 +174,20 @@ class FContentType(_Rex):
         return False
 
 
-class FRequestContentType(_Rex):
+class FContentTypeRequest(_Rex):
     code = "tq"
     help = "Request Content-Type header"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         return _check_content_type(self.re, f.request)
 
 
-class FResponseContentType(_Rex):
+class FContentTypeResponse(_Rex):
     code = "ts"
     help = "Response Content-Type header"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.response:
             return _check_content_type(self.re, f.response)
@@ -160,6 +199,7 @@ class FHead(_Rex):
     help = "Header"
     flags = re.MULTILINE
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.request and self.re.search(bytes(f.request.headers)):
             return True
@@ -173,6 +213,7 @@ class FHeadRequest(_Rex):
     help = "Request header"
     flags = re.MULTILINE
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.request and self.re.search(bytes(f.request.headers)):
             return True
@@ -183,6 +224,7 @@ class FHeadResponse(_Rex):
     help = "Response header"
     flags = re.MULTILINE
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.response and self.re.search(bytes(f.response.headers)):
             return True
@@ -192,13 +234,19 @@ class FBod(_Rex):
     code = "b"
     help = "Body"
 
+    @only(HTTPFlow, TCPFlow)
     def __call__(self, f):
-        if f.request and f.request.raw_content:
-            if self.re.search(f.request.get_content(strict=False)):
-                return True
-        if f.response and f.response.raw_content:
-            if self.re.search(f.response.get_content(strict=False)):
-                return True
+        if isinstance(f, HTTPFlow):
+            if f.request and f.request.raw_content:
+                if self.re.search(f.request.get_content(strict=False)):
+                    return True
+            if f.response and f.response.raw_content:
+                if self.re.search(f.request.get_content(strict=False)):
+                    return True
+        elif isinstance(f, TCPFlow):
+            for msg in f.messages:
+                if self.re.search(msg.content):
+                    return True
         return False
 
 
@@ -206,20 +254,32 @@ class FBodRequest(_Rex):
     code = "bq"
     help = "Request body"
 
+    @only(HTTPFlow, TCPFlow)
     def __call__(self, f):
-        if f.request and f.request.raw_content:
-            if self.re.search(f.request.get_content(strict=False)):
-                return True
+        if isinstance(f, HTTPFlow):
+            if f.request and f.request.raw_content:
+                if self.re.search(f.request.get_content(strict=False)):
+                    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, TCPFlow)
     def __call__(self, f):
-        if f.response and f.response.raw_content:
-            if self.re.search(f.response.get_content(strict=False)):
-                return True
+        if isinstance(f, HTTPFlow):
+            if f.response and f.response.raw_content:
+                if self.re.search(f.response.get_content(strict=False)):
+                    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):
@@ -227,6 +287,7 @@ class FMethod(_Rex):
     help = "Method"
     flags = re.IGNORECASE
 
+    @only(HTTPFlow)
     def __call__(self, f):
         return bool(self.re.search(f.request.data.method))
 
@@ -236,6 +297,7 @@ class FDomain(_Rex):
     help = "Domain"
     flags = re.IGNORECASE
 
+    @only(HTTPFlow)
     def __call__(self, f):
         return bool(self.re.search(f.request.data.host))
 
@@ -252,6 +314,7 @@ class FUrl(_Rex):
             toks = toks[1:]
         return klass(*toks)
 
+    @only(HTTPFlow)
     def __call__(self, f):
         return self.re.search(f.request.url)
 
@@ -284,6 +347,7 @@ class FCode(_Int):
     code = "c"
     help = "HTTP response code"
 
+    @only(HTTPFlow)
     def __call__(self, f):
         if f.response and f.response.status_code == self.num:
             return True
@@ -331,26 +395,28 @@ class FNot(_Token):
 
 
 filt_unary = [
+    FAsset,
+    FErr,
+    FHTTP,
     FReq,
     FResp,
-    FAsset,
-    FErr
+    FTCP,
 ]
 filt_rex = [
-    FHeadRequest,
-    FHeadResponse,
-    FHead,
+    FBod,
     FBodRequest,
     FBodResponse,
-    FBod,
-    FMethod,
-    FDomain,
-    FUrl,
-    FRequestContentType,
-    FResponseContentType,
     FContentType,
-    FSrc,
+    FContentTypeRequest,
+    FContentTypeResponse,
+    FDomain,
     FDst,
+    FHead,
+    FHeadRequest,
+    FHeadResponse,
+    FMethod,
+    FSrc,
+    FUrl,
 ]
 filt_int = [
     FCode
diff --git a/mitmproxy/models/tcp.py b/mitmproxy/models/tcp.py
index e33475c2..6650141d 100644
--- a/mitmproxy/models/tcp.py
+++ b/mitmproxy/models/tcp.py
@@ -7,6 +7,8 @@ from typing import List
 import netlib.basetypes
 from mitmproxy.models.flow import Flow
 
+import six
+
 
 class TCPMessage(netlib.basetypes.Serializable):
 
@@ -53,3 +55,22 @@ class TCPFlow(Flow):
 
     def __repr__(self):
         return "<TCPFlow ({} messages)>".format(len(self.messages))
+
+    def match(self, f):
+        """
+            Match this flow against a compiled filter expression. Returns True
+            if matched, False if not.
+
+            If f is a string, it will be compiled as a filter expression. If
+            the expression is invalid, ValueError is raised.
+        """
+        if isinstance(f, six.string_types):
+            from .. import filt
+
+            f = filt.parse(f)
+            if not f:
+                raise ValueError("Invalid filter expression.")
+        if f:
+            return f(self)
+
+        return True
diff --git a/test/mitmproxy/test_filt.py b/test/mitmproxy/test_filt.py
index 9fe36b2a..69f042bb 100644
--- a/test/mitmproxy/test_filt.py
+++ b/test/mitmproxy/test_filt.py
@@ -1,6 +1,8 @@
 from six.moves import cStringIO as StringIO
-from mitmproxy import filt
 from mock import patch
+
+from mitmproxy import filt
+
 from . import tutils
 
 
@@ -73,7 +75,7 @@ class TestParsing:
         self._dump(a)
 
 
-class TestMatching:
+class TestMatchingHTTPFlow:
 
     def req(self):
         return tutils.tflow()
@@ -87,6 +89,11 @@ class TestMatching:
     def q(self, q, o):
         return filt.parse(q)(o)
 
+    def test_http(self):
+        s = self.req()
+        assert self.q("~http", s)
+        assert not self.q("~tcp", s)
+
     def test_asset(self):
         s = self.resp()
         assert not self.q("~a", s)
@@ -247,6 +254,186 @@ class TestMatching:
         assert not self.q("!~c 201 !~c 200", s)
 
 
+class TestMatchingTCPFlow:
+
+    def flow(self):
+        return tutils.ttcpflow()
+
+    def err(self):
+        return tutils.ttcpflow(err=True)
+
+    def q(self, q, o):
+        return filt.parse(q)(o)
+
+    def test_tcp(self):
+        f = self.flow()
+        assert self.q("~tcp", f)
+        assert not self.q("~http", f)
+
+    def test_ferr(self):
+        e = self.err()
+        assert self.q("~e", e)
+
+    def test_body(self):
+        f = self.flow()
+
+        # 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)
+
+        # Messages sent by client
+        assert self.q("~bq hello", f)
+        assert not self.q("~bq me", f)
+        assert not self.q("~bq nonexistent", 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()
+        assert self.q("~src address", f)
+        assert not self.q("~src foobar", f)
+        assert self.q("~src :22", f)
+        assert not self.q("~src :99", f)
+        assert self.q("~src address:22", f)
+
+    def test_dst(self):
+        f = self.flow()
+        f.server_conn = tutils.tserver_conn()
+        assert self.q("~dst address", f)
+        assert not self.q("~dst foobar", f)
+        assert self.q("~dst :22", f)
+        assert not self.q("~dst :99", f)
+        assert self.q("~dst address:22", f)
+
+    def test_and(self):
+        f = self.flow()
+        f.server_conn = tutils.tserver_conn()
+        assert self.q("~b hello & ~b me", f)
+        assert not self.q("~src wrongaddress & ~b hello", f)
+        assert self.q("(~src :22 & ~dst :22) & ~b hello", f)
+        assert not self.q("(~src address:22 & ~dst :22) & ~b nonexistent", f)
+        assert not self.q("(~src address:22 & ~dst :99) & ~b hello", f)
+
+    def test_or(self):
+        f = self.flow()
+        f.server_conn = tutils.tserver_conn()
+        assert self.q("~b hello | ~b me", f)
+        assert self.q("~src :22 | ~b me", f)
+        assert not self.q("~src :99 | ~dst :99", f)
+        assert self.q("(~src :22 | ~dst :22) | ~b me", f)
+
+    def test_not(self):
+        f = self.flow()
+        assert not self.q("! ~src :22", f)
+        assert self.q("! ~src :99", f)
+        assert self.q("!~src :99 !~src :99", f)
+        assert not self.q("!~src :99 !~src :22", f)
+
+    def test_request(self):
+        f = self.flow()
+        assert not self.q("~q", f)
+
+    def test_response(self):
+        f = self.flow()
+        assert not self.q("~s", f)
+
+    def test_headers(self):
+        f = self.flow()
+        assert not self.q("~h whatever", f)
+
+        # Request headers
+        assert not self.q("~hq whatever", f)
+
+        # Response headers
+        assert not self.q("~hs whatever", f)
+
+    def test_content_type(self):
+        f = self.flow()
+        assert not self.q("~t whatever", f)
+
+        # Request content-type
+        assert not self.q("~tq whatever", f)
+
+        # Response content-type
+        assert not self.q("~ts whatever", f)
+
+    def test_code(self):
+        f = self.flow()
+        assert not self.q("~c 200", f)
+
+    def test_domain(self):
+        f = self.flow()
+        assert not self.q("~d whatever", f)
+
+    def test_method(self):
+        f = self.flow()
+        assert not self.q("~m whatever", f)
+
+    def test_url(self):
+        f = self.flow()
+        assert not self.q("~u whatever", f)
+
+
+class TestMatchingDummyFlow:
+
+    def flow(self):
+        return tutils.tdummyflow()
+
+    def err(self):
+        return tutils.tdummyflow(err=True)
+
+    def q(self, q, o):
+        return filt.parse(q)(o)
+
+    def test_filters(self):
+        e = self.err()
+        f = self.flow()
+        f.server_conn = tutils.tserver_conn()
+
+        assert not self.q("~a", f)
+
+        assert not self.q("~b whatever", f)
+        assert not self.q("~bq whatever", f)
+        assert not self.q("~bs whatever", f)
+
+        assert not self.q("~c 0", f)
+
+        assert not self.q("~d whatever", f)
+
+        assert self.q("~dst address", f)
+        assert not self.q("~dst nonexistent", f)
+
+        assert self.q("~e", e)
+        assert not self.q("~e", f)
+
+        assert not self.q("~http", f)
+
+        assert not self.q("~h whatever", f)
+        assert not self.q("~hq whatever", f)
+        assert not self.q("~hs whatever", f)
+
+        assert not self.q("~m whatever", f)
+
+        assert not self.q("~s", f)
+
+        assert self.q("~src address", f)
+        assert not self.q("~src nonexistent", f)
+
+        assert not self.q("~tcp", f)
+
+        assert not self.q("~t whatever", f)
+        assert not self.q("~tq whatever", f)
+        assert not self.q("~ts whatever", f)
+
+        assert not self.q("~u whatever", f)
+
+        assert not self.q("~q", f)
+
+
 @patch('traceback.extract_tb')
 def test_pyparsing_bug(extract_tb):
     """https://github.com/mitmproxy/mitmproxy/issues/1087"""
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index f73616d1..10163401 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -293,7 +293,7 @@ class TestServerPlaybackState:
         assert s._hash(r) == s._hash(r2)
 
 
-class TestFlow(object):
+class TestHTTPFlow(object):
 
     def test_copy(self):
         f = tutils.tflow(resp=True)
@@ -443,6 +443,20 @@ class TestFlow(object):
         assert f.response.raw_content == b"abarb"
 
 
+class TestTCPFlow:
+
+    def test_match(self):
+        f = tutils.ttcpflow()
+        assert not f.match("~b nonexistent")
+        assert f.match(None)
+        assert not f.match("~b nonexistent")
+
+        f = tutils.ttcpflow(err=True)
+        assert f.match("~e")
+
+        tutils.raises(ValueError, f.match, "~")
+
+
 class TestState:
 
     def test_backup(self):
diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py
index d0a09035..d743a9e6 100644
--- a/test/mitmproxy/tutils.py
+++ b/test/mitmproxy/tutils.py
@@ -4,18 +4,19 @@ import tempfile
 import argparse
 import sys
 
-from mitmproxy.models.tcp import TCPMessage
-from six.moves import cStringIO as StringIO
 from contextlib import contextmanager
-
 from unittest.case import SkipTest
 
+from six.moves import cStringIO as StringIO
+
 import netlib.utils
 import netlib.tutils
 from mitmproxy import controller
 from mitmproxy.models import (
     ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow, TCPFlow
 )
+from mitmproxy.models.tcp import TCPMessage
+from mitmproxy.models.flow import Flow
 
 
 def _skip_windows(*args):
@@ -47,6 +48,27 @@ def skip_appveyor(fn):
         return fn
 
 
+class DummyFlow(Flow):
+    """A flow that is neither HTTP nor TCP."""
+
+    def __init__(self, client_conn, server_conn, live=None):
+        super(DummyFlow, self).__init__("dummy", client_conn, server_conn, live)
+
+
+def tdummyflow(client_conn=True, server_conn=True, err=None):
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if err is True:
+        err = terr()
+
+    f = DummyFlow(client_conn, server_conn)
+    f.error = err
+    f.reply = controller.DummyReply()
+    return f
+
+
 def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
     if client_conn is True:
         client_conn = tclient_conn()