"""Pruebas sobre la red: cada afirmación de los apuntes que se puede comprobar corriéndola."""
import os
import re
import tempfile
import time
import unittest

from netlab import HERE, Lab, require_answers, run, wait_until

LAB = Lab("red", "w12-red", "redes-w12-red")
RESOLVER, ROOT, TLD, AUTH = "10.212.2.10", "10.212.3.10", "10.212.4.10", "10.212.5.10"
WEB = "10.212.6.10"
PROXY = "10.212.7.10:3128"
ATTEMPTS = 6
URL = "http://www.wayra.test"
SERVERS = {"raiz": ROOT, "tld": TLD, "auth": AUTH, "resolver": RESOLVER, "web": WEB}


def setUpModule():
    unittest.addModuleCleanup(LAB.destroy)
    LAB.deploy()
    done = run(["sh", os.path.join(HERE, "red", "start.sh")])
    if done.returncode != 0:
        raise RuntimeError(done.stdout + done.stderr)
    for node in ("raiz", "tld", "auth", "resolver", "web", "cache"):
        if not wait_until(lambda: "en 10.212" in LAB.out(node, "cat", f"/tmp/{node}.log"), 20):
            raise RuntimeError(f"{node} no arrancó: " + LAB.out(node, "cat", f"/tmp/{node}.log"))


def dig(*args):
    return LAB.out("cliente", "dig", "+time=3", "+tries=1", *args)


def answers(*args):
    return LAB.out("cliente", "dig", "+short", "+time=3", "+tries=1", *args).split()


def query_time(*args):
    return int(re.search(r"Query time: (\d+) msec", dig(*args)).group(1))


def log_of(node):
    return LAB.out(node, "cat", f"/tmp/{node}.log")


def restart_resolver():
    """Un resolver nuevo no tiene nada en cache."""
    LAB.sh("resolver", "pkill", "-f", "dns_resolver")
    wait_until(lambda: "dns_resolver" not in LAB.out("resolver", "ps"), 5)
    LAB.daemon("resolver", f"cd /w12 && exec python3 -u dns_resolver.py {RESOLVER} {ROOT} "
                           ">/tmp/resolver.log 2>&1")
    if not wait_until(lambda: "resolver en" in log_of("resolver"), 10):
        raise RuntimeError(log_of("resolver"))


def curl(*args):
    return LAB.out("cliente", "curl", "-s", "--max-time", "10", *args)


def capture(node, name, seconds, expression):
    """Captura en segundo plano, en el nodo. Devuelve cuando tcpdump ya está escuchando."""
    LAB.daemon(node, f"timeout {seconds} tcpdump -i eth1 --immediate-mode -U -w /tmp/{name}.pcap {expression} "
                     f">/tmp/{name}.err 2>&1")
    if not wait_until(lambda: "listening" in LAB.out(node, "cat", f"/tmp/{name}.err"), 5):
        raise RuntimeError("tcpdump no arrancó")


def tshark(node, name, *args):
    """Lee la captura de un nodo con tshark, que corre en el contenedor del curso."""
    with tempfile.TemporaryDirectory() as folder:
        local = os.path.join(folder, "x.pcap")
        copied = run(["docker", "cp", "-q", f"{LAB.prefix}{node}:/tmp/{name}.pcap", local])
        if copied.returncode != 0:
            raise RuntimeError(copied.stderr)
        return run(["tshark", "-r", local, *args]).stdout


class A_Hierarchy(unittest.TestCase):
    def test_trace_goes_root_tld_auth(self):
        out = dig("+trace", "www.wayra.test")
        where = [out.index(f"from {ip}#53") for ip in (ROOT, TLD, AUTH)]
        self.assertEqual(where, sorted(where), out)
        self.assertIn("10.212.6.10", out)

    def test_manual_walk(self):
        root = dig(f"@{ROOT}", "+norecurse", "www.wayra.test")
        self.assertIn("ANSWER: 0", root)
        self.assertIn(TLD, root)
        tld = dig(f"@{TLD}", "+norecurse", "www.wayra.test")
        self.assertIn("ANSWER: 0", tld)
        self.assertIn(AUTH, tld)
        auth = dig(f"@{AUTH}", "+norecurse", "www.wayra.test")
        self.assertRegex(auth, r"flags: qr aa;")
        self.assertIn("ANSWER: 1", auth)

    def test_cold_query_walks_and_second_one_stays_in_the_resolver(self):
        restart_resolver()
        before = {node: log_of(node).count("www.wayra.test.") for node in ("raiz", "tld", "auth")}
        cold = query_time("www.wayra.test")
        after_cold = {node: log_of(node).count("www.wayra.test.") for node in before}
        for node in before:
            self.assertEqual(after_cold[node], before[node] + 1, node)
        warm = min(query_time("www.wayra.test") for _ in range(3))
        self.assertEqual({node: log_of(node).count("www.wayra.test.") for node in before}, after_cold)
        # Tres RTT de 20 ms contra casi nada. El margen es grande porque la máquina es compartida.
        self.assertGreaterEqual(cold, 50)
        self.assertLess(warm, cold / 2)

    def test_after_learning_the_tld_a_new_name_skips_the_root(self):
        restart_resolver()
        answers("www.wayra.test")
        before = log_of("raiz").count("\n")
        self.assertEqual(answers("correo.wayra.test"), ["10.212.6.25"])
        self.assertEqual(log_of("raiz").count("\n"), before)


