/*
 * 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 <net/auth.h>
#include <crypto/ecdh.h>
#include <otconfig.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>

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

/*
 * 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("waidpid() failed\n");
        return -1;
    }

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

    return -1;
}

static int
passwd_auth(struct ostp_client *c, const unsigned char *session_key)
{
    int error;
    struct session_auth auth;
    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;
    }
    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 int
session_run(struct ostp_listener *lp, const unsigned char *session_key)
{
    struct ostp_client *c;
    char buf[4096];
    size_t len;

    while (1) {
        for (int i = 1; i < MAX_CLIENTS; ++i) {
            c = &lp->clients[i];
            if (c->sockfd <= 0)
                continue;
            if (FD_ISSET(c->sockfd, &lp->client_fds) <= 0)
                continue;

            len = recv_frame(c->sockfd, sizeof(buf) - 1, session_key, buf);
            if (len < 0) {
                printf("recv_frame() failure, packet lost\n");
                continue;
            }
            if (len == 0) {
                return 0;
            }
            if (lp->on_recv != NULL) {
                lp->on_recv(c, buf, len);
            }
        }
    }
}

static void *
session_td(void *args)
{
    struct session_td_args *tmp = args;
    int error;

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

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

    send_motd(tmp->c, tmp->session_key);
    session_run(tmp->lp, tmp->session_key);
    free(args);
    return NULL;
}

int
handle_srq(struct ostp_client *c, struct ostp_listener *lp, struct session_request *srq)
{
    struct x25519_keypair keypair;
    struct session_td_args *sargs;
    unsigned char *session_key;
    int error;

    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;
    }

    printf("Generating keys...\n");

    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;
    }

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

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

    sargs->c = c;
    sargs->lp = lp;
    sargs->session_key = session_key;
    error = pthread_create(&c->td, NULL, session_td, sargs);
    if (error != 0) {
        return error;
    }

    return 0;
}