#include <unity.h>

#include <string>

#include "ipv4.h"
#include "net_probe.h"

using namespace roro::net;

void setUp() {}
void tearDown() {}

void test_what_was_typed() {
    PingArgs p;
    TEST_ASSERT_EQUAL_STRING("", parsePing("example.org", p).c_str());
    TEST_ASSERT_EQUAL_STRING("example.org", p.host.c_str());
    TEST_ASSERT_EQUAL_INT(4, p.count);
    TEST_ASSERT_EQUAL_INT(56, p.size);
    TEST_ASSERT_EQUAL_STRING("", parsePing("  10.9.0.1   10  1372 ", p).c_str());
    TEST_ASSERT_EQUAL_STRING("10.9.0.1", p.host.c_str());
    TEST_ASSERT_EQUAL_INT(10, p.count);
    TEST_ASSERT_EQUAL_INT(1372, p.size);
    TEST_ASSERT_EQUAL_STRING("ping <host> [count] [size]", parsePing("", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a count from 1 to 100", parsePing("a.example 0", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a count from 1 to 100", parsePing("a.example lots", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a size from 0 to 1400 bytes", parsePing("a.example 4 9000", p).c_str());
    TEST_ASSERT_EQUAL_INT(10, p.count);  // a refused line changes nothing

    PortArgs t;
    TEST_ASSERT_EQUAL_STRING("", parsePort("irc.libera.chat 6697", t).c_str());
    TEST_ASSERT_EQUAL_STRING("irc.libera.chat", t.host.c_str());
    TEST_ASSERT_EQUAL_UINT16(6697, t.port);
    TEST_ASSERT_EQUAL_STRING("", parsePort("10.9.0.1:2323", t).c_str());
    TEST_ASSERT_EQUAL_UINT16(2323, t.port);
    TEST_ASSERT_EQUAL_STRING("port <host> <port>", parsePort("10.9.0.1", t).c_str());
    TEST_ASSERT_EQUAL_STRING("a port from 1 to 65535", parsePort("10.9.0.1 70000", t).c_str());

    LookupArgs l;
    TEST_ASSERT_EQUAL_STRING("", parseLookup("roro9stack.net", l).c_str());
    TEST_ASSERT_EQUAL_STRING("", l.server.c_str());
    TEST_ASSERT_EQUAL_STRING("", parseLookup("roro9stack.net 9.9.9.9", l).c_str());
    TEST_ASSERT_EQUAL_STRING("9.9.9.9", l.server.c_str());
    TEST_ASSERT_EQUAL_STRING("the server as an address: 9.9.9.9", parseLookup("roro9stack.net dns.quad9.net", l).c_str());
}

void test_an_echo_request_and_its_answers() {
    uint8_t echo[64];
    size_t len = buildEcho(echo, sizeof echo, 0xBEEF, 7, 32);
    TEST_ASSERT_EQUAL_size_t(40, len);
    TEST_ASSERT_EQUAL_UINT8(8, echo[0]);
    TEST_ASSERT_EQUAL_UINT16(0, inetChecksum(echo, len));  // a packet with its checksum in sums to nothing
    TEST_ASSERT_EQUAL_size_t(0, buildEcho(echo, 20, 1, 1, 32));

    // The reply, as a raw socket hands it: an IP header, then the same with type 0.
    uint8_t reply[20 + 40] = {0x45, 0, 0, 60, 0, 0, 0, 0, 64, 1};
    std::copy(echo, echo + len, reply + 20);
    reply[20] = 0;
    IcmpAnswer a = parseIcmp(reply, sizeof reply);
    TEST_ASSERT_TRUE(a.kind == IcmpAnswer::Kind::Echo);
    TEST_ASSERT_EQUAL_HEX16(0xBEEF, a.id);
    TEST_ASSERT_EQUAL_UINT16(7, a.seq);

    // A router on the way: "time exceeded", with the start of our packet inside it.
    uint8_t exceeded[20 + 8 + 20 + 8] = {0x45, 0, 0, 56, 0, 0, 0, 0, 250, 1};
    exceeded[20] = 11;
    uint8_t* inner = exceeded + 28;
    inner[0] = 0x45;
    inner[9] = 1;
    std::copy(echo, echo + 8, inner + 20);
    a = parseIcmp(exceeded, sizeof exceeded);
    TEST_ASSERT_TRUE(a.kind == IcmpAnswer::Kind::TimeExceeded);
    TEST_ASSERT_EQUAL_HEX16(0xBEEF, a.id);
    TEST_ASSERT_EQUAL_UINT16(7, a.seq);
    exceeded[20] = 3;
    TEST_ASSERT_TRUE(parseIcmp(exceeded, sizeof exceeded).kind == IcmpAnswer::Kind::Unreachable);

    // Not ours to read: someone else's ping to us, a cut packet, UDP.
    TEST_ASSERT_TRUE(parseIcmp(reply, 22).kind == IcmpAnswer::Kind::Other);
    reply[20] = 8;
    TEST_ASSERT_TRUE(parseIcmp(reply, sizeof reply).kind == IcmpAnswer::Kind::Other);
    reply[20] = 0;
    reply[9] = 17;
    TEST_ASSERT_TRUE(parseIcmp(reply, sizeof reply).kind == IcmpAnswer::Kind::Other);
    TEST_ASSERT_TRUE(parseIcmp(exceeded, 40).kind == IcmpAnswer::Kind::Other);
}

void test_the_summary() {
    PingStats s;
    s.sent = 4;
    TEST_ASSERT_EQUAL_STRING("0/4 back, 100% lost", s.summary().c_str());
    s.add(23);
    s.add(12);
    s.add(41);
    TEST_ASSERT_EQUAL_STRING("3/4 back, 25% lost, 12/25/41 ms", s.summary().c_str());
}

void test_a_dns_query() {
    uint8_t q[64];
    size_t len = buildDnsQuery(q, sizeof q, 0x1234, "roro9stack.net");
    const uint8_t want[] = {0x12, 0x34, 0x01, 0x00, 0, 1, 0, 0, 0, 0, 0, 0, 10, 'r', 'o', 'r', 'o', '9', 's', 't', 'a', 'c', 'k', 3, 'n', 'e', 't', 0, 0, 1, 0, 1};
    TEST_ASSERT_EQUAL_size_t(sizeof want, len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(want, q, sizeof want);
    TEST_ASSERT_EQUAL_size_t(len, buildDnsQuery(q, sizeof q, 0x1234, "roro9stack.net."));  // a final dot is the same name
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, ""));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, "a..b"));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, std::string(64, 'a') + ".net"));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, 20, 1, "roro9stack.net"));
}

