// The MIT License (MIT)
//
// Copyright (C) 2026 Scott Duensing
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.

// Secure serial link  -- ties rs232, packet, and security into one API
//
// Handshake protocol:
//   1. Both sides send their DH public key (128 bytes) as a single packet
//   2. On receiving the remote key, compute shared secret immediately
//   3. Derive separate TX/RX cipher keys based on public key ordering
//   4. Transition to READY  -- all subsequent encrypted packets use the keys
//
// The handshake uses the packet layer's reliable delivery, so lost packets
// are automatically retransmitted. Both sides can send their public key
// simultaneously  -- there's no initiator/responder distinction.
//
// Directionality: the side with the lexicographically lower public key
// uses master XOR 0xAA for TX and master XOR 0x55 for RX. The other
// side uses the reverse. This is critical: if both sides used the same
// key and counter for CTR mode, they'd produce identical keystreams,
// and XOR'ing two ciphertexts would reveal the XOR of the plaintexts.
// The XOR-derived directional keys ensure each direction has a unique
// keystream even though both start their counters at zero.
//
// Channel header byte: bit 7 = encrypted, bits 6..0 = channel (0-127)

#include <stdint.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "secLink.h"
#include "../rs232/rs232.h"
#include "../packet/packet.h"
#include "../security/security.h"


// ========================================================================
// Internal defines
// ========================================================================

#define STATE_INIT          0
#define STATE_HANDSHAKE     1
#define STATE_READY         2
#define STATE_ERROR         3

#define TX_KEY_XOR          0xAA
#define RX_KEY_XOR          0x55

#define ENCRYPT_FLAG        0x80
#define CHANNEL_MASK        0x7F

// Handshake wait bound.  The peer may need several seconds for its DH
// keypair on slow hardware, but a silent peer (proxy or BBS not running)
// must not hang the single-threaded caller forever.
#define HANDSHAKE_TIMEOUT_MS    30000


// ========================================================================
// Types
// ========================================================================

struct SecLinkS {
    int           com;
    PktConnT     *pkt;
    SecDhT       *dh;
    SecCipherT   *txCipher;
    SecCipherT   *rxCipher;
    SecLinkRecvT  userCallback;
    void         *userCtx;
    int           state;
    uint8_t       myPub[SEC_DH_KEY_SIZE];
    uint8_t       remoteKey[SEC_DH_KEY_SIZE];
    bool          gotRemoteKey;
    bool          txEncryptBusy;
};


// ========================================================================
// Prototypes (alphabetical)
// ========================================================================

static void completeHandshake(SecLinkT *link);
static void internalRecv(void *ctx, const uint8_t *data, int len);
void        secLinkClose(SecLinkT *link);
int         secLinkGetPending(SecLinkT *link);
int         secLinkHandshake(SecLinkT *link);
bool        secLinkIsReady(SecLinkT *link);
SecLinkT   *secLinkOpen(int com, int32_t bps, int dataBits, char parity, int stopBits, int handshake, SecLinkRecvT callback, void *ctx);
int         secLinkPoll(SecLinkT *link);
int         secLinkSend(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt, bool block);
int         secLinkSendBuf(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt);


// ========================================================================
// Static functions (alphabetical)
// ========================================================================

// Called when we've received the remote's public key. Computes the DH
// shared secret, derives directional cipher keys, and transitions to READY.
// After this, the DH context (containing the private key) is destroyed
// immediately  -- forward secrecy principle: even if the long-term state is
// compromised later, past session keys can't be recovered.
static void completeHandshake(SecLinkT *link) {
    uint8_t masterKey[SEC_XTEA_KEY_SIZE];
    uint8_t txKey[SEC_XTEA_KEY_SIZE];
    uint8_t rxKey[SEC_XTEA_KEY_SIZE];
    bool    weAreLower;
    int     rc;

    // An invalid remote key (out of range, e.g. all zeros) must fail the
    // handshake -- ignoring these return codes would leave masterKey as
    // uninitialized stack memory and enter READY with garbage key material.
    rc = secDhComputeSecret(link->dh, link->remoteKey, SEC_DH_KEY_SIZE);
    if (rc != SEC_SUCCESS) {
        secDhDestroy(link->dh);
        link->dh    = 0;
        link->state = STATE_ERROR;
        return;
    }

    rc = secDhDeriveKey(link->dh, masterKey, SEC_XTEA_KEY_SIZE);
    if (rc != SEC_SUCCESS) {
        secDhDestroy(link->dh);
        link->dh    = 0;
        link->state = STATE_ERROR;
        return;
    }

    // Derive directional keys so each side encrypts with a unique key
    weAreLower = (memcmp(link->myPub, link->remoteKey, SEC_DH_KEY_SIZE) < 0);
    for (int i = 0; i < SEC_XTEA_KEY_SIZE; i++) {
        txKey[i] = masterKey[i] ^ (weAreLower ? TX_KEY_XOR : RX_KEY_XOR);
        rxKey[i] = masterKey[i] ^ (weAreLower ? RX_KEY_XOR : TX_KEY_XOR);
    }

    // A repeat handshake replaces the session ciphers. Destroy the previous
    // contexts (securely zeroing their key material) before overwriting the
    // pointers, or every re-handshake would leak both cipher blocks.
    secCipherDestroy(link->txCipher);
    secCipherDestroy(link->rxCipher);

    link->txCipher = secCipherCreate(txKey);
    link->rxCipher = secCipherCreate(rxKey);

    memset(masterKey, 0, sizeof(masterKey));
    memset(txKey, 0, sizeof(txKey));
    memset(rxKey, 0, sizeof(rxKey));

    // Destroy DH context  -- private key no longer needed
    secDhDestroy(link->dh);
    link->dh = 0;

    // If either cipher failed to allocate, abort rather than entering READY
    // with a NULL cipher (which would silently transmit cleartext).
    if (!link->txCipher || !link->rxCipher) {
        secCipherDestroy(link->txCipher);
        link->txCipher = 0;
        secCipherDestroy(link->rxCipher);
        link->rxCipher = 0;
        link->state = STATE_ERROR;
        return;
    }

    link->state = STATE_READY;
}


