#include "services/net_tools.h"

#include <Arduino.h>

#include <NetworkClientSecure.h>

#include <lwip/etharp.h>
#include <lwip/dns.h>
#include <lwip/netdb.h>
#include <lwip/netif.h>
#include <lwip/priv/tcp_priv.h>
#include <lwip/sockets.h>
#include <lwip/tcpip.h>
#include <lwip/udp.h>
#include <mbedtls/x509_crt.h>
#include <sys/time.h>

#include <esp_random.h>

#include <algorithm>
#include <memory>

#include "ipv4.h"
#include "net_probe.h"
#include "platform/ca_roots.h"

namespace roro {

namespace {
constexpr int kMaxHops = 20;
constexpr uint32_t kPingEveryMs = 1000, kPingWaitMs = 1000, kHopWaitMs = 2000, kPortWaitMs = 5000, kDnsWaitMs = 3000;
// A TLS handshake peaks at about 52 KB of heap; below this it isn't tried.
constexpr size_t kTlsNeedsFree = 70 * 1024;

struct LwipLock {
    LwipLock() { LOCK_TCPIP_CORE(); }
    ~LwipLock() { UNLOCK_TCPIP_CORE(); }
};

std::string text(uint32_t networkOrder) { return net::formatIpv4(lwip_ntohl(networkOrder)); }

// A name or an address, as an address in network order. False: it doesn't resolve.
bool resolve(const std::string& host, uint32_t& out) {
    uint32_t ip;
    if (net::parseIpv4(host, ip)) {
        out = lwip_htonl(ip);
        return true;
    }
    struct addrinfo hints = {};
    hints.ai_family = AF_INET;
    struct addrinfo* found = nullptr;
    if (lwip_getaddrinfo(host.c_str(), nullptr, &hints, &found) != 0 || !found) return false;
    out = reinterpret_cast<struct sockaddr_in*>(found->ai_addr)->sin_addr.s_addr;
    lwip_freeaddrinfo(found);
    return true;
}

void waitMs(int socket, uint32_t ms) {
    struct timeval tv = {static_cast<time_t>(ms / 1000), static_cast<suseconds_t>((ms % 1000) * 1000)};
    lwip_setsockopt(socket, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
}
}  // namespace

struct NetTools::Job {
    enum class Kind { Ping, Trace, Port, Lookup, Tls, Ntp } kind;
    NetTools* owner;
    Console::Origin from;
    net::PingArgs ping;
    net::PortArgs port;
    net::LookupArgs lookup;

