/*
 * 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 <net/stpsession.h>
#include <crypto/aes.h>
#include <arpa/inet.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

static int
decrypt_frame(const struct msg_frame *frame, const unsigned char *key, size_t len,
    void *res)
{
    struct aes_message am;
    unsigned char *tmp;
    int error;

    /* Construct an AES message with the message frame */
    am.ciphertext = (unsigned char *)frame->payload;
    am.ciphertext_len = frame->len;
    memcpy(am.tag, frame->tag, AES_GCM_TAG_SIZE);
    memcpy(am.iv, frame->iv, AES_IV_SIZE);
    memcpy(am.ciphertext, frame->payload, frame->len);

    /* Decrypt and deserialize the frame payload */
    if ((error = aes256_decrypt(&am, key, &tmp)) < 0) {
        return error;
    }
    memcpy(res, tmp, len);
    free(tmp);
    return 0;
}

int
send_frame(int client_fd, void *data, size_t len, const unsigned char *session_key)
{
    struct aes_message am;
    struct msg_frame msg_frame;
    int error;

    error = aes256_encrypt((void *)data, len, session_key, &am);
    if (error < 0) {
        printf("send_frame: Failed to encrypt data\n");
        return error;
    }

    memcpy(msg_frame.iv, am.iv, AES_IV_SIZE);
    memcpy(msg_frame.tag, am.tag, AES_GCM_TAG_SIZE);
    memcpy(msg_frame.payload, am.ciphertext, am.ciphertext_len);
    msg_frame.len = am.ciphertext_len;

    /* Send in the message frame */
    error = send(client_fd, &msg_frame, sizeof(msg_frame), 0);
    if (error < 0) {
        printf("send_frame: Failed sending frame to server\n");
        aes256_free_msg(&am);
        return error;
    }
    return 0;
}

int
recv_frame(int client_fd, size_t len, const unsigned char *session_key, void *res)
{
    struct msg_frame frame;
    int error, ret;

    /* Get the message frame over the network */
    ret = recv(client_fd, &frame, sizeof(struct msg_frame), MSG_WAITALL);
    if (ret < 0) {
        printf("Failed to get frame payload from peer\n");
        return error;
    }
    if (ret == 0) {
        printf("Connection closed by peer\n");
        return -1;
    }

    /* Decrypt the payload */
    if ((error = decrypt_frame(&frame, session_key, len, res)) < 0) {
        printf("Failed to decrypt frame payload, rejecting...\n");
        return error;
    }

    return (ret < len) ? ret : len;
}