From 6a2668d8651b52a688ce5771d9a834473d334a64 Mon Sep 17 00:00:00 2001
From: Shadab Zafar <dufferzafar0@gmail.com>
Date: Fri, 15 Jul 2016 16:55:09 +0530
Subject: Add ~http and ~tcp filters

---
 mitmproxy/filt.py           | 22 +++++++++++++++++++++-
 test/mitmproxy/test_filt.py | 10 ++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/mitmproxy/filt.py b/mitmproxy/filt.py
index e0e5606b..11765dcd 100644
--- a/mitmproxy/filt.py
+++ b/mitmproxy/filt.py
@@ -81,6 +81,24 @@ 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"
@@ -384,7 +402,9 @@ filt_unary = [
     FReq,
     FResp,
     FAsset,
-    FErr
+    FErr,
+    FHTTP,
+    FTCP,
 ]
 filt_rex = [
     FHeadRequest,
diff --git a/test/mitmproxy/test_filt.py b/test/mitmproxy/test_filt.py
index 8330b68c..46e86b57 100644
--- a/test/mitmproxy/test_filt.py
+++ b/test/mitmproxy/test_filt.py
@@ -87,6 +87,11 @@ class TestMatchingHTTPFlow:
     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)
@@ -258,6 +263,11 @@ class TestMatchingTCPFlow:
     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)
-- 
cgit v1.2.3