class B_Records(unittest.TestCase):
    def test_types(self):
        self.assertEqual(answers("www.wayra.test", "A"), ["10.212.6.10"])
        self.assertEqual(answers("www.wayra.test", "AAAA"), ["2001:db8:212:6::10"])
        self.assertEqual(answers("wayra.test", "NS"), ["ns1.wayra.test."])
        self.assertEqual(answers("wayra.test", "MX"), ["10", "correo.wayra.test."])
        self.assertIn("ejemplo", " ".join(answers("wayra.test", "TXT")))

    def test_cname(self):
        self.assertEqual(answers("blog.wayra.test"), ["www.wayra.test.", "10.212.6.10"])

    def test_nxdomain_carries_soa(self):
        out = dig("noexiste.wayra.test")
        self.assertIn("status: NXDOMAIN", out)
        self.assertRegex(out, r"wayra\.test\.\s+\d+\s+IN\s+SOA")

    def test_ttl_goes_down_and_then_the_resolver_asks_again(self):
        restart_resolver()
        ttl = lambda: int(LAB.out("cliente", "dig", "+noall", "+answer", "rapido.wayra.test").split()[1])
        first = ttl()
        asked = log_of("auth").count("rapido.wayra.test.")
        time.sleep(2)
        second = ttl()
        self.assertLess(second, first)
        self.assertEqual(log_of("auth").count("rapido.wayra.test."), asked)
        time.sleep(3.5)
        ttl()
        self.assertEqual(log_of("auth").count("rapido.wayra.test."), asked + 1)

    def test_resolver_without_recursion_only_tells_what_it_has(self):
        restart_resolver()
        out = dig(f"@{RESOLVER}", "+norecurse", "nunca-preguntado.wayra.test")
        self.assertIn("ANSWER: 0", out)
        self.assertNotIn("nunca-preguntado", log_of("auth"))