// Internal packet-layer callback. Routes incoming packets based on state:
//   - During handshake: expects a 128-byte DH public key
//   - When ready: strips the channel header, decrypts if flagged, and
//     forwards plaintext to the user callback
// The channel header is always one byte, so the minimum valid data packet
// is 1 byte (header only, zero-length payload).
static void internalRecv(void *ctx, const uint8_t *data, int len) {
    SecLinkT *link = (SecLinkT *)ctx;

    if (link->state == STATE_HANDSHAKE && !link->gotRemoteKey) {
        // During handshake, expect exactly one 128-byte public key
        if (len == SEC_DH_KEY_SIZE) {
            memcpy(link->remoteKey, data, len);
            link->gotRemoteKey = true;
            completeHandshake(link);
        }
    } else if (link->state == STATE_READY || link->state == STATE_INIT) {
        // Need at least the channel header byte
        if (len < 1) {
            return;
        }

        uint8_t       hdr        = data[0];
        uint8_t       channel    = hdr & CHANNEL_MASK;
        bool          encrypted  = (hdr & ENCRYPT_FLAG) != 0;
        const uint8_t *payload   = data + 1;
        int            payloadLen = len - 1;

        if (encrypted && link->state == STATE_READY) {
            uint8_t plaintext[PKT_MAX_PAYLOAD];
            if (payloadLen > (int)sizeof(plaintext)) {
                payloadLen = (int)sizeof(plaintext);
            }
            memcpy(plaintext, payload, payloadLen);
            secCipherCrypt(link->rxCipher, plaintext, payloadLen);
            if (link->userCallback) {
                link->userCallback(link->userCtx, plaintext, payloadLen, channel);
            }
        } else if (!encrypted) {
            if (link->userCallback) {
                link->userCallback(link->userCtx, payload, payloadLen, channel);
            }
        }
        // encrypted but not READY: silently drop
    }
}


// ========================================================================
// Public functions (alphabetical)
// ========================================================================

void secLinkClose(SecLinkT *link) {
    if (!link) {
        return;
    }

    secCipherDestroy(link->txCipher);
    secCipherDestroy(link->rxCipher);
    secDhDestroy(link->dh);
    if (link->pkt) {
        pktClose(link->pkt);
    }
    rs232Close(link->com);

    memset(link, 0, sizeof(SecLinkT));
    free(link);
}


int secLinkGetPending(SecLinkT *link) {
    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    return pktGetPending(link->pkt);
}


int secLinkHandshake(SecLinkT *link) {
    clock_t start;
    clock_t timeout;
    int     len;
    int     rc;

    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    // Generate DH keypair (idempotent: free any context left by a prior call)
    secDhDestroy(link->dh);
    link->dh = 0;

    link->dh = secDhCreate();
    if (!link->dh) {
        return SECLINK_ERR_ALLOC;
    }

    rc = secDhGenerateKeys(link->dh);
    if (rc != SEC_SUCCESS) {
        secDhDestroy(link->dh);
        link->dh = 0;
        return SECLINK_ERR_HANDSHAKE;
    }

    // Export our public key
    len = SEC_DH_KEY_SIZE;
    secDhGetPublicKey(link->dh, link->myPub, &len);

    // Enter handshake state and send our key
    link->state        = STATE_HANDSHAKE;
    link->gotRemoteKey = false;

    rc = pktSend(link->pkt, link->myPub, SEC_DH_KEY_SIZE, true);
    if (rc != PKT_SUCCESS) {
        secDhDestroy(link->dh);
        link->dh = 0;
        return SECLINK_ERR_HANDSHAKE;
    }

    // Poll until the callback completes the handshake (or signals an error).
    // The wait is bounded: pktPoll only reports PKT_ERR_DISCONNECTED for a
    // closed port, never for an open but silent line, so without a timeout
    // a missing peer would hang the caller forever.
    start   = clock();
    timeout = (clock_t)HANDSHAKE_TIMEOUT_MS * CLOCKS_PER_SEC / 1000;

    while (link->state != STATE_READY && link->state != STATE_ERROR) {
        if (clock() - start >= timeout) {
            secDhDestroy(link->dh);
            link->dh    = 0;
            link->state = STATE_INIT;
            return SECLINK_ERR_TIMEOUT;
        }

        if (pktPoll(link->pkt) == PKT_ERR_DISCONNECTED) {
            secDhDestroy(link->dh);
            link->dh = 0;
            return SECLINK_ERR_HANDSHAKE;
        }
    }

    // completeHandshake sets STATE_ERROR (and frees the DH context) when a
    // cipher could not be created; surface that as a handshake failure.
    if (link->state == STATE_ERROR) {
        return SECLINK_ERR_HANDSHAKE;
    }

    return SECLINK_SUCCESS;
}


