// testSsh.c -- exercises the SSH/SFTP library against a throwaway, unprivileged sshd on
// loopback. The test generates a host key and a client key, starts sshd on an ephemeral port
// authenticating the CURRENT user by that key, then drives sshConnect / sshAuthKey / sshExec
// and the sftp* file operations from a Lua context. Not part of `make test` (it needs a system
// sshd + ssh-keygen); run it with `make ssh-test`.

#define _GNU_SOURCE

#include "calog.h"
#include "calogSsh.h"

#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <pwd.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define CHECK(cond, msg) checkImpl((cond), (msg), __FILE__, __LINE__)

#define PUMP_LIMIT   6000
#define RESULT_SLOTS 12

static CalogT         *calog      = NULL;
static _Atomic int64_t results[RESULT_SLOTS];
static _Atomic int32_t doneCount   = 0;
static _Atomic int32_t errorCount   = 0;
static int32_t         testsRun     = 0;
static int32_t         testsFailed  = 0;

static void    checkImpl(bool condition, const char *message, const char *file, int32_t line);
static int     freePort(void);
static int32_t nativeDone(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static int32_t nativeReport(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static void    onError(uint64_t contextId, const char *message, void *userData);
static void    pumpUntilDone(int32_t target);
static bool    runCommand(const char *path, char *const argv[]);
static bool    waitForListen(int port);


static void checkImpl(bool condition, const char *message, const char *file, int32_t line) {
    testsRun++;
    if (!condition) {
        testsFailed++;
        printf("FAIL %s:%d  %s\n", file, line, message);
    }
}


static int freePort(void) {
    struct sockaddr_in addr;
    socklen_t          len;
    int                fd;
    int                port;

    fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }
    memset(&addr, 0, sizeof(addr));
    addr.sin_family      = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_port        = 0;
    if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        close(fd);
        return -1;
    }
    len = sizeof(addr);
    if (getsockname(fd, (struct sockaddr *)&addr, &len) != 0) {
        close(fd);
        return -1;
    }
    port = (int)ntohs(addr.sin_port);
    close(fd);
    return port;
}


static int32_t nativeDone(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)args;
    (void)argCount;
    (void)userData;
    atomic_fetch_add(&doneCount, 1);
    calogValueNil(result);
    return calogOkE;
}


static int32_t nativeReport(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)userData;
    calogValueNil(result);
    if (argCount != 2 || args[0].type != calogIntE) {
        return calogFail(result, calogErrArgE, "report expects (tag, value)");
    }
    if (args[0].as.i < 0 || args[0].as.i >= RESULT_SLOTS) {
        return calogFail(result, calogErrArgE, "report: tag out of range");
    }
    atomic_store(&results[args[0].as.i], (args[1].type == calogIntE) ? args[1].as.i : (int64_t)args[1].as.r);
    return calogOkE;
}


static void onError(uint64_t contextId, const char *message, void *userData) {
    (void)contextId;
    (void)userData;
    fprintf(stderr, "  [script error] %s\n", (message != NULL) ? message : "(null)");
    atomic_fetch_add(&errorCount, 1);
}


static void pumpUntilDone(int32_t target) {
    struct timespec ts = { 0, 1000000 };
    int             i;

    for (i = 0; i < PUMP_LIMIT; i++) {
        calogPump(calog);
        if (atomic_load(&doneCount) >= target) {
            calogPump(calog);
            return;
        }
        nanosleep(&ts, NULL);
    }
}


static bool runCommand(const char *path, char *const argv[]) {
    pid_t pid;
    int   status;

    pid = fork();
    if (pid < 0) {
        return false;
    }
    if (pid == 0) {
        int devnull;
        devnull = open("/dev/null", 1);
        if (devnull >= 0) {
            dup2(devnull, 1);
            dup2(devnull, 2);
        }
        execv(path, argv);
        _exit(127);
    }
    if (waitpid(pid, &status, 0) != pid) {
        return false;
    }
    return WIFEXITED(status) && WEXITSTATUS(status) == 0;
}