class C_Transport(unittest.TestCase):
    def test_large_answer_makes_dig_retry_over_tcp(self):
        out = dig(f"@{AUTH}", "grande.wayra.test", "TXT")
        self.assertIn("Truncated, retrying in TCP mode", out)
        self.assertIn("(TCP)", out)
        self.assertIn("ANSWER: 1", out.split("Truncated")[1])

    def test_large_answer_also_works_through_the_resolver(self):
        out = dig("grande.wayra.test", "TXT")
        self.assertIn("Truncated, retrying in TCP mode", out)
        self.assertIn("(TCP)", out)
        self.assertIn("status: NOERROR", out.split("Truncated")[1])
        self.assertRegex(log_of("auth"), r"tcp 10\.212\.2\.10:\d+ id=\d+ grande")

    def test_ignoring_truncation_shows_the_empty_answer(self):
        out = dig(f"@{AUTH}", "+ignore", "grande.wayra.test", "TXT")
        self.assertRegex(out, r"flags: qr aa tc")
        self.assertIn("ANSWER: 0", out)

    def test_capture_matches_dig(self):
        restart_resolver()
        capture("cliente", "dns", 6, "udp port 53")
        out = dig("blog.wayra.test")
        msg_id = int(re.search(r"id: (\d+)", out).group(1))
        time.sleep(0.5)
        LAB.sh("cliente", "pkill", "tcpdump")
        time.sleep(0.5)
        fields = tshark("cliente", "dns", "-T", "fields", "-E", "separator=|", "-e", "dns.id",
                        "-e", "dns.flags.response", "-e", "dns.qry.name", "-e", "dns.a",
                        "-e", "dns.cname")
        rows = [line.split("|") for line in fields.strip().splitlines()]
        queries = [row for row in rows if row[1] in ("0", "False")]
        replies = [row for row in rows if row[1] in ("1", "True")]
        self.assertEqual(len(queries), 1, fields)
        self.assertEqual(len(replies), 1, fields)
        self.assertEqual(int(queries[0][0], 16), msg_id)
        self.assertEqual(int(replies[0][0], 16), msg_id)
        self.assertEqual(replies[0][2:], ["blog.wayra.test", "10.212.6.10", "www.wayra.test"])

    def test_upstream_queries_change_id_and_port(self):
        restart_resolver()
        capture("resolver", "up", 8, f"udp and dst port 53 and dst host {AUTH}")
        answers("www.wayra.test")
        for index in range(6):
            answers(f"azar{index}.wayra.test")
        time.sleep(0.5)
        LAB.sh("resolver", "pkill", "tcpdump")
        time.sleep(0.5)
        fields = tshark("resolver", "up", "-T", "fields", "-E", "separator=|",
                        "-e", "dns.id", "-e", "udp.srcport")
        rows = [line.split("|") for line in fields.strip().splitlines()]
        self.assertEqual(len(rows), 7, fields)
        self.assertGreaterEqual(len({row[0] for row in rows}), 6)
        self.assertGreaterEqual(len({row[1] for row in rows}), 6)

    def test_dns_ask_reads_the_same_as_dig(self):
        out = LAB.out("cliente", "/w12/bin/dns_ask", RESOLVER, "www.wayra.test")
        self.assertIn("rd=1 ra=1 rcode=0", out)
        self.assertRegex(out, r"www\.wayra\.test\. ttl=\d+ A 10\.212\.6\.10")
        out = LAB.out("cliente", "/w12/bin/dns_ask", ROOT, "www.wayra.test", "--norecurse")
        self.assertIn("rd=0 ra=0 rcode=0", out)
        self.assertIn("respuestas=0 autoridad=1 adicionales=1", out)

    def test_dns_ask_reports_a_bad_address(self):
        done = LAB.sh("cliente", "/w12/bin/dns_ask", "no-es-ip", "www.wayra.test")
        self.assertEqual(done.returncode, 2)


def nc(request):
    return LAB.out("cliente", "sh", "-c", f"printf '{request}' | nc -w1 {WEB} 80")


class D_Http(unittest.TestCase):
    def test_by_hand_request(self):
        out = nc(r"GET / HTTP/1.1\r\nHost: www.wayra.test\r\n\r\n")
        self.assertTrue(out.startswith("HTTP/1.1 200 OK"), out)
        self.assertIn("wayra</html>", out)

    def test_host_chooses_the_site(self):
        self.assertIn("el blog", nc(r"GET / HTTP/1.1\r\nHost: blog.wayra.test\r\n\r\n"))

    def test_this_toy_server_accepts_a_missing_host(self):
        # RFC 9112 le exige 400 a un servidor de verdad. Los apuntes muestran la diferencia.
        self.assertTrue(nc(r"GET / HTTP/1.1\r\n\r\n").startswith("HTTP/1.1 200"))

    def test_status_404(self):
        self.assertTrue(nc(r"GET /nada HTTP/1.1\r\nHost: a\r\n\r\n").startswith("HTTP/1.1 404"))

    def test_head_has_headers_but_no_body(self):
        out = curl("-I", f"{URL}/")
        self.assertIn("Content-Length: 67", out)
        self.assertNotIn("<html>", out)

    def test_unknown_method_is_answered_501(self):
        self.assertIn(" 501 ", nc(r"DELETE / HTTP/1.1\r\nHost: a\r\n\r\n").splitlines()[0])

    def test_cookie_carries_the_state(self):
        self.assertIn("visita 1", curl(f"{URL}/visitas"))
        self.assertIn("visita 1", curl(f"{URL}/visitas"))  # sin cookie, no nos recuerda
        self.assertIn("visita 8", curl("-H", "Cookie: visitas=7", f"{URL}/visitas"))
        self.assertIn("visitas=8", curl("-D-", "-H", "Cookie: visitas=7", f"{URL}/visitas"))

    def test_conditional_request(self):
        first = curl("-D-", f"{URL}/")
        self.assertIn("200 OK", first)
        etag = re.search(r"ETag: (\S+)", first).group(1)
        again = curl("-D-", "-H", f"If-None-Match: {etag}", f"{URL}/")
        self.assertIn("304 Not Modified", again)
        self.assertNotIn("<html>", again)
        other = curl("-D-", "-H", 'If-None-Match: "otro"', f"{URL}/")
        self.assertIn("200 OK", other)

    def timing(self, *args):
        fmt = "%{time_connect} %{time_starttransfer} %{time_total} %{num_connects}\n"
        out = curl(*args, "-o", "/dev/null", "-w", fmt)
        return [[float(value) for value in line.split()] for line in out.strip().splitlines()]

    def test_connect_costs_one_rtt_and_first_byte_another(self):
        # El RTT hacia web es de 50 ms: conexión a 50 ms, primer byte a 100 ms.
        best = min(self.timing(f"{URL}/")[0] for _ in range(3))
        connect, first_byte, _total, connections = best
        self.assertEqual(connections, 1)
        self.assertGreaterEqual(connect, 0.045)
        self.assertGreaterEqual(first_byte, 0.090)
        self.assertLess(connect, first_byte)

    def test_persistent_connection_saves_handshakes(self):
        urls = [f"{URL}/{name}" for name in ("", "a.png", "b.png", "c.png")]
        outputs = [arg for _ in urls for arg in ("-o", "/dev/null")]
        persistent, separate = [], []
        for _ in range(3):
            out = curl(*outputs, "-w", "%{num_connects} %{time_total}\n", *urls)
            rows = [line.split() for line in out.strip().splitlines()]
            self.assertEqual([row[0] for row in rows], ["1", "0", "0", "0"], out)
            persistent.append(sum(float(row[1]) for row in rows))
            separate.append(sum(self.timing("-H", "Connection: close", url)[0][2] for url in urls))
        # 5 RTT contra 8 RTT. Se compara lo mejor de tres: una máquina compartida solo agrega tiempo.
        self.assertGreaterEqual(min(persistent), 0.22)
        self.assertGreaterEqual(min(separate), 0.36)
        self.assertGreater(min(separate), min(persistent) * 1.2)

    def test_connection_close_is_honored(self):
        out = curl("-D-", "-o", "/dev/null", "-H", "Connection: close", f"{URL}/")
        self.assertIn("Connection: close", out)


