From 471e196e085f4161cb0667e0b70a940cb1de67d2 Mon Sep 17 00:00:00 2001
From: isra17 <isra017@gmail.com>
Date: Mon, 13 Jul 2015 17:07:39 -0400
Subject: Add ~src ~dst REGEX filters

This filter allow to match on the request source and destination address
in the form of `<IP>:<Port>`.

Also fixed the parsing grammar to add a `WordEnd` after each filter
name. That way, `~src` doesn't match `~s` instead and keep the behavior
consistent with `~hq` != `~h`.
---
 libmproxy/filt.py | 21 ++++++++++++++++++---
 test/test_filt.py | 17 +++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/libmproxy/filt.py b/libmproxy/filt.py
index 1983586b..bd17a807 100644
--- a/libmproxy/filt.py
+++ b/libmproxy/filt.py
@@ -241,6 +241,19 @@ class FUrl(_Rex):
     def __call__(self, f):
         return re.search(self.expr, f.request.url)
 
+class FSrc(_Rex):
+    code = "src"
+    help = "Match source address"
+
+    def __call__(self, f):
+        return f.client_conn and re.search(self.expr, repr(f.client_conn.address))
+
+class FDst(_Rex):
+    code = "dst"
+    help = "Match destination address"
+
+    def __call__(self, f):
+        return f.server_conn and re.search(self.expr, repr(f.server_conn.address))
 
 class _Int(_Action):
     def __init__(self, num):
@@ -313,6 +326,8 @@ filt_rex = [
     FRequestContentType,
     FResponseContentType,
     FContentType,
+    FSrc,
+    FDst,
 ]
 filt_int = [
     FCode
@@ -324,7 +339,7 @@ def _make():
     # ones.
     parts = []
     for klass in filt_unary:
-        f = pp.Literal("~%s" % klass.code)
+        f = pp.Literal("~%s" % klass.code) + pp.WordEnd()
         f.setParseAction(klass.make)
         parts.append(f)
 
@@ -333,12 +348,12 @@ def _make():
         pp.QuotedString("\"", escChar='\\') |\
         pp.QuotedString("'", escChar='\\')
     for klass in filt_rex:
-        f = pp.Literal("~%s" % klass.code) + rex.copy()
+        f = pp.Literal("~%s" % klass.code) + pp.WordEnd() + rex.copy()
         f.setParseAction(klass.make)
         parts.append(f)
 
     for klass in filt_int:
-        f = pp.Literal("~%s" % klass.code) + pp.Word(pp.nums)
+        f = pp.Literal("~%s" % klass.code) + pp.WordEnd() + pp.Word(pp.nums)
         f.setParseAction(klass.make)
         parts.append(f)
 
diff --git a/test/test_filt.py b/test/test_filt.py
index 3ad17dfe..bcdf6e4c 100644
--- a/test/test_filt.py
+++ b/test/test_filt.py
@@ -241,6 +241,23 @@ class TestMatching:
         assert self.q("~c 200", s)
         assert not self.q("~c 201", s)
 
+    def test_src(self):
+        q = self.req()
+        assert self.q("~src address", q)
+        assert not self.q("~src foobar", q)
+        assert self.q("~src :22", q)
+        assert not self.q("~src :99", q)
+        assert self.q("~src address:22", q)
+
+    def test_dst(self):
+        q = self.req()
+        q.server_conn = tutils.tserver_conn()
+        assert self.q("~dst address", q)
+        assert not self.q("~dst foobar", q)
+        assert self.q("~dst :22", q)
+        assert not self.q("~dst :99", q)
+        assert self.q("~dst address:22", q)
+
     def test_and(self):
         s = self.resp()
         assert self.q("~c 200 & ~h head", s)
-- 
cgit v1.2.3