static bool waitForListen(int port) {
    int i;

    for (i = 0; i < 200; i++) {
        struct sockaddr_in addr;
        struct timespec    ts = { 0, 20000000 };
        int                fd;
        int                rc;
        fd = socket(AF_INET, SOCK_STREAM, 0);
        if (fd < 0) {
            return false;
        }
        memset(&addr, 0, sizeof(addr));
        addr.sin_family      = AF_INET;
        addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        addr.sin_port        = htons((uint16_t)port);
        rc = connect(fd, (struct sockaddr *)&addr, sizeof(addr));
        close(fd);
        if (rc == 0) {
            return true;
        }
        nanosleep(&ts, NULL);
    }
    return false;
}


int main(void) {
    CalogContextT *ctx;
    struct passwd *pw;
    const char    *tmp;
    const char    *user;
    char           dir[256];
    char           hostKey[320];
    char           clientKey[320];
    char           clientPub[328];
    char           authKeys[320];
    char           remoteFile[384];
    char           sshdArg[400];
    char           akoArg[400];
    char           script[8192];
    char          *cpCmd[5];
    char          *sshdCmd[16];
    pid_t          sshd;
    int            port;
    int32_t        i;

    tmp = getenv("TMPDIR");
    if (tmp == NULL) {
        tmp = "/tmp";
    }
    pw = getpwuid(getuid());
    if (pw == NULL) {
        printf("cannot determine current user\n");
        return 1;
    }
    user = pw->pw_name;
    snprintf(dir, sizeof(dir), "%s/calogSshTest_%ld", tmp, (long)getpid());
    {
        char *mkdirCmd[4];
        mkdirCmd[0] = (char *)"/bin/mkdir";
        mkdirCmd[1] = (char *)"-p";
        mkdirCmd[2] = dir;
        mkdirCmd[3] = NULL;
        if (!runCommand("/bin/mkdir", mkdirCmd)) {
            printf("mkdir failed\n");
            return 1;
        }
    }
    snprintf(hostKey, sizeof(hostKey), "%s/host_ed25519", dir);
    snprintf(clientKey, sizeof(clientKey), "%s/client_ed25519", dir);
    snprintf(clientPub, sizeof(clientPub), "%s.pub", clientKey);
    snprintf(authKeys, sizeof(authKeys), "%s/authorized_keys", dir);
    snprintf(remoteFile, sizeof(remoteFile), "%s/remote.txt", dir);

    {
        char *kh[9];
        char *kc[9];
        kh[0] = (char *)"/usr/bin/ssh-keygen"; kh[1] = (char *)"-q"; kh[2] = (char *)"-t"; kh[3] = (char *)"ed25519";
        kh[4] = (char *)"-f"; kh[5] = hostKey; kh[6] = (char *)"-N"; kh[7] = (char *)""; kh[8] = NULL;
        kc[0] = (char *)"/usr/bin/ssh-keygen"; kc[1] = (char *)"-q"; kc[2] = (char *)"-t"; kc[3] = (char *)"ed25519";
        kc[4] = (char *)"-f"; kc[5] = clientKey; kc[6] = (char *)"-N"; kc[7] = (char *)""; kc[8] = NULL;
        if (!runCommand("/usr/bin/ssh-keygen", kh) || !runCommand("/usr/bin/ssh-keygen", kc)) {
            printf("ssh-keygen failed\n");
            return 1;
        }
    }
    // authorized_keys = the client public key.
    cpCmd[0] = (char *)"/bin/cp";
    cpCmd[1] = clientPub;
    cpCmd[2] = authKeys;
    cpCmd[3] = NULL;
    if (!runCommand("/bin/cp", cpCmd)) {
        printf("cp authorized_keys failed\n");
        return 1;
    }

    port = freePort();
    if (port < 0) {
        printf("no free port\n");
        return 1;
    }
    snprintf(sshdArg, sizeof(sshdArg), "%d", port);
    snprintf(akoArg, sizeof(akoArg), "AuthorizedKeysFile=%s", authKeys);
    sshdCmd[0]  = (char *)"/usr/sbin/sshd";
    sshdCmd[1]  = (char *)"-D";
    sshdCmd[2]  = (char *)"-p";
    sshdCmd[3]  = sshdArg;
    sshdCmd[4]  = (char *)"-h";
    sshdCmd[5]  = hostKey;
    sshdCmd[6]  = (char *)"-o";
    sshdCmd[7]  = (char *)"PidFile=none";
    sshdCmd[8]  = (char *)"-o";
    sshdCmd[9]  = (char *)"UsePAM=no";
    sshdCmd[10] = (char *)"-o";
    sshdCmd[11] = (char *)"StrictModes=no";
    sshdCmd[12] = (char *)"-o";
    sshdCmd[13] = akoArg;
    sshdCmd[14] = NULL;
    sshd = fork();
    if (sshd < 0) {
        printf("fork sshd failed\n");
        return 1;
    }
    if (sshd == 0) {
        int devnull;
        devnull = open("/dev/null", 1);
        if (devnull >= 0) {
            dup2(devnull, 1);
            dup2(devnull, 2);
        }
        execv("/usr/sbin/sshd", sshdCmd);
        _exit(127);
    }
    if (!waitForListen(port)) {
        printf("sshd did not start listening\n");
        kill(sshd, SIGTERM);
        waitpid(sshd, NULL, 0);
        return 1;
    }

    calog = calogCreate();
    if (calog == NULL) {
        printf("calog create failed\n");
        kill(sshd, SIGTERM);
        waitpid(sshd, NULL, 0);
        return 1;
    }
    calogSetErrorHandler(calog, onError, NULL);
    calogRegister(calog, "report", nativeReport, NULL);
    calogRegister(calog, "done", nativeDone, NULL);
    calogSshRegister(calog);
    for (i = 0; i < RESULT_SLOTS; i++) {
        atomic_store(&results[i], -1);
    }

    snprintf(script, sizeof(script),
        "local h = sshConnect('127.0.0.1', %d)\n"
        "report(1, h and 1 or 0)\n"
        "local ok = sshAuthKey(h, '%s', '%s', '%s', '')\n"
        "report(2, ok and 1 or 0)\n"
        "local r = sshExec(h, 'echo hello')\n"
        "report(3, (r.stdout == 'hello\\n' and r.exitCode == 0) and 1 or 0)\n"
        "sftpPut(h, '%s', 'sftp payload')\n"
        "report(4, sftpGet(h, '%s') == 'sftp payload' and 1 or 0)\n"
        "local st = sftpStat(h, '%s')\n"
        "report(5, (st ~= nil and st.size == 12 and st.isDir == false) and 1 or 0)\n"
        "local found = 0\n"
        "for _, e in ipairs(sftpList(h, '%s')) do if e.name == 'remote.txt' then found = 1 end end\n"
        "report(6, found)\n"
        "sftpRemove(h, '%s')\n"
        "report(7, sftpStat(h, '%s') == nil and 1 or 0)\n"
        "local big = sshExec(h, 'seq 1 20000')\n"
        "report(8, (big.exitCode == 0 and #big.stdout > 100000 and big.stdout:sub(-6) == '20000\\n') and 1 or 0)\n"
        "sshClose(h)\n"
        "done()",
        port, user, clientKey, clientPub, remoteFile, remoteFile, remoteFile, dir, remoteFile, remoteFile);

    ctx = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(ctx, script);
    pumpUntilDone(1);

    CHECK(atomic_load(&results[1]) == 1, "sshConnect returns a handle");
    CHECK(atomic_load(&results[2]) == 1, "sshAuthKey authenticates with the client key");
    CHECK(atomic_load(&results[3]) == 1, "sshExec runs a command and returns stdout + exitCode");
    CHECK(atomic_load(&results[4]) == 1, "sftpPut then sftpGet round-trips a file");
    CHECK(atomic_load(&results[5]) == 1, "sftpStat reports the file size and that it is not a directory");
    CHECK(atomic_load(&results[6]) == 1, "sftpList includes the written file");
    CHECK(atomic_load(&results[7]) == 1, "sftpRemove deletes the file (stat then returns nil)");
    CHECK(atomic_load(&results[8]) == 1, "sshExec drains a large multi-realloc output completely");
    CHECK(atomic_load(&errorCount) == 0, "no uncaught errors");

    calogContextClose(ctx);
    calogDestroy(calog);
    calogSshShutdown();

    kill(sshd, SIGTERM);
    waitpid(sshd, NULL, 0);
    {
        char *rmCmd[4];
        rmCmd[0] = (char *)"/bin/rm";
        rmCmd[1] = (char *)"-rf";
        rmCmd[2] = dir;
        rmCmd[3] = NULL;
        (void)runCommand("/bin/rm", rmCmd);
    }

    printf("\n%d checks, %d failed\n", testsRun, testsFailed);
    fflush(stdout);
    if (testsFailed != 0) {
        return 1;
    }
    return 0;
}