def through_proxy(path):
    out = curl("-D-", "-x", PROXY, f"{URL}{path}")
    state = re.search(r"X-Cache: (\w+)", out)
    return (state.group(1) if state else None), out


def origin_requests():
    return int(LAB.out("cliente", "curl", "-s", f"{URL}/contadores").split()[0])


class E_Proxy(unittest.TestCase):
    def test_fresh_stale_revalidated_changed(self):
        self.assertEqual(through_proxy("/noticia")[0], "MISS")
        before = origin_requests()
        state, out = through_proxy("/noticia")
        self.assertEqual(state, "HIT")
        self.assertIn("Age:", out)
        # La consulta del contador cuenta una petición; el HIT no agrega ninguna.
        self.assertEqual(origin_requests(), before + 1)
        time.sleep(3.2)
        self.assertEqual(through_proxy("/noticia")[0], "REVALIDATED")
        curl(f"{URL}/cambiar")
        self.assertEqual(through_proxy("/noticia")[0], "HIT")  # sigue fresca: dato viejo
        time.sleep(3.2)
        state, out = through_proxy("/noticia")
        self.assertEqual(state, "MISS")
        self.assertRegex(out, r"versión \d+")

    def test_no_store_never_hits(self):
        self.assertEqual(through_proxy("/privado")[0], "BYPASS")
        self.assertEqual(through_proxy("/privado")[0], "BYPASS")

    def test_without_indications_it_is_not_stored(self):
        self.assertEqual(through_proxy("/hora")[0], "BYPASS")

    def test_hit_is_faster_than_going_to_the_origin(self):
        through_proxy("/a.png")
        fmt = "%{time_total}\n"
        hit = min(float(curl("-o", "/dev/null", "-w", fmt, "-x", PROXY, f"{URL}/a.png"))
                  for _ in range(3))
        direct = min(float(curl("-o", "/dev/null", "-w", fmt, f"{URL}/a.png")) for _ in range(3))
        self.assertLess(hit, direct / 2)


def streams_overlap(log):
    """¿Salió un DATA del stream 3 antes del último DATA del stream 1?"""
    data = re.findall(r"-> DATA stream=(\d+)", log)
    if sorted(set(data)) != ["1", "3"]:
        return False
    return data.index("3") < len(data) - 1 - data[::-1].index("1")


