#include "net_probe.h"

#include <algorithm>
#include <cstring>

#include "ipv4.h"

namespace roro::net {

namespace {
std::vector<std::string> words(const std::string& text) {
    std::vector<std::string> out;
    size_t at = 0;
    while (at < text.size()) {
        while (at < text.size() && text[at] == ' ') at++;
        size_t end = text.find(' ', at);
        if (end == std::string::npos) end = text.size();
        if (end > at) out.push_back(text.substr(at, end - at));
        at = end;
    }
    return out;
}
bool number(const std::string& s, long& out) {
    if (s.empty() || s.size() > 6) return false;
    out = 0;
    for (char c : s) {
        if (c < '0' || c > '9') return false;
        out = out * 10 + (c - '0');
    }
    return true;
}
uint16_t be16(const uint8_t* p) { return static_cast<uint16_t>((p[0] << 8) | p[1]); }
}  // namespace

std::string parsePing(const std::string& args, PingArgs& out) {
    static const char* const kUsage = "ping <host> [count] [size]";
    auto w = words(args);
    if (w.empty() || w.size() > 3 || !validHost(w[0])) return kUsage;
    PingArgs a;
    a.host = w[0];
    long n;
    if (w.size() > 1) {
        if (!number(w[1], n) || n < 1 || n > 100) return "a count from 1 to 100";
        a.count = static_cast<int>(n);
    }
    if (w.size() > 2) {
        if (!number(w[2], n) || n > 1400) return "a size from 0 to 1400 bytes";
        a.size = static_cast<int>(n);
    }
    out = a;
    return "";
}

std::string parsePort(const std::string& args, PortArgs& out) {
    static const char* const kUsage = "port <host> <port>";
    auto w = words(args);
    if (w.size() == 1) {  // host:port
        size_t colon = w[0].rfind(':');
        if (colon == std::string::npos) return kUsage;
        w = {w[0].substr(0, colon), w[0].substr(colon + 1)};
    }
    long n;
    if (w.size() != 2 || !validHost(w[0])) return kUsage;
    if (!number(w[1], n) || n < 1 || n > 65535) return "a port from 1 to 65535";
    out.host = w[0];
    out.port = static_cast<uint16_t>(n);
    return "";
}

std::string parseLookup(const std::string& args, LookupArgs& out) {
    static const char* const kUsage = "nslookup <name> [server's address]";
    auto w = words(args);
    uint32_t ip;
    if (w.empty() || w.size() > 2 || !validHost(w[0])) return kUsage;
    if (w.size() == 2 && !parseIpv4(w[1], ip)) return "the server as an address: 9.9.9.9";
    out.name = w[0];
    out.server = w.size() == 2 ? w[1] : "";
    return "";
}

uint16_t inetChecksum(const uint8_t* data, size_t len) {
    uint32_t sum = 0;
    for (size_t i = 0; i + 1 < len; i += 2) sum += static_cast<uint32_t>((data[i] << 8) | data[i + 1]);
    if (len & 1) sum += static_cast<uint32_t>(data[len - 1] << 8);
    while (sum >> 16) sum = (sum & 0xFFFF) + (sum >> 16);
    return static_cast<uint16_t>(~sum);
}

size_t buildEcho(uint8_t* out, size_t max, uint16_t id, uint16_t seq, size_t payload) {
    size_t len = 8 + payload;
    if (len > max) return 0;
    out[0] = 8;  // echo request
    out[1] = 0;
    out[2] = out[3] = 0;
    out[4] = static_cast<uint8_t>(id >> 8);
    out[5] = static_cast<uint8_t>(id);
    out[6] = static_cast<uint8_t>(seq >> 8);
    out[7] = static_cast<uint8_t>(seq);
    for (size_t i = 0; i < payload; i++) out[8 + i] = static_cast<uint8_t>('a' + i % 26);
    uint16_t sum = inetChecksum(out, len);
    out[2] = static_cast<uint8_t>(sum >> 8);
    out[3] = static_cast<uint8_t>(sum);
    return len;
}

IcmpAnswer parseIcmp(const uint8_t* packet, size_t len) {
    IcmpAnswer a;
    if (len < 20 || (packet[0] >> 4) != 4) return a;
    size_t header = static_cast<size_t>(packet[0] & 0x0F) * 4;
    if (header < 20 || len < header + 8 || packet[9] != 1) return a;  // not ICMP
    const uint8_t* icmp = packet + header;
    size_t left = len - header;
    if (icmp[0] == 0 && icmp[1] == 0) {  // echo reply
        a.kind = IcmpAnswer::Kind::Echo;
        a.id = be16(icmp + 4);
        a.seq = be16(icmp + 6);
        return a;
    }
    if (icmp[0] != 11 && icmp[0] != 3) return a;
    // Inside: the IP header of the packet it is about, and that packet's first 8 bytes.
    if (left < 8 + 20) return a;
    const uint8_t* inner = icmp + 8;
    size_t innerHeader = static_cast<size_t>(inner[0] & 0x0F) * 4;
    if ((inner[0] >> 4) != 4 || innerHeader < 20 || left < 8 + innerHeader + 8 || inner[9] != 1 || inner[innerHeader] != 8) return a;
    a.kind = icmp[0] == 11 ? IcmpAnswer::Kind::TimeExceeded : IcmpAnswer::Kind::Unreachable;
    a.id = be16(inner + innerHeader + 4);
    a.seq = be16(inner + innerHeader + 6);
    return a;
}

void PingStats::add(uint32_t ms) {
    minMs = back ? std::min(minMs, ms) : ms;
    maxMs = std::max(maxMs, ms);
    sumMs += ms;
    back++;
}

std::string PingStats::summary() const {
    int lost = sent ? (sent - back) * 100 / sent : 0;
    std::string s = std::to_string(back) + "/" + std::to_string(sent) + " back, " + std::to_string(lost) + "% lost";  // short: a Shell line is 38 characters
    if (back) s += ", " + std::to_string(minMs) + "/" + std::to_string(sumMs / static_cast<uint32_t>(back)) + "/" + std::to_string(maxMs) + " ms";
    return s;
}

size_t buildDnsQuery(uint8_t* out, size_t max, uint16_t id, const std::string& name) {
    if (name.empty() || name.size() > 253 || 12 + name.size() + 2 + 4 > max) return 0;
    std::memset(out, 0, 12);
    out[0] = static_cast<uint8_t>(id >> 8);
    out[1] = static_cast<uint8_t>(id);
    out[2] = 0x01;  // recursion wanted
    out[5] = 1;     // one question
    size_t at = 12;
    for (size_t from = 0; from <= name.size();) {
        size_t dot = name.find('.', from);
        if (dot == std::string::npos) dot = name.size();
        size_t n = dot - from;
        if (n == 0 && dot == name.size()) break;  // a final dot
        if (n == 0 || n > 63) return 0;
        out[at++] = static_cast<uint8_t>(n);
        std::memcpy(out + at, name.data() + from, n);
        at += n;
        from = dot + 1;
    }
    out[at++] = 0;
    out[at++] = 0;
    out[at++] = 1;  // A
    out[at++] = 0;
    out[at++] = 1;  // IN
    return at;
}

namespace {
// Reads a name at `at`, following the pointers DNS shortens names with. Where the name ends in
// the message (not where a pointer led), or 0 if it is broken.
size_t readName(const uint8_t* m, size_t len, size_t at, std::string* out) {
    size_t end = 0;
    int jumps = 0;
    while (at < len) {
        uint8_t n = m[at];
        if (n == 0) return end ? end : at + 1;
        if ((n & 0xC0) == 0xC0) {
            if (at + 1 >= len || ++jumps > 8) return 0;
            if (!end) end = at + 2;
            at = static_cast<size_t>(((n & 0x3F) << 8) | m[at + 1]);
            continue;
        }
        if (n > 63 || at + 1 + n > len) return 0;
        if (out) {
            if (!out->empty()) *out += '.';
            out->append(reinterpret_cast<const char*>(m + at + 1), n);
        }
        at += 1 + static_cast<size_t>(n);
    }
    return 0;
}
}  // namespace

bool parseDnsAnswer(const uint8_t* m, size_t len, uint16_t id, DnsAnswer& out) {
    if (len < 12 || be16(m) != id || !(m[2] & 0x80)) return false;
    DnsAnswer a;
    a.truncated = m[2] & 0x02;
    a.rcode = m[3] & 0x0F;
    int questions = be16(m + 4), answers = be16(m + 6);
    size_t at = 12;
    for (int i = 0; i < questions; i++) {
        at = readName(m, len, at, nullptr);
        if (!at || at + 4 > len) return false;
        at += 4;
    }
    for (int i = 0; i < answers; i++) {
        at = readName(m, len, at, nullptr);
        if (!at || at + 10 > len) return false;
        uint16_t type = be16(m + at), size = be16(m + at + 8);
        at += 10;
        if (at + size > len) return false;
        if (type == 1 && size == 4) a.addresses.push_back((static_cast<uint32_t>(m[at]) << 24) | (m[at + 1] << 16) | (m[at + 2] << 8) | m[at + 3]);
        if (type == 5) {
            std::string name;
            if (readName(m, len, at, &name)) a.alias = name;
        }
        at += size;
    }
    out = a;
    return true;
}

void buildNtpRequest(uint8_t out[kNtpPacket]) {
    std::memset(out, 0, kNtpPacket);
    out[0] = 0x23;  // no warning, version 4, a client
}

bool parseNtpAnswer(const uint8_t* p, size_t len, NtpAnswer& out) {
    if (len < kNtpPacket || (p[0] & 0x07) != 4) return false;  // not a server's
    if (p[1] == 0 || p[1] > 15) return false;                  // "kiss of death", or not synchronised
    uint32_t secs = (static_cast<uint32_t>(p[40]) << 24) | (p[41] << 16) | (p[42] << 8) | p[43];
    uint32_t frac = (static_cast<uint32_t>(p[44]) << 24) | (p[45] << 16) | (p[46] << 8) | p[47];
    if (!secs) return false;
    // NTP counts from 1900 and wraps in 2036: a small number is the era after.
    constexpr int64_t k1900To1970 = 2208988800LL;
    int64_t since1900 = secs < 0x80000000u ? static_cast<int64_t>(secs) + 4294967296LL : static_cast<int64_t>(secs);
    out.stratum = p[1];
    out.seconds = since1900 - k1900To1970;
    out.millis = static_cast<uint32_t>((static_cast<uint64_t>(frac) * 1000) >> 32);
    return true;
}

std::string clockOffset(int64_t ownMs, int64_t serverMs) {
    int64_t diff = ownMs - serverMs, size = diff < 0 ? -diff : diff;
    if (size < 100) return "right, to 0.1 s";
    std::string amount = size < 10000 ? std::to_string(size / 1000) + "." + std::to_string(size % 1000 / 100) + " s"
                         : size < 120000 ? std::to_string(size / 1000) + " s"
                         : size < 7200000 ? std::to_string(size / 60000) + " min"
                         : size < 172800000LL ? std::to_string(size / 3600000) + " h" : std::to_string(size / 86400000LL) + " days";
    return amount + (diff > 0 ? " ahead" : " behind");
}

std::string certName(const std::string& dn) {
    for (const char* key : {"CN=", "O="}) {
        size_t at = 0;
        while ((at = dn.find(key, at)) != std::string::npos) {
            if (at == 0 || dn[at - 1] == ' ' || dn[at - 1] == ',') {
                size_t from = at + std::strlen(key), end = dn.find(", ", from);
                std::string name = dn.substr(from, end == std::string::npos ? std::string::npos : end - from);
                // An old kind of string comes out as "#" and hex, type and length first: read it.
                if (name.size() > 5 && name[0] == '#' && name.size() % 2 == 1) {
                    std::string plain;
                    for (size_t i = 5; i + 1 < name.size(); i += 2) {
                        auto digit = [](char c) { return c >= '0' && c <= '9' ? c - '0' : c >= 'A' && c <= 'F' ? c - 'A' + 10 : c >= 'a' && c <= 'f' ? c - 'a' + 10 : -1; };
                        int hi = digit(name[i]), lo = digit(name[i + 1]);
                        if (hi < 0 || lo < 0 || hi * 16 + lo < 0x20 || hi * 16 + lo > 0x7E) return name;
                        plain += static_cast<char>(hi * 16 + lo);
                    }
                    return plain;
                }
                return name;
            }
            at++;
        }
    }
    return dn;
}

namespace {
// Days since a fixed day long ago (the civil calendar, leap years and all).
long dayNumber(int y, int m, int d) {
    y -= m <= 2;
    long era = (y >= 0 ? y : y - 399) / 400;
    long yoe = y - era * 400, doy = (153 * (m + (m > 2 ? -3 : 9)) + 2) / 5 + d - 1;
    return era * 146097 + yoe * 365 + yoe / 4 - yoe / 100 + doy;
}
}  // namespace

int daysBetween(int y1, int m1, int d1, int y2, int m2, int d2) { return static_cast<int>(dayNumber(y2, m2, d2) - dayNumber(y1, m1, d1)); }

const char* portLabel(uint16_t port, bool tcp) {
    if (tcp) return port == 3232 ? "updates" : port == 2323 ? "Debug Console" : port == 80 ? "sharing" : "";
    return port == 68 ? "DHCP" : port == 123 ? "NTP" : "";
}

}  // namespace roro::net