bool secLinkIsReady(SecLinkT *link) {
    if (!link) {
        return false;
    }
    return link->state == STATE_READY;
}


SecLinkT *secLinkOpen(int com, int32_t bps, int dataBits, char parity, int stopBits, int handshake, SecLinkRecvT callback, void *ctx) {
    SecLinkT *link;
    int       rc;

    link = (SecLinkT *)calloc(1, sizeof(SecLinkT));
    if (!link) {
        return 0;
    }

    link->com          = com;
    link->userCallback = callback;
    link->userCtx      = ctx;
    link->state        = STATE_INIT;

    // Open serial port
    rc = rs232Open(com, bps, dataBits, parity, stopBits, handshake);
    if (rc != RS232_SUCCESS) {
        free(link);
        return 0;
    }

    // Open packet layer with our internal callback
    link->pkt = pktOpen(com, 0, internalRecv, link);
    if (!link->pkt) {
        rs232Close(com);
        free(link);
        return 0;
    }

    return link;
}


int secLinkPoll(SecLinkT *link) {
    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    return pktPoll(link->pkt);
}


int secLinkSend(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt, bool block) {
    uint8_t buf[PKT_MAX_PAYLOAD];
    int     rc;

    if (!link || !data) {
        return SECLINK_ERR_PARAM;
    }
    if (channel > SECLINK_MAX_CHANNEL) {
        return SECLINK_ERR_PARAM;
    }
    if (encrypt && link->state != STATE_READY) {
        return SECLINK_ERR_NOT_READY;
    }
    if (len <= 0 || len > SECLINK_MAX_PAYLOAD) {
        return SECLINK_ERR_PARAM;
    }

    // A blocking pktSend pumps pktPoll, whose receive callback may legally
    // re-enter secLinkSend. A second encrypted frame must not reach the wire
    // before the outer frame that was already encrypted -- the receiver
    // decrypts in arrival order, so the reorder would desynchronize the
    // XTEA-CTR keystream. Refuse the re-entrant encrypted send instead.
    if (encrypt && link->txEncryptBusy) {
        return SECLINK_ERR_SEND;
    }

    // CRITICAL: check window space BEFORE encrypting. If we encrypted first
    // and then the send failed, the cipher counter would have advanced but
    // the data wouldn't have been sent, permanently desynchronizing the
    // TX cipher state from the remote's RX cipher state.
    if (!block && !pktCanSend(link->pkt)) {
        return SECLINK_ERR_SEND;
    }

    // Build channel header byte
    buf[0] = channel & CHANNEL_MASK;
    if (encrypt) {
        buf[0] |= ENCRYPT_FLAG;
    }

    // Copy payload after header
    memcpy(buf + SECLINK_CHAN_HDR_SIZE, data, len);

    // Encrypt the payload portion only (not the header). The busy flag
    // covers the window from counter advance until the frame is queued.
    if (encrypt) {
        link->txEncryptBusy = true;
        secCipherCrypt(link->txCipher, buf + SECLINK_CHAN_HDR_SIZE, len);
    }

    rc = pktSend(link->pkt, buf, len + SECLINK_CHAN_HDR_SIZE, block);

    if (encrypt) {
        link->txEncryptBusy = false;
    }

    if (rc != PKT_SUCCESS) {
        return SECLINK_ERR_SEND;
    }

    return SECLINK_SUCCESS;
}


// Convenience function for sending large buffers. Splits the data into
// SECLINK_MAX_PAYLOAD (254 byte) chunks and sends each one, blocking until
// the send window has room. The receiver sees multiple packets on the same
// channel and must reassemble if needed. Always blocking because the caller
// expects the entire buffer to be sent when this returns.
int secLinkSendBuf(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt) {
    int offset = 0;

    if (!link || !data) {
        return SECLINK_ERR_PARAM;
    }
    if (len <= 0) {
        return SECLINK_ERR_PARAM;
    }

    while (offset < len) {
        int chunk = len - offset;
        int rc;

        if (chunk > SECLINK_MAX_PAYLOAD) {
            chunk = SECLINK_MAX_PAYLOAD;
        }

        rc = secLinkSend(link, data + offset, chunk, channel, encrypt, true);
        if (rc != SECLINK_SUCCESS) {
            return rc;
        }

        offset += chunk;
    }

    return SECLINK_SUCCESS;
}