void test_a_dns_answer() {
    // www.example.org is an alias of example.org, which has two addresses. Names after the first
    // are pointers back into the message, as servers send them.
    const uint8_t msg[] = {
        0x12, 0x34, 0x81, 0x80, 0, 1, 0, 3, 0, 0, 0, 0,
        3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'o', 'r', 'g', 0, 0, 1, 0, 1,  // the question, at 12
        0xC0, 12, 0, 5, 0, 1, 0, 0, 1, 0, 0, 2, 0xC0, 16,                                       // CNAME -> example.org (pointer to 16)
        0xC0, 16, 0, 1, 0, 1, 0, 0, 1, 0, 0, 4, 93, 184, 216, 34,
        0xC0, 16, 0, 1, 0, 1, 0, 0, 1, 0, 0, 4, 93, 184, 216, 35,
    };
    DnsAnswer a;
    TEST_ASSERT_TRUE(parseDnsAnswer(msg, sizeof msg, 0x1234, a));
    TEST_ASSERT_EQUAL_INT(0, a.rcode);
    TEST_ASSERT_FALSE(a.truncated);
    TEST_ASSERT_EQUAL_size_t(2, a.addresses.size());
    TEST_ASSERT_EQUAL_STRING("93.184.216.34", formatIpv4(a.addresses[0]).c_str());
    TEST_ASSERT_EQUAL_STRING("93.184.216.35", formatIpv4(a.addresses[1]).c_str());
    TEST_ASSERT_EQUAL_STRING("example.org", a.alias.c_str());

    TEST_ASSERT_FALSE(parseDnsAnswer(msg, sizeof msg, 0x9999, a));  // someone else's answer
    for (size_t cut = 0; cut < sizeof msg; cut++) parseDnsAnswer(msg, cut, 0x1234, a);  // cut anywhere: no crash
    TEST_ASSERT_FALSE(parseDnsAnswer(msg, sizeof msg - 3, 0x1234, a));

    uint8_t none[sizeof msg];
    std::copy(msg, msg + sizeof msg, none);
    none[3] = 0x83;  // no such name
    none[7] = 0;
    TEST_ASSERT_TRUE(parseDnsAnswer(none, 33, 0x1234, a));
    TEST_ASSERT_EQUAL_INT(3, a.rcode);
    TEST_ASSERT_EQUAL_size_t(0, a.addresses.size());

    // A pointer that points at itself must not hang the reader.
    uint8_t loop[] = {0x12, 0x34, 0x81, 0x80, 0, 1, 0, 0, 0, 0, 0, 0, 0xC0, 12, 0, 1, 0, 1};
    TEST_ASSERT_FALSE(parseDnsAnswer(loop, sizeof loop, 0x1234, a));
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_what_was_typed);
    RUN_TEST(test_an_echo_request_and_its_answers);
    RUN_TEST(test_the_summary);
    RUN_TEST(test_a_dns_query);
    RUN_TEST(test_a_dns_answer);
    return UNITY_END();
}