"""Pruebas de DHCP y NAT en nat/nat.clab.yml."""
import os
import re
import time
import unittest

from netlab import HERE, Lab, answer_path, wait_until

LAB = Lab("nat", "w04-nat", "redes-w04-nat")
GW_PUBLIC = "100.64.0.2"
EXT = "100.64.9.10"


def setUpModule():
    unittest.addModuleCleanup(LAB.destroy)
    LAB.deploy()
    ok = wait_until(lambda: "default" in LAB.out("gw", "ip", "route")
                    and "100.64.9.0/24" in LAB.out("isp", "ip", "route")
                    and LAB.sh("isp", "ping", "-c", "1", "-W", "1", EXT).returncode == 0)
    if not ok:
        raise RuntimeError("la topología no quedó lista")
    LAB.daemon("ext", "nc -lk -p 9000 -e /bin/cat")


def reset_nat():
    LAB.sh("gw", "nft", "flush", "ruleset")
    LAB.sh("gw", "conntrack", "-F")


class A_Dhcp(unittest.TestCase):
    def test_lease_with_four_messages_by_broadcast(self):
        self.assertEqual(LAB.out("a3", "ip", "-4", "addr", "show", "dev", "eth1").count("inet "), 0)
        LAB.daemon("gw", "python3 /opt/dhcp_server.py eth1 10.200.50.1 "
                         "10.200.50.100 10.200.50.110 > /tmp/dhcp.log 2>&1")
        time.sleep(1)
        at_client = LAB.capture("sw", "eth3", "udp port 67 or udp port 68", 4, seconds=8,
                                flags=("-n", "-e", "-v"))
        at_other = LAB.capture("a2", "eth1", "udp port 67 or udp port 68", 4, seconds=8,
                               flags=("-n", "-e", "-v"))
        done = LAB.sh("a3", "udhcpc", "-i", "eth1", "-n", "-q", "-f", timeout=30)
        self.assertEqual(done.returncode, 0, done.stderr)
        output = at_client.result()
        for name in ("Discover", "Offer", "Request", "ACK"):
            self.assertIn(f"DHCP-Message (53), length 1: {name}", output)
        self.assertEqual(output.count("> ff:ff:ff:ff:ff:ff"), 4)
        self.assertIn("0.0.0.0.68 > 255.255.255.255.67", output)
        # Otro host de la LAN también ve los mensajes: es broadcast.
        self.assertIn("Discover", at_other.result())

        self.assertIn("10.200.50.100/24", LAB.out("a3", "ip", "-4", "addr", "show", "dev", "eth1"))
        self.assertIn("default via 10.200.50.1", LAB.out("a3", "ip", "route"))
        self.assertEqual(LAB.sh("a3", "ping", "-c", "1", "-W", "2", "10.200.50.1").returncode, 0)


class B_Nat(unittest.TestCase):
    def setUp(self):
        reset_nat()
        self.addCleanup(reset_nat)

    def load_masquerade(self):
        LAB.load_nft("gw", os.path.join(HERE, "nat", "masquerade.nft"))

    def test_without_nat_private_source_does_not_come_back(self):
        capture = LAB.capture("gw", "eth2", "icmp", 2, seconds=3, flags=("-n",))
        done = LAB.sh("a1", "ping", "-c", "1", "-W", "1", EXT)
        self.assertNotEqual(done.returncode, 0)
        self.assertIn(f"10.200.50.10 > {EXT}", capture.result())

    def test_masquerade_rewrites_source_on_the_way_out_and_back(self):
        self.load_masquerade()
        inside = LAB.capture("gw", "eth1", "icmp", 2, flags=("-n",))
        outside = LAB.capture("gw", "eth2", "icmp", 2, flags=("-n",))
        at_ext = LAB.capture("ext", "eth1", "icmp", 2, flags=("-n",))
        self.assertEqual(LAB.sh("a1", "ping", "-c", "1", "-W", "2", EXT).returncode, 0)
        self.assertIn(f"10.200.50.10 > {EXT}", inside.result())
        self.assertIn(f"{GW_PUBLIC} > {EXT}", outside.result())
        at_ext = at_ext.result()
        self.assertIn(f"{GW_PUBLIC} > {EXT}", at_ext)
        self.assertNotIn("10.200.50", at_ext)

    def test_same_source_port_from_two_hosts_gets_two_external_ports(self):
        self.load_masquerade()
        syns = LAB.capture("ext", "eth1", "tcp[tcpflags] & tcp-syn != 0 and dst port 9000", 2,
                           seconds=6, flags=("-n",))
        LAB.daemon("a1", f"sleep 8 | nc -p 40000 {EXT} 9000")
        LAB.daemon("a2", f"sleep 8 | nc -p 40000 {EXT} 9000")
        ports = re.findall(rf"{GW_PUBLIC}\.(\d+) > ", syns.result())
        self.assertEqual(len(set(ports)), 2)
        table = LAB.out("gw", "conntrack", "-L", "-p", "tcp")
        self.assertRegex(table, r"src=10\.200\.50\.10 .*sport=40000")
        self.assertRegex(table, r"src=10\.200\.50\.11 .*sport=40000")

    def test_private_address_is_unreachable_from_outside(self):
        self.load_masquerade()
        capture = LAB.capture("ext", "eth1", "icmp", 1, flags=("-n",))
        done = LAB.sh("ext", "curl", "-sS", "--connect-timeout", "2",
                      "http://10.200.50.10:8080/")
        self.assertNotEqual(done.returncode, 0)
        self.assertIn("10.200.50.10 unreachable", capture.result())

    def test_public_address_refuses_until_port_is_redirected(self):
        self.load_masquerade()
        LAB.daemon("a1", "nc -lk -p 8080 -e /bin/echo hola-desde-a1")
        time.sleep(0.5)
        capture = LAB.capture("gw", "eth2", "tcp port 8080", 2, flags=("-n",))
        before = LAB.sh("ext", "curl", "-sS", "--connect-timeout", "2",
                        f"http://{GW_PUBLIC}:8080/")
        self.assertNotEqual(before.returncode, 0)
        self.assertIn("Flags [R.]", capture.result())

        LAB.load_nft("gw", os.path.join(HERE, "nat", "dnat_demo.nft"))
        after = LAB.sh("ext", "sh", "-c", f"echo x | nc -w2 {GW_PUBLIC} 8080")
        self.assertIn("hola-desde-a1", after.stdout)
        table = LAB.out("gw", "conntrack", "-L", "-p", "tcp", "--dport", "8080")
        self.assertRegex(table, rf"dst={GW_PUBLIC} sport=\d+ dport=8080 src=10\.200\.50\.10")

    def test_lab_step_5_port_forward(self):
        path = answer_path("portforward.nft")
        self.load_masquerade()
        LAB.load_nft("gw", path)
        LAB.daemon("a2", "nc -lk -p 8000 -e /bin/echo hola-desde-a2")
        time.sleep(0.5)
        done = LAB.sh("ext", "sh", "-c", f"echo x | nc -w2 {GW_PUBLIC} 2222")
        self.assertIn("hola-desde-a2", done.stdout)
        # a1 no recibe nada por ese puerto, y sigue saliendo con NAT.
        self.assertEqual(LAB.sh("a1", "ping", "-c", "1", "-W", "2", EXT).returncode, 0)


if __name__ == "__main__":
    unittest.main()