def test_firewall_block_packet(self):
"""Verify firewall blocks a packet that doesn't match a rule."""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(
direction="inbound", protocol="tcp", port="80",
ip_address="192.168.1.2"
)
)
self.assertFalse(
fw.accept_packet(
direction="outbound", protocol="tcp", port=80,
ip_address="192.168.1.2"
)
)
self.assertFalse(
fw.accept_packet(
direction="inbound", protocol="udp", port=80,
ip_address="192.168.1.2"
)
)
self.assertFalse(
fw.accept_packet(
direction="inbound", protocol="udp", port=81,
ip_address="192.168.1.2"
)
)
self.assertFalse(
fw.accept_packet(
direction="outbound", protocol="tcp", port=80,
ip_address="192.168.1.3"
)
)
def test_firewall_allow_range_ipaddr_packet(self):
"""
Verify firewall allows a packet that matches a rule with ranged IP
addresses.
"""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(
direction="inbound", protocol="tcp", port="80",
ip_address="0.0.0.0-255.255.255.255"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=80,
ip_address="0.0.0.0"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=80,
ip_address="255.255.255.255"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=80,
ip_address="192.168.1.2"
)
)
def test_firewall_allow_range_port_packet(self):
"""
Verify firewall allows a packet that matches a rule with ranged port
numbers.
"""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(
direction="inbound", protocol="tcp", port="1-65535",
ip_address="192.168.1.2"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=1,
ip_address="192.168.1.2"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=65535,
ip_address="192.168.1.2"
)
)
self.assertTrue(
fw.accept_packet(
direction="inbound", protocol="tcp", port=30000,
ip_address="192.168.1.2"
)
)
def __init__(self, csv_file_path: Optional[str] = None):
"""
Initialize the firewall by reading and storing the firewall rules of
the CSV file.
"""
# initialize the data structure to store firewall rules
num_buckets = 64
self.num_ports_bucket = 65536 // num_buckets
self.fw_rules = {
"inbound": {
"tcp": [set() for i in range(num_buckets)],
"udp": [set() for i in range(num_buckets)],
},
"outbound": {
"tcp": [set() for i in range(num_buckets)],
"udp": [set() for i in range(num_buckets)],
},
}
# read firewall rules from CSV file and add them to the data structure
if csv_file_path:
with open(csv_file_path, "r") as csv_file:
csv_reader = csv.reader(csv_file)
for csv_fw_rule in csv_reader:
fw_rule = FirewallRule(*csv_fw_rule)
self.add_fw_rule(fw_rule)
def test_no_add_duplicate_ipaddr_rules(self):
"""Verify that duplicate range IP address rules cannot be added."""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(direction="inbound",
protocol="tcp",
port="80",
ip_address="192.168.1.2-192.168.2.2"))
bucket_num = 80 // fw.num_ports_bucket
self.assertEqual(len(fw.fw_rules["inbound"]["tcp"][bucket_num]), 1)
fw.add_fw_rule(
FirewallRule(direction="inbound",
protocol="tcp",
port="80",
ip_address="192.168.1.2-192.168.2.2"))
self.assertEqual(len(fw.fw_rules["inbound"]["tcp"][bucket_num]), 1)
def test_no_add_duplicate_range_port_rules(self):
"""Verify that duplicate range port rules cannot be added."""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(
direction="inbound", protocol="tcp", port="50-100",
ip_address="192.168.1.2"
)
)
self.assertEqual(len(fw.fw_rules), 1)
fw.add_fw_rule(
FirewallRule(
direction="inbound", protocol="tcp", port="50-100",
ip_address="192.168.1.2"
)
)
self.assertEqual(len(fw.fw_rules), 1)
def test_no_add_duplicate_range_port_rules(self):
"""Verify that duplicate range port rules cannot be added."""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(direction="inbound",
protocol="tcp",
port="50-2000",
ip_address="192.168.1.2"))
start_bucket = 50 // fw.num_ports_bucket
end_bucket = 2000 // fw.num_ports_bucket
for bucket_num in range(start_bucket, end_bucket + 1):
self.assertEqual(len(fw.fw_rules["inbound"]["tcp"][bucket_num]), 1)
fw.add_fw_rule(
FirewallRule(direction="inbound",
protocol="tcp",
port="50-2000",
ip_address="192.168.1.2"))
for bucket_num in range(start_bucket, end_bucket + 1):
self.assertEqual(len(fw.fw_rules["inbound"]["tcp"][bucket_num]), 1)
def __init__(self, csv_file_path: Optional[str] = None):
"""
Initialize the firewall by reading and storing the firewall rules of
the CSV file.
"""
self.fw_rules = set()
if csv_file_path:
with open(csv_file_path, "r") as csv_file:
csv_reader = csv.reader(csv_file)
for csv_fw_rule in csv_reader:
self.add_fw_rule(FirewallRule(*csv_fw_rule))
def test_firewall_block_range_ipaddr_packet(self):
"""
Verify firewall blocks a packet that doesn't match a rule with ranged
IP addresses.
"""
fw = Firewall()
fw.add_fw_rule(
FirewallRule(direction="inbound",
protocol="tcp",
port="80",
ip_address="192.168.1.2-192.168.2.1"))
self.assertFalse(
fw.accept_packet(direction="inbound",
protocol="tcp",
port=80,
ip_address="192.168.1.1"))
self.assertFalse(
fw.accept_packet(direction="inbound",
protocol="tcp",
port=91,
ip_address="192.168.2.2"))