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

#define KEY_BYTE_WIDTH 32

static void
log_pubkey(uint8_t pubkey[KEY_BYTE_WIDTH])
{
    for (size_t i = 0; i < KEY_BYTE_WIDTH; ++i) {
        printf("%02X ", pubkey[i] & 0xFF);
        if (i != 0 && i % 4 == 0) {
            printf("\n");
        }
    }

    printf("\n");
}

static void
send_motd(int client_fd, const unsigned char *session_key)
{
    char motd[] = MOTD;

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

/*
 * 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(int client_fd, 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(client_fd, 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(client_fd, &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(client_fd, &auth, sizeof(auth), session_key);
    if (error < 0) {
        printf("Failed to ACK user authentication with frame\n");
        return error;
    }
    return 0;
}

static int
client_echo(int client_fd, const unsigned char *session_key)
{
    char buf[4096];
    int error;

    error = recv_frame(client_fd, sizeof(buf) - 1, session_key, buf);
    if (error < 0) {
        return error;
    }

    /* Echo frame to all clients */
    for (size_t i = 1; i < MAX_CLIENTS; ++i) {
        if (clients[i] <= 0)
            continue;

        send_frame(clients[i], buf, sizeof(buf), session_key);
    }

    return 0;
}
/*
 * Verify the session request packet and handle
 * the rest.
 *
 * @client_fd: File descriptor for client socket.
 * @srq: Session request packet.
 */
int
handle_srq(int client_fd, struct session_request *srq)
{
    struct x25519_keypair keypair;
    unsigned char *session_key;
    int error;

    if (REQUIRE_USER_AUTH && !ISSET(srq->options, SESSION_REQ_USER)) {
        printf("%x\n", srq->options);
        printf("User authentication enforced but client 'U' bit not set\n");
        printf("Closing connection...\n");
        return -1;
    }

    printf("Got public key from peer: \n");
    log_pubkey(srq->pubkey);
    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(client_fd, 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;
    }

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

    send_motd(client_fd, session_key);
    free_session_key(session_key);
    return 0;
}