zmap-mark-ii/test/test_zblocklist.py
2024-02-21 18:23:18 -06:00

139 lines
3.9 KiB
Python

import unittest
import subprocess
import os
import sys
executable_path = None
class ZBlocklistTest(unittest.TestCase):
BLOCKLIST = [
"10.0.0.0/8 # private subnet",
"192.168.0.0/16 # private subnet",
"128.255.0.0/16 # university of iowa",
"141.212.120.0/24 # halderman lab"
]
ALLOWLIST = [
"141.212.0.0/16 # university of michigan",
]
IPS = [
"61.193.80.24",
"195.19.1.6",
"114.34.253.25",
"180.69.174.9",
"38.134.130.203",
"192.168.1.50",
"98.125.221.180",
"197.160.60.150",
"47.139.63.128",
"95.224.78.221",
"170.114.52.252",
"10.0.0.5",
"128.255.134.1",
"141.212.120.10",
"141.212.12.6"
]
IPS_MINUS_BL = [
"61.193.80.24",
"195.19.1.6",
"114.34.253.25",
"180.69.174.9",
"38.134.130.203",
"98.125.221.180",
"197.160.60.150",
"47.139.63.128",
"95.224.78.221",
"170.114.52.252",
"141.212.12.6"
]
WL_IPS = [
"141.212.120.10",
"141.212.12.6"
]
WL_IPS_MINUS_BL = [
"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
with open("/tmp/blocklist", "w") as fd:
for line in self.BLOCKLIST:
fd.write("%s\n" % line)
with open("/tmp/allowlist", "w") as fd:
for line in self.ALLOWLIST:
fd.write("%s\n" % line)
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/blocklist"):
os.remove("/tmp/blocklist")
if os.path.exists("/tmp/allowlist"):
os.remove("/tmp/allowlist")
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, allowlist, blocklist, ipsfile="/tmp/ips", numtimestocat=1):
cmd = "cat"
for _ in range(0, numtimestocat):
cmd += " %s" % ipsfile
cmd += " | %s" % self.path
if allowlist:
cmd = cmd + " -w %s" % allowlist
if blocklist:
cmd = cmd + " -b %s" % blocklist
results = subprocess.check_output(cmd, shell=True)
ips = results.rstrip().split("\n")
return ips
def testValidBlocklist(self):
res = self.execute(None, "/tmp/blocklist")
self.assertEqual(set(res), set(self.IPS_MINUS_BL))
def testValidAllowlist(self):
res = self.execute("/tmp/allowlist", None)
self.assertEqual(set(res), set(self.WL_IPS))
def testValidAllowAndBlockList(self):
res = self.execute("/tmp/allowlist", "/tmp/blocklist")
self.assertEqual(set(res), set(self.WL_IPS_MINUS_BL))
def testDuplicateChecking(self):
res = self.execute(None, "/tmp/blocklist", 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/blocklist", ipsfile="/tmp/ips-commented")
self.assertEqual(set(res), set(self.IPS_MINUS_BL))
if __name__ == "__main__":
if len(sys.argv) != 2:
print("USAGE: %s zblocklist" % sys.argv[0])
sys.exit(1)
executable_path = sys.argv[1]
assert(os.path.exists(executable_path))
unittest.main(argv=sys.argv[:1])