/*
 * Copyright (c) 2023-2024 Ian Marco Moffett and the Osmora Team.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of Hyra nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/wait.h>
#include <sys/socket.h>
#include <ostp/net/auth.h>
#include <ostp/net/stpsession.h>
#include <ostp/crypto/ecdh.h>
#include <ostp/otconfig.h>
#include <ostp/defs.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdbool.h>
#include <errno.h>
#include <string.h>

struct session_td_args {
    struct ostp_client *c;
    struct ostp_listener *lp;
    unsigned char *session_key;
};

/*
 * keypair: Master session keypair.
 * g_have_link: True when one or more link(s) are established.
 */
static struct x25519_keypair keypair;
bool g_have_link = false;

/*
 * Check a password to see if it matches with
 * the hash in /etc/shadow by using the pwcheck
 * script. Returns 0 on success.
 */
static int
pwcheck(char *username, char *pw)
{
    char *pwcheck = "/usr/local/bin/pwcheck";
    pid_t pid;
    char *args[] = {pwcheck, username, pw, NULL};
    int status;

    pid = fork();
    if (pid == 0) {
        execv(pwcheck, args);
    }

    if (waitpid(pid, &status, 0) < 0) {
        printf("waitpid() failed\n");
        return -1;
    }

    if (WIFEXITED(status)) {
        return WEXITSTATUS(status);
    }

    return -1;
}

static int
passwd_auth(struct ostp_client *c, struct ostp_listener *lp, const unsigned char *session_key)
{
    int error;
    struct session_auth auth;
    struct ostp_session *session;
    const size_t LEN = sizeof(auth);

    if (!REQUIRE_USER_AUTH) {
        return 0;
    }

    error = recv_frame(c->sockfd, sizeof(auth), session_key, &auth);
    if (error < 0) {
        return error;
    }

    if (pwcheck(auth.username, auth.password) != 0) {
        printf("Got bad password for %s\n", auth.username);
        auth.code = AUTH_BAD_PW;
        error = send_frame(c->sockfd, &auth, sizeof(auth), session_key);
        if (error < 0) {
            printf("Failed to ACK user authentication with frame\n");
        }
        return -1;
    }

    auth.code = AUTH_SUCCESS;
    error = send_frame(c->sockfd, &auth, sizeof(auth), session_key);
    if (error < 0) {
        printf("Failed to ACK user authentication with frame\n");
        return error;
    }

    session = &c->session;
    memcpy(session->username, auth.username, sizeof(auth.username));
    return 0;
}

static void
send_motd(struct ostp_client *c, const unsigned char *session_key)
{
    char motd[] = MOTD;

    printf("Sending MOTD...\n");
    if (send_frame(c->sockfd, motd, sizeof(motd), session_key) < 0) {
        printf("Failed to session MOTD\n");
    }
}

static inline void
populate_peer(struct ostp_client *c, struct peer *pp)
{
    memset(pp, 0, sizeof(*pp));
    pp->port = htons(OSTP_PORT);
    pp->pad[0] = 0xFF;
    memcpy(pp->host, c->host, IP_LEN_MAX);
}

static int
send_peerlist(struct ostp_listener *lp, struct ostp_client *c, struct ostp_session *s)
{
    struct ostp_client *c_cur, *c_tmp;
    struct peer peer_cur;
    struct peer_block pb_cur;
    struct ostp_session *sp;
    size_t peer_count;
    int i, j, error;
    int seq;

    memset(&pb_cur, 0, sizeof(pb_cur));
    peer_count = lp->client_count;
    seq = (peer_count <= 16) ? 1 : (peer_count / PB_N_PEERS);
    pb_cur.seq = seq;
    printf("peer_count=%d, seq=%d\n", peer_count, pb_cur.seq);

    /*
     * Grab peers and prepare blocks until seq == 0
     *
     *
     *        lp->clients
     *   +-------------------+
     * i=| | | | | | | | | | |
     *   +-------------------+
     *    0,1,2,3,4,5,6,7,8,9 <--+
     *                           |
     *     pb_cur.peers          |
     *   +---------------+       |  lp->clients can be bigger
     * j=| | | | | | | | |       |
     *   +---------------+       +
     *    0,1,2,3,4,5,6,7 <---+-/
     *
     *
     *  lp->clients is one big contigious array
     *  of connected clients and is used internally
     *  by the server to manage connected users.
     *
     *  During a session request, the client may request
     *  peer-to-peer mode by setting the 'P' bit. If accepted,
     *  we'll need to send one or more peer blocks back, containing
     *  a list of `PB_N_PEERS' peers max per block.
     *
     */
    for (i = 0; i < PB_N_BLOCKS && seq > 0; ++i) {
        for (j = 1; j < PB_N_PEERS && j < peer_count; ++j) {
            populate_peer(&lp->clients[j], &peer_cur);
            pb_cur.peers[j] = peer_cur;
        }

        --seq;
        pb_cur.seq = seq;
        sp = &c->session;
        session_send(&pb_cur, sizeof(pb_cur), sp);

        if (j >= lp->client_count) {
            break;
        }
    }

    return 0;
}

int
handle_srq(struct ostp_client *c, struct ostp_listener *lp, struct session_request *srq)
{
    struct session_td_args *sargs;
    struct ostp_session *session;
    int error;

    /* Do we require the SRQ 'U' bit? */
    if (REQUIRE_USER_AUTH && !ISSET(srq->options, SESSION_REQ_USER)) {
        printf("User authentication enforced but client 'U' bit not set\n");
        printf("Closing connection...\n");
        return -1;
    }

    /* Don't use P2P if not enabled */
    if (!ENABLE_P2P && !ISSET(srq->options, SESSION_REQ_P2P)) {
        printf("P2P not enabled but client 'P' bit set\n");
        printf("Closing connection...\n");
        return -1;
    }

    /* Generate a new keypair if we have no link */
    if (!g_have_link) {
        if (gen_x25519_keypair(&keypair) < 0) {
            printf("Key generation failed!\n");
            return -1;
        }
    }

    /* Send back our our public key */
    error = send(c->sockfd, keypair.pubkey, keypair.pubkey_len, 0);
    if (error < 0) {
        perror("Failed to send public key");
        return error;
    }

    /* Setup client session descriptor */
    session = &c->session;
    session->sockfd = c->sockfd;

    printf("Deriving session key...\n");
    error = gen_session_key(keypair.privkey, srq->pubkey, &session->session_key);
    if (error < 0) {
        return error;
    }

    sargs = malloc(sizeof(*sargs));
    if (sargs == NULL) {
        printf("Failed to allocate session args\n");
        return errno;
    }

    /* Try user auth, not needed if REQUIRE_USER_AUTH is 0 */
    if (passwd_auth(c, lp, session->session_key) != 0) {
        free_session_key(session->session_key);
        return -1;
    }

    /* Handle P2P requests */
    if (ISSET(srq->options, SESSION_REQ_P2P)) {
        /*
         * Not an error but we want to send this down
         * the callstack so everything is cleaned up
         * and terminated properly.
         *
         * TODO: Figure something else out here maybe...
         */
        send_peerlist(lp, c, session);
        return -1;
    }

    /* Handle any requested session parameters */
    if ((error = negotiate_spw(c, session->session_key)) < 0) {
        free_session_key(session->session_key);
        return -1;
    }

    send_motd(c, session->session_key);
    c->authed = 1;

    if (lp->on_connect != NULL)
        lp->on_connect(c);

    g_have_link = true;
    return 0;
}