# Django imports.
from django.test import TestCase

# Application imports.
from conntrackt.models import Entity
from conntrackt import utils


class GenerateEntityIptablesTest(TestCase):

    fixtures = ['test-data.json']

    def test_generated_iptables(self):
        """
        Tests if the entity's iptables are generated properly or not.
        """

        entity = Entity.objects.get(pk=1)
        generated = utils.generate_entity_iptables(entity)

        expected = """*filter
:INPUT DROP [0:0]
# Accept all incoming related traffic.
-A INPUT -m state --state RELATED,ESTABLISHED -j ACCEPT

# Accept all incoming traffic on loopback interface.
-A INPUT -i lo -j ACCEPT

# MySQL.
-A INPUT -s 192.168.1.3/255.255.255.255 -d 192.168.1.1/255.255.255.255 -p tcp -m tcp --dport 3306 -j ACCEPT

# Ping.
-A INPUT -s 192.168.1.2/255.255.255.255 -d 192.168.1.1/255.255.255.255 -p icmp -m icmp --icmp-type 8 -j ACCEPT

# SSH.
-A INPUT -s 192.168.1.2/255.255.255.255 -d 192.168.1.1/255.255.255.255 -p tcp -m tcp --dport 22 -j ACCEPT
-A INPUT -s 192.168.2.0/255.255.255.0 -d 192.168.1.1/255.255.255.255 -p tcp -m tcp --dport 22 -j ACCEPT

:OUTPUT ACCEPT [0:0]
:FORWARD DROP [0:0]
COMMIT
*nat
:PREROUTING ACCEPT [0:0]
:INPUT ACCEPT [0:0]
:OUTPUT ACCEPT [0:0]
:POSTROUTING ACCEPT [0:0]
COMMIT
"""
        self.assertEqual(generated, expected)