class F_Http2(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        LAB.daemon("web", f"cd /w12 && exec python3 -u h2_frames.py {WEB} 8080 >/tmp/h2.log 2>&1")
        if not wait_until(lambda: "h2_frames en" in LAB.out("web", "cat", "/tmp/h2.log"), 10):
            raise RuntimeError(LAB.out("web", "cat", "/tmp/h2.log"))
        capture("cliente", "h2", 15, "tcp port 8080")
        base = f"http://{WEB}:8080"
        # A veces curl abre el segundo stream recién cuando termina el primero, y entonces no hay
        # mezcla que ver: es decisión de curl, no del servidor. Se repite hasta ver dos abiertos.
        for _attempt in range(ATTEMPTS):
            seen = len(LAB.out("web", "cat", "/tmp/h2.log"))
            cls.curl_out = curl("--http2-prior-knowledge", "--parallel", "-o", "/dev/null", "-o", "/dev/null",
                                "-w", "%{http_version} %{http_code} %{num_connects} %{size_download}\n",
                                f"{base}/uno", f"{base}/dos")
            time.sleep(0.3)
            cls.log = LAB.out("web", "cat", "/tmp/h2.log")[seen:]
            if streams_overlap(cls.log):
                break
        time.sleep(0.5)
        LAB.sh("cliente", "pkill", "tcpdump")
        time.sleep(0.5)
        LAB.sh("web", "pkill", "-f", "h2_frames")

    def test_curl_used_http2_over_one_connection(self):
        rows = [line.split() for line in self.curl_out.strip().splitlines()]
        self.assertEqual([row[:2] for row in rows], [["2", "200"]] * 2, self.curl_out)
        self.assertEqual(sum(int(row[2]) for row in rows), 1, self.curl_out)
        self.assertEqual([row[3] for row in rows], ["2000", "2000"], self.curl_out)

    def test_streams_take_turns(self):
        self.assertTrue(streams_overlap(self.log), self.log)

    def test_second_headers_are_much_smaller(self):
        sizes = dict(re.findall(r"<- HEADERS stream=(\d+) largo=(\d+)", self.log))
        self.assertLess(int(sizes["3"]), int(sizes["1"]) / 2, self.log)

    def test_tshark_reads_the_same_frames(self):
        out = tshark("cliente", "h2", "-d", "tcp.port==8080,http2", "-Y", "http2",
                     "-T", "fields", "-e", "http2.headers.path")
        self.assertIn("/uno", out)
        self.assertIn("/dos", out)


class G_Delegation(unittest.TestCase):
    """El laboratorio: delegar killa.test. Se salta hasta que completen las dos zonas."""

    @classmethod
    def setUpClass(cls):
        folder = require_answers("test.zone", "killa.zone")
        LAB.sh("tld", "pkill", "-f", "dns_auth")
        wait_until(lambda: "dns_auth" not in LAB.out("tld", "ps"), 5)
        LAB.daemon("tld", f"cd /w12 && exec python3 -u dns_auth.py {TLD} {folder}/test.zone >/tmp/tld.log 2>&1")
        LAB.daemon("killa", f"cd /w12 && exec python3 -u dns_auth.py 10.212.8.10 {folder}/killa.zone "
                            ">/tmp/killa.log 2>&1")
        for node in ("tld", "killa"):
            if not wait_until(lambda: "autoritativo" in log_of(node), 10):
                raise RuntimeError(log_of(node))
        restart_resolver()

    def test_parent_delegates_with_glue(self):
        out = dig(f"@{TLD}", "+norecurse", "www.killa.test")
        self.assertIn("ANSWER: 0", out)
        self.assertRegex(out, r"killa\.test\.\s+\d+\s+IN\s+NS\s+ns1\.killa\.test\.")
        self.assertRegex(out, r"ns1\.killa\.test\.\s+\d+\s+IN\s+A\s+10\.212\.8\.10")

    def test_child_answers_with_authority(self):
        out = dig("@10.212.8.10", "+norecurse", "www.killa.test")
        self.assertRegex(out, r"flags: qr aa;")
        self.assertRegex(out, r"www\.killa\.test\.\s+120\s+IN\s+A\s+10\.212\.6\.10")

    def test_resolver_finds_every_record(self):
        self.assertEqual(answers("www.killa.test"), ["10.212.6.10"])
        self.assertEqual(answers("luna.killa.test"), ["www.killa.test.", "10.212.6.10"])
        self.assertEqual(answers("killa.test", "MX"), ["20", "correo.killa.test."])
        self.assertEqual(answers("correo.killa.test"), ["10.212.6.25"])
        self.assertEqual(answers("killa.test", "NS"), ["ns1.killa.test."])
        self.assertTrue(answers("killa.test", "TXT"))


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