    bool stopped() const { return owner->stop_; }
    // Sends one echo request and waits for what answers it. The time in ms, or -1; `from` and
    // `kind` say who answered and how.
    int echo(int socket, uint32_t to, uint16_t id, uint16_t seq, int size, uint32_t waitFor, uint32_t& from, net::IcmpAnswer::Kind& kind);
    void runPing();
    void runTrace();
    void runPort();
    void runLookup();
    void runTls();
    void runNtp();
};

int NetTools::Job::echo(int socket, uint32_t to, uint16_t id, uint16_t seq, int size, uint32_t waitFor, uint32_t& from, net::IcmpAnswer::Kind& kind) {
    uint8_t packet[8 + 1400], answer[128];
    size_t len = net::buildEcho(packet, sizeof packet, id, seq, static_cast<size_t>(size));
    struct sockaddr_in dest = {};
    dest.sin_family = AF_INET;
    dest.sin_addr.s_addr = to;
    uint32_t sent = micros();
    if (lwip_sendto(socket, packet, len, 0, reinterpret_cast<struct sockaddr*>(&dest), sizeof dest) < 0) return -2;
    while (!stopped()) {
        uint32_t gone = (micros() - sent) / 1000;
        if (gone >= waitFor) break;
        waitMs(socket, std::min<uint32_t>(waitFor - gone, 200));  // short waits: `cancel` is noticed
        struct sockaddr_in who = {};
        socklen_t wholen = sizeof who;
        int n = lwip_recvfrom(socket, answer, sizeof answer, 0, reinterpret_cast<struct sockaddr*>(&who), &wholen);
        if (n <= 0) continue;
        net::IcmpAnswer a = net::parseIcmp(answer, static_cast<size_t>(n));
        if (a.kind == net::IcmpAnswer::Kind::Other || a.id != id || a.seq != seq) continue;  // somebody else's
        from = who.sin_addr.s_addr;
        kind = a.kind;
        return static_cast<int>((micros() - sent + 500) / 1000);
    }
    return -1;
}

void NetTools::Job::runPing() {
    uint32_t to;
    if (!resolve(ping.host, to)) return (void)console.printf("ping: %s doesn't resolve\n", ping.host.c_str());
    int s = lwip_socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (s < 0) return (void)console.println("ping: error no socket");
    console.printf("ping: %s, %d bytes\n", text(to).c_str(), ping.size);
    uint16_t id = static_cast<uint16_t>(esp_random());
    net::PingStats stats;
    for (int seq = 1; seq <= ping.count && !stopped(); seq++) {
        uint32_t started = millis(), from = 0;
        net::IcmpAnswer::Kind kind = net::IcmpAnswer::Kind::Other;
        stats.sent++;
        int ms = echo(s, to, id, static_cast<uint16_t>(seq), ping.size, kPingWaitMs, from, kind);
        if (ms == -2) console.printf("ping: %d not sent: no route, or too big\n", seq);
        else if (ms < 0) console.printf("ping: %d no answer\n", seq);
        else if (kind == net::IcmpAnswer::Kind::Echo) {
            stats.add(static_cast<uint32_t>(ms));
            console.printf("ping: %d %d ms\n", seq, ms);
        } else {
            console.printf("ping: %d %s says %s\n", seq, text(from).c_str(), kind == net::IcmpAnswer::Kind::TimeExceeded ? "too many hops" : "unreachable");
        }
        while (seq < ping.count && !stopped() && millis() - started < kPingEveryMs) delay(50);
    }
    lwip_close(s);
    console.printf("ping: %s%s\n", stopped() ? "stopped, " : "", stats.summary().c_str());
}

// An echo request allowed one hop, then two, then three: each router that drops it says so, and
// that is the list.
void NetTools::Job::runTrace() {
    uint32_t to;
    if (!resolve(ping.host, to)) return (void)console.printf("traceroute: %s doesn't resolve\n", ping.host.c_str());
    int s = lwip_socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (s < 0) return (void)console.println("traceroute: error no socket");
    console.printf("traceroute: to %s, %d hops at most\n", text(to).c_str(), kMaxHops);
    uint16_t id = static_cast<uint16_t>(esp_random());
    bool arrived = false;
    for (int hop = 1; hop <= kMaxHops && !stopped() && !arrived; hop++) {
        int ttl = hop;
        lwip_setsockopt(s, IPPROTO_IP, IP_TTL, &ttl, sizeof ttl);
        uint32_t from = 0;
        net::IcmpAnswer::Kind kind = net::IcmpAnswer::Kind::Other;
        int ms = echo(s, to, id, static_cast<uint16_t>(hop), 32, kHopWaitMs, from, kind);
        if (ms < 0) console.printf("traceroute: %2d *\n", hop);
        else console.printf("traceroute: %2d %s %d ms%s\n", hop, text(from).c_str(), ms, kind == net::IcmpAnswer::Kind::Unreachable ? " unreachable" : "");
        arrived = ms >= 0 && kind != net::IcmpAnswer::Kind::TimeExceeded;
    }
    lwip_close(s);
    console.printf("traceroute: %s\n", stopped() ? "stopped" : arrived ? "arrived" : "not reached");
}

void NetTools::Job::runPort() {
    uint32_t to;
    if (!resolve(port.host, to)) return (void)console.printf("port: %s doesn't resolve\n", port.host.c_str());
    int s = lwip_socket(AF_INET, SOCK_STREAM, 0);
    if (s < 0) return (void)console.println("port: error no socket");
    lwip_fcntl(s, F_SETFL, lwip_fcntl(s, F_GETFL, 0) | O_NONBLOCK);
    struct sockaddr_in dest = {};
    dest.sin_family = AF_INET;
    dest.sin_port = lwip_htons(port.port);
    dest.sin_addr.s_addr = to;
    uint32_t started = millis();
    int error = lwip_connect(s, reinterpret_cast<struct sockaddr*>(&dest), sizeof dest) == 0 ? 0 : errno;
    bool answered = error == 0;
    while (error == EINPROGRESS && !answered && !stopped() && millis() - started < kPortWaitMs) {
        fd_set writable, failed;
        FD_ZERO(&writable);
        FD_ZERO(&failed);
        FD_SET(s, &writable);
        FD_SET(s, &failed);
        struct timeval tv = {0, 200000};
        if (lwip_select(s + 1, nullptr, &writable, &failed, &tv) > 0) {
            socklen_t len = sizeof error;
            lwip_getsockopt(s, SOL_SOCKET, SO_ERROR, &error, &len);
            answered = true;
        }
    }
    uint32_t ms = millis() - started;
    lwip_close(s);
    std::string where = text(to) + ":" + std::to_string(port.port);
    if (answered && error == 0) console.printf("port: %s open, %lu ms\n", where.c_str(), (unsigned long)ms);
    else if (answered && (error == ECONNREFUSED || error == ECONNRESET)) console.printf("port: %s refused, %lu ms: the host is there, nothing listens\n", where.c_str(), (unsigned long)ms);
    else if (answered || error != EINPROGRESS) console.printf("port: %s no route to it (error %d)\n", where.c_str(), error);
    else if (stopped()) console.println("port: stopped");
    else console.printf("port: %s no answer in %lu s: down, or filtered\n", where.c_str(), (unsigned long)(kPortWaitMs / 1000));
}

// Asks one server directly, so that the answer says which server and how long, which the
// system's own resolver doesn't.
void NetTools::Job::runLookup() {
    uint32_t server = 0;
    if (!lookup.server.empty()) resolve(lookup.server, server);
    else {
        LwipLock lock;
        const ip_addr_t* first = dns_getserver(0);
        if (first && IP_IS_V4(first)) server = ip_2_ip4(first)->addr;
    }
    if (!server) return (void)console.println("nslookup: no DNS server is set");
    int s = lwip_socket(AF_INET, SOCK_DGRAM, 0);
    if (s < 0) return (void)console.println("nslookup: error no socket");
    uint8_t query[300], answer[512];
    uint16_t id = static_cast<uint16_t>(esp_random());
    size_t len = net::buildDnsQuery(query, sizeof query, id, lookup.name);
    struct sockaddr_in dest = {};
    dest.sin_family = AF_INET;
    dest.sin_port = lwip_htons(53);
    dest.sin_addr.s_addr = server;
    uint32_t started = millis();
    net::DnsAnswer result;
    bool got = false;
    if (len && lwip_sendto(s, query, len, 0, reinterpret_cast<struct sockaddr*>(&dest), sizeof dest) >= 0) {
        while (!got && !stopped() && millis() - started < kDnsWaitMs) {
            waitMs(s, 200);
            int n = lwip_recv(s, answer, sizeof answer, 0);
            if (n > 0) got = net::parseDnsAnswer(answer, static_cast<size_t>(n), id, result);
        }
    }
    uint32_t ms = millis() - started;
    lwip_close(s);
    if (!got) return (void)console.printf("nslookup: no answer from %s in %lu s\n", text(server).c_str(), (unsigned long)(kDnsWaitMs / 1000));
    console.printf("nslookup: %s answered in %lu ms\n", text(server).c_str(), (unsigned long)ms);
    if (!result.alias.empty()) console.printf("nslookup: %s is %s\n", lookup.name.c_str(), result.alias.c_str());
    for (uint32_t ip : result.addresses) console.printf("nslookup: %s\n", net::formatIpv4(ip).c_str());
    if (result.rcode == 3) console.printf("nslookup: there is no %s\n", lookup.name.c_str());
    else if (result.rcode) console.printf("nslookup: the server refused (code %d)\n", result.rcode);
    else if (result.addresses.empty()) console.printf("nslookup: %s has no IPv4 address%s\n", lookup.name.c_str(), result.truncated ? " in a first packet" : "");
}

// A handshake that checks nothing, to see the certificate whatever it is; then the certificate is
// checked here, against this device's own roots and the name asked for, and the answer is said in
// words. It is what the Update Service's connection would have decided.
void NetTools::Job::runTls() {
    if (ESP.getFreeHeap() < kTlsNeedsFree)
        return (void)console.printf("tls: not enough memory (%u KB free, %u needed): close IRC or a Gemini page\n", (unsigned)(ESP.getFreeHeap() / 1024),
                                    (unsigned)(kTlsNeedsFree / 1024));
    NetworkClientSecure tls;
    tls.setInsecure();
    uint32_t started = millis();
    if (!tls.connect(port.host.c_str(), port.port, 8000)) {
        char why[100] = "";
        tls.lastError(why, sizeof why);
        return (void)console.printf("tls: no handshake with %s:%u in %lu ms: %s\n", port.host.c_str(), (unsigned)port.port, (unsigned long)(millis() - started),
                                    why[0] ? why : "no connection");
    }
    console.printf("tls: %s:%u answered in %lu ms\n", port.host.c_str(), (unsigned)port.port, (unsigned long)(millis() - started));
    const mbedtls_x509_crt* cert = tls.getPeerCertificate();
    if (!cert) {
        tls.stop();
        return (void)console.println("tls: it showed no certificate");
    }
    char dn[200];
    std::string subject = mbedtls_x509_dn_gets(dn, sizeof dn, &cert->subject) > 0 ? net::certName(dn) : "?";
    std::string issuer = mbedtls_x509_dn_gets(dn, sizeof dn, &cert->issuer) > 0 ? net::certName(dn) : "?";
    console.printf("tls: for %s, by %s\n", subject.c_str(), issuer.c_str());
    const mbedtls_x509_time& from = cert->valid_from;
    const mbedtls_x509_time& to = cert->valid_to;
    time_t now = time(nullptr);
    struct tm today;
    gmtime_r(&now, &today);
    bool clock = today.tm_year + 1900 >= 2024;
    int left = net::daysBetween(today.tm_year + 1900, today.tm_mon + 1, today.tm_mday, to.year, to.mon, to.day);
    console.printf("tls: valid %04d-%02d-%02d to %04d-%02d-%02d", from.year, from.mon, from.day, to.year, to.mon, to.day);
    if (!clock) console.println(" (this clock isn't set)");
    else if (left >= 0) console.printf(", %d days left\n", left);
    else console.printf(", EXPIRED %d days ago\n", -left);

    mbedtls_x509_crt roots;
    mbedtls_x509_crt_init(&roots);
    uint32_t flags = 0;
    bool parsed = mbedtls_x509_crt_parse(&roots, reinterpret_cast<const unsigned char*>(kTrustedRootsPem), sizeof kTrustedRootsPem) == 0;
    int verdict = parsed ? mbedtls_x509_crt_verify(const_cast<mbedtls_x509_crt*>(cert), &roots, nullptr, port.host.c_str(), &flags, nullptr, nullptr) : -1;
    mbedtls_x509_crt_free(&roots);
    if (verdict == 0) console.println("tls: this device trusts it");
    else {
        std::string why;
        if ((flags & MBEDTLS_X509_BADCERT_EXPIRED) || (clock && left < 0)) why += ", expired";
        if (flags & MBEDTLS_X509_BADCERT_FUTURE) why += ", not valid yet";
        if (flags & MBEDTLS_X509_BADCERT_CN_MISMATCH) why += ", not for that name";
        if (flags & MBEDTLS_X509_BADCERT_NOT_TRUSTED) why += ", not signed by a root this device has";
        if (why.empty()) why = ", it doesn't check out";
        console.printf("tls: NOT trusted here: %s\n", why.c_str() + 2);
    }
    uint8_t sha[32];
    if (tls.getFingerprintSHA256(sha)) {
        char hex[65];
        for (int i = 0; i < 32; i++) std::snprintf(hex + i * 2, 3, "%02x", sha[i]);
        console.printf("tls: sha256 %s\n", hex);
    }
    tls.stop();
}

// Asks a time server and compares with this device's clock, allowing for half the round trip.
void NetTools::Job::runNtp() {
    uint32_t to;
    if (!resolve(port.host, to)) return (void)console.printf("ntp: %s doesn't resolve\n", port.host.c_str());
    int s = lwip_socket(AF_INET, SOCK_DGRAM, 0);
    if (s < 0) return (void)console.println("ntp: error no socket");
    uint8_t packet[net::kNtpPacket];
    net::buildNtpRequest(packet);
    struct sockaddr_in dest = {};
    dest.sin_family = AF_INET;
    dest.sin_port = lwip_htons(123);
    dest.sin_addr.s_addr = to;
    uint32_t sent = micros();
    net::NtpAnswer answer;
    bool got = false;
    struct timeval own = {};
    if (lwip_sendto(s, packet, sizeof packet, 0, reinterpret_cast<struct sockaddr*>(&dest), sizeof dest) >= 0) {
        while (!got && !stopped() && (micros() - sent) / 1000 < kDnsWaitMs) {
            waitMs(s, 200);
            int n = lwip_recv(s, packet, sizeof packet, 0);
            if (n <= 0) continue;
            gettimeofday(&own, nullptr);
            got = net::parseNtpAnswer(packet, static_cast<size_t>(n), answer);
        }
    }
    uint32_t tripMs = (micros() - sent) / 1000;
    lwip_close(s);
    if (!got) return (void)console.printf("ntp: no answer from %s (%s) in %lu s\n", port.host.c_str(), text(to).c_str(), (unsigned long)(kDnsWaitMs / 1000));
    console.printf("ntp: %s (%s), stratum %d, %lu ms away\n", port.host.c_str(), text(to).c_str(), answer.stratum, (unsigned long)tripMs);
    int64_t ownMs = static_cast<int64_t>(own.tv_sec) * 1000 + own.tv_usec / 1000, serverMs = answer.seconds * 1000 + answer.millis + tripMs / 2;
    if (own.tv_sec < 1700000000) console.println("ntp: this clock isn't set");
    else console.printf("ntp: this clock is %s\n", net::clockOffset(ownMs, serverMs).c_str());
}

void NetTools::task(void* arg) {
    Job* job = static_cast<Job*>(arg);
    {
        Console::As as(job->from);  // the lines go to the console that asked (the Shell shows only its own)
        switch (job->kind) {
            case Job::Kind::Ping: job->runPing(); break;
            case Job::Kind::Trace: job->runTrace(); break;
            case Job::Kind::Port: job->runPort(); break;
            case Job::Kind::Lookup: job->runLookup(); break;
            case Job::Kind::Tls: job->runTls(); break;
            case Job::Kind::Ntp: job->runNtp(); break;
        }
    }
    NetTools* owner = job->owner;
    delete job;
    owner->busy_ = false;
    vTaskDelete(nullptr);
}

void NetTools::start(Job* job) {
    job->owner = this;
    job->from = console.origin();
    stop_ = false;
    busy_ = true;
    // A TLS handshake needs far more stack than a ping.
    if (xTaskCreate(task, "nettool", job->kind == Job::Kind::Tls ? 12288 : 6144, job, 1, nullptr) != pdPASS) {
        busy_ = false;
        delete job;
        console.println("net: error not enough memory for it");
    }
}

bool NetTools::command(const std::string& line) {
    size_t space = line.find(' ');
    std::string name = line.substr(0, space), args = space == std::string::npos ? "" : line.substr(space + 1);
    if (name == "ifconfig") {
        LwipLock lock;
        for (struct netif* n = netif_list; n; n = n->next) {
            if (n->name[0] == 'l' && n->name[1] == 'o') continue;
            const char* label = n->name[0] == 's' && n->name[1] == 't' ? "wifi" : n->name[0] == 'w' && n->name[1] == 'g' ? "vpn" : nullptr;
            char raw[4] = {n->name[0], n->name[1], static_cast<char>('0' + n->num % 10), 0};
            bool up = netif_is_up(n) && netif_is_link_up(n);
            console.printf("ifconfig: %s %s/%d", label ? label : raw, text(netif_ip4_addr(n)->addr).c_str(), __builtin_popcount(netif_ip4_netmask(n)->addr));
            if (netif_ip4_gw(n)->addr) console.printf(" gw %s", text(netif_ip4_gw(n)->addr).c_str());
            console.printf(" mtu %u, %s%s\n", (unsigned)n->mtu, up ? "up" : "down", n == netif_default ? ", default route" : "");
        }
        std::string servers;
        for (u8_t i = 0; i < DNS_MAX_SERVERS; i++) {
            const ip_addr_t* d = dns_getserver(i);
            if (d && IP_IS_V4(d) && ip_2_ip4(d)->addr) servers += " " + text(ip_2_ip4(d)->addr);
        }
        console.printf("ifconfig: dns%s\n", servers.empty() ? " none" : servers.c_str());
        return true;
    }
    if (name == "arp") {
        LwipLock lock;
        int found = 0;
        for (size_t i = 0; i < ARP_TABLE_SIZE; i++) {
            ip4_addr_t* ip;
            struct netif* nif;
            struct eth_addr* mac;
            if (!etharp_get_entry(i, &ip, &nif, &mac)) continue;
            found++;
            console.printf("arp: %-15s %02x:%02x:%02x:%02x:%02x:%02x\n", text(ip->addr).c_str(), mac->addr[0], mac->addr[1], mac->addr[2], mac->addr[3], mac->addr[4],
                           mac->addr[5]);
        }
        if (!found) console.println("arp: nobody heard yet on this network");
        return true;
    }
    if (name == "netstat") {
        LwipLock lock;
        for (struct tcp_pcb_listen* p = tcp_listen_pcbs.listen_pcbs; p; p = p->next) {
            const char* label = net::portLabel(p->local_port, true);
            console.printf("netstat: tcp %u listens%s%s%s\n", (unsigned)p->local_port, *label ? " (" : "", label, *label ? ")" : "");
        }
        for (struct tcp_pcb* p = tcp_active_pcbs; p; p = p->next) {
            std::string local = IP_IS_V4(&p->local_ip) ? text(ip_2_ip4(&p->local_ip)->addr) : "::", remote = IP_IS_V4(&p->remote_ip) ? text(ip_2_ip4(&p->remote_ip)->addr) : "::";
            console.printf("netstat: tcp %s:%u - %s:%u\n", local.c_str(), (unsigned)p->local_port, remote.c_str(), (unsigned)p->remote_port);
        }
        for (struct udp_pcb* p = udp_pcbs; p; p = p->next) {
            const char* label = net::portLabel(p->local_port, false);
            console.printf("netstat: udp %u%s%s%s\n", (unsigned)p->local_port, *label ? " (" : "", label, *label ? ")" : "");
        }
        return true;
    }
    if (name != "ping" && name != "traceroute" && name != "port" && name != "nslookup" && name != "tls" && name != "ntp") return false;
    std::unique_ptr<Job> job(new Job());
    std::string why;
    if (name == "ping") {
        job->kind = Job::Kind::Ping;
        why = net::parsePing(args, job->ping);
    } else if (name == "traceroute") {
        job->kind = Job::Kind::Trace;
        why = net::parsePing(args, job->ping);
        if (!why.empty() || args.find(' ') != std::string::npos) why = "traceroute <host>";
    } else if (name == "port") {
        job->kind = Job::Kind::Port;
        why = net::parsePort(args, job->port);
    } else if (name == "tls") {  // the port may be left out: 443
        job->kind = Job::Kind::Tls;
        bool onlyHost = !args.empty() && args.find(' ') == std::string::npos && args.find(':') == std::string::npos;
        why = net::parsePort(onlyHost ? args + " 443" : args, job->port);
        if (!why.empty() && why.find("port <") == 0) why = "tls <host> [port]";
    } else if (name == "ntp") {  // the server may be left out: the first one in use
        job->kind = Job::Kind::Ntp;
        job->port.host = !args.empty() ? args : ntpServer ? ntpServer() : "";
        if (job->port.host.empty() || !net::validHost(job->port.host)) why = "ntp [server]";
    } else {
        job->kind = Job::Kind::Lookup;
        why = net::parseLookup(args, job->lookup);
    }
    if (!why.empty()) return console.printf("%s: %s\n", name.c_str(), why.c_str()), true;
    if (busy_) return console.printf("%s: another one is running: `cancel` stops it\n", name.c_str()), true;
    start(job.release());
    return true;
}

}  // namespace roro