allow comment chars.

zmap · Nov 13, 2014 · ebd488a · ebd488a
1 parent 7a60e2f
commit ebd488a
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 5 deletions.
diff --git a/src/zblacklist.c b/src/zblacklist.c
@@ -43,6 +43,18 @@
 //	uint32_t duplicates;
 //};
 
+#undef  MIN
+#define MIN(X,Y) ((X) < (Y) ? (X) : (Y))
+
+static inline char* zmin(char *a, char *b) {
+    if (a && !b)
+        return a;
+    else if (b && !a)
+        return b;
+    else
+        return MIN(a,b);
+}
+
 struct zbl_conf {
 	char *blacklist_filename;
 	char *whitelist_filename;
@@ -160,7 +172,11 @@ int main(int argc, char **argv)
 	char line[1000]; 
 	while (fgets(line, sizeof(line), stdin) != NULL) {
 		// remove new line
-		char *n = strchr(line, '\n');
+		char *n = zmin(zmin(zmin(zmin(strchr(line, '\n'), 
+                        strchr(line, ',')), 
+                        strchr(line, '\t')),
+                        strchr(line, ' ')), 
+                        strchr(line, '#'));
 		assert(n);
 		n[0] = 0;
 		log_trace("zblacklist", "input value %s", line);

diff --git a/test/test_zblacklist.py b/test/test_zblacklist.py
@@ -58,6 +58,13 @@ class ZBlacklistTest(unittest.TestCase):
         "141.212.12.6"
     ]
 
+    COMMENT_STRS = [
+        "# some comment here",
+        " # some comment here",
+        ",google.com,data",
+        "\t#some comment here"
+    ] 
+
     def setUp(self):
         global executable_path
         self.path = executable_path
@@ -70,6 +77,10 @@ def setUp(self):
         with open("/tmp/ips", "w") as fd:
             for line in self.IPS:
                 fd.write("%s\n" % line)
+        with open("/tmp/ips-commented", "w") as fd:
+            for line in self.IPS:
+                for comment in self.COMMENT_STRS:
+                    fd.write("%s%s\n" % (line, comment))
 
     def tearDown(self):
         if os.path.exists("/tmp/blacklist"):
@@ -78,11 +89,14 @@ def tearDown(self):
             os.remove("/tmp/whitelist")
         if os.path.exists("/tmp/ips"):
             os.remove("/tmp/ips")
+        if os.path.exists("/tmp/ips-commented"):
+            os.remove("/tmp/ips-commented")
+
 
-    def execute(self, whitelist, blacklist, numtimestocat=1):
+    def execute(self, whitelist, blacklist, ipsfile="/tmp/ips", numtimestocat=1):
         cmd = "cat"
-        for _ in range(0,numtimestocat):
-            cmd += " /tmp/ips"
+        for _ in range(0, numtimestocat):
+            cmd += " %s" % ipsfile
         cmd += " | %s" % self.path
         if whitelist:
             cmd = cmd + " -w %s" % whitelist
@@ -105,11 +119,15 @@ def testValidWhiteAndBlackList(self):
         self.assertEqual(set(res), set(self.WL_IPS_MINUS_BL))
 
     def testDuplicateChecking(self):
-        res = self.execute(None, "/tmp/blacklist", 5)
+        res = self.execute(None, "/tmp/blacklist", numtimestocat=5)
         self.assertEqual(len(res), len(self.IPS_MINUS_BL))
         self.assertEqual(set(res), set(self.IPS_MINUS_BL))
 
+    def testCommentCharacters(self):
+        res = self.execute(None, "/tmp/blacklist", ipsfile="/tmp/ips-commented")
+        self.assertEqual(set(res), set(self.IPS_MINUS_BL))
 
+
 if __name__ == "__main__":
     if len(sys.argv) != 2:
         print "USAGE: %s zblacklist" % sys.argv[0]