// SecLink proxy — bridges an 86Box serial connection to a telnet BBS
//
// 86Box (DOS terminal) ←→ TCP ←→ proxy ←→ TCP ←→ BBS
//                         secLink protocol       plain telnet
//
// Usage: proxy [listen_port] [bbs_host] [bbs_port]
// Defaults: 2323 10.1.0.244 2023

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <poll.h>
#include <fcntl.h>

#include "sockShim.h"
#include "../seclink/secLink.h"
#include "../security/security.h"


// ========================================================================
// Defines
// ========================================================================

#define DEFAULT_LISTEN_PORT     2323
#define DEFAULT_BBS_HOST        "10.1.0.244"
#define DEFAULT_BBS_PORT        2023

#define CHANNEL_TERMINAL        0
#define POLL_TIMEOUT_MS         10

// Telnet protocol bytes
#define TEL_IAC                 255
#define TEL_DONT                254
#define TEL_DO                  253
#define TEL_WONT                252
#define TEL_WILL                251
#define TEL_SB                  250
#define TEL_SE                  240

// Telnet options we accept
#define TELOPT_ECHO             1
#define TELOPT_SGA              3
#define TELOPT_TTYPE            24
#define TELOPT_NAWS             31

// Telnet parser states
#define TS_DATA                 0
#define TS_IAC                  1
#define TS_WILL                 2
#define TS_WONT                 3
#define TS_DO                   4
#define TS_DONT                 5
#define TS_SB                   6
#define TS_SB_IAC               7


// ========================================================================
// Static globals
// ========================================================================

static volatile bool sRunning  = true;
static volatile bool sGotEnter = false;
static int           sTelState = TS_DATA;
static int           sClientFd = -1;


// ========================================================================
// Static prototypes (alphabetical)
// ========================================================================

static int  connectToBbs(const char *host, int port);
static int  createListenSocket(int port);
static void onRecvFromDos(void *ctx, const uint8_t *data, int len, uint8_t channel);
static void seedRng(void);
static void sigHandler(int sig);
static void telnetRespond(int bbsFd, uint8_t cmd, uint8_t opt);
static int  telnetFilter(int bbsFd, const uint8_t *in, int inLen, uint8_t *out);


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

static int connectToBbs(const char *host, int port) {
    struct addrinfo  hints;
    struct addrinfo *res;
    char             portStr[16];
    int              fd;

    memset(&hints, 0, sizeof(hints));
    hints.ai_family   = AF_INET;
    hints.ai_socktype = SOCK_STREAM;

    snprintf(portStr, sizeof(portStr), "%d", port);

    if (getaddrinfo(host, portStr, &hints, &res) != 0) {
        return -1;
    }

    fd = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (fd < 0) {
        freeaddrinfo(res);
        return -1;
    }

    if (connect(fd, res->ai_addr, res->ai_addrlen) < 0) {
        close(fd);
        freeaddrinfo(res);
        return -1;
    }

    freeaddrinfo(res);
    return fd;
}


static int createListenSocket(int port) {
    struct sockaddr_in addr;
    int                fd;
    int                opt = 1;

    fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }

    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    memset(&addr, 0, sizeof(addr));
    addr.sin_family      = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port        = htons(port);

    if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        close(fd);
        return -1;
    }

    if (listen(fd, 1) < 0) {
        close(fd);
        return -1;
    }

    return fd;
}


static void onRecvFromDos(void *ctx, const uint8_t *data, int len, uint8_t channel) {
    int bbsFd = *(int *)ctx;

    (void)channel;

    // Check for ENTER before BBS is connected
    if (!sGotEnter) {
        for (int i = 0; i < len; i++) {
            if (data[i] == '\r' || data[i] == '\n') {
                sGotEnter = true;
                printf("Got ENTER from terminal.\n");
                break;
            }
        }
        return;
    }

    int sent = 0;
    while (sent < len) {
        ssize_t n = write(bbsFd, data + sent, len - sent);
        if (n <= 0) {
            break;
        }
        sent += (int)n;
    }
}


static void seedRng(void) {
    uint8_t entropy[32];
    FILE   *f;

    f = fopen("/dev/urandom", "rb");
    if (f) {
        if (fread(entropy, 1, sizeof(entropy), f) < sizeof(entropy)) {
            fprintf(stderr, "Warning: short read from /dev/urandom\n");
        }
        fclose(f);
    }

    secRngSeed(entropy, sizeof(entropy));
}


static void sigHandler(int sig) {
    (void)sig;
    sRunning = false;
    // Shutdown the client socket to break any blocking reads in the packet layer
    if (sClientFd >= 0) {
        shutdown(sClientFd, SHUT_RDWR);
    }
}


// Send a telnet negotiation response
static void telnetRespond(int bbsFd, uint8_t cmd, uint8_t opt) {
    uint8_t resp[3] = {TEL_IAC, cmd, opt};
    ssize_t n = write(bbsFd, resp, 3);
    (void)n;
    printf("  TEL TX: %s %d\n",
           cmd == TEL_WILL ? "WILL" : cmd == TEL_WONT ? "WONT" :
           cmd == TEL_DO ? "DO" : "DONT", opt);
}


// Filter telnet IAC sequences from BBS data.
// Handles negotiation by responding appropriately.
// Returns the number of clean data bytes written to 'out'.
static int telnetFilter(int bbsFd, const uint8_t *in, int inLen, uint8_t *out) {
    int outLen = 0;

    for (int i = 0; i < inLen; i++) {
        uint8_t b = in[i];

        switch (sTelState) {
            case TS_DATA:
                if (b == TEL_IAC) {
                    sTelState = TS_IAC;
                } else {
                    out[outLen++] = b;
                }
                break;

            case TS_IAC:
                switch (b) {
                    case TEL_IAC:
                        // Escaped 0xFF — emit literal
                        out[outLen++] = 0xFF;
                        sTelState = TS_DATA;
                        break;
                    case TEL_WILL:
                        sTelState = TS_WILL;
                        break;
                    case TEL_WONT:
                        sTelState = TS_WONT;
                        break;
                    case TEL_DO:
                        sTelState = TS_DO;
                        break;
                    case TEL_DONT:
                        sTelState = TS_DONT;
                        break;
                    case TEL_SB:
                        sTelState = TS_SB;
                        break;
                    default:
                        // Unknown command, skip
                        sTelState = TS_DATA;
                        break;
                }
                break;

            case TS_WILL:
                // Server offers to do something — accept ECHO and SGA, refuse others
                printf("  TEL RX: WILL %d\n", b);
                if (b == TELOPT_ECHO || b == TELOPT_SGA) {
                    telnetRespond(bbsFd, TEL_DO, b);
                } else {
                    telnetRespond(bbsFd, TEL_DONT, b);
                }
                sTelState = TS_DATA;
                break;

            case TS_WONT:
                printf("  TEL RX: WONT %d\n", b);
                telnetRespond(bbsFd, TEL_DONT, b);
                sTelState = TS_DATA;
                break;

            case TS_DO:
                // Server asks us to do something — accept TTYPE and NAWS, refuse others
                printf("  TEL RX: DO %d\n", b);
                if (b == TELOPT_TTYPE || b == TELOPT_NAWS) {
                    telnetRespond(bbsFd, TEL_WILL, b);
                } else {
                    telnetRespond(bbsFd, TEL_WONT, b);
                }
                sTelState = TS_DATA;
                break;

            case TS_DONT:
                printf("  TEL RX: DONT %d\n", b);
                telnetRespond(bbsFd, TEL_WONT, b);
                sTelState = TS_DATA;
                break;

            case TS_SB:
                // Inside subnegotiation — skip until IAC SE
                if (b == TEL_IAC) {
                    sTelState = TS_SB_IAC;
                }
                break;

            case TS_SB_IAC:
                if (b == TEL_SE) {
                    sTelState = TS_DATA;
                } else {
                    sTelState = TS_SB;
                }
                break;
        }
    }

    return outLen;
}


// ========================================================================
// Main
// ========================================================================

int main(int argc, char *argv[]) {
    int              listenPort = DEFAULT_LISTEN_PORT;
    const char      *bbsHost   = DEFAULT_BBS_HOST;
    int              bbsPort   = DEFAULT_BBS_PORT;
    int              listenFd;
    int              clientFd;
    int              bbsFd;
    SecLinkT        *link;
    struct sockaddr_in clientAddr;
    socklen_t        clientLen;
    struct pollfd    fds[2];
    int              rc;

    if (argc > 1 && (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0)) {
        printf("Usage: %s [listen_port] [bbs_host] [bbs_port]\n", argv[0]);
        printf("Defaults: %d %s %d\n", DEFAULT_LISTEN_PORT, DEFAULT_BBS_HOST, DEFAULT_BBS_PORT);
        return 0;
    }

    if (argc > 1) {
        listenPort = atoi(argv[1]);
    }
    if (argc > 2) {
        bbsHost = argv[2];
    }
    if (argc > 3) {
        bbsPort = atoi(argv[3]);
    }

    signal(SIGPIPE, SIG_IGN);
    struct sigaction sa;
    sa.sa_handler = sigHandler;
    sa.sa_flags   = 0;
    sigemptyset(&sa.sa_mask);
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);

    // Listen for 86Box connection
    listenFd = createListenSocket(listenPort);
    if (listenFd < 0) {
        fprintf(stderr, "Failed to listen on port %d: %s\n", listenPort, strerror(errno));
        return 1;
    }
    printf("Listening on port %d...\n", listenPort);

    // Accept connection from 86Box (poll so Ctrl+C works)
    struct pollfd listenPoll = {listenFd, POLLIN, 0};
    while (sRunning) {
        int pr = poll(&listenPoll, 1, 500);
        if (pr > 0) {
            break;
        }
    }
    clientLen = sizeof(clientAddr);
    clientFd  = accept(listenFd, (struct sockaddr *)&clientAddr, &clientLen);
    close(listenFd);
    if (!sRunning || clientFd < 0) {
        if (sRunning) {
            fprintf(stderr, "Accept failed: %s\n", strerror(errno));
        }
        return sRunning ? 1 : 0;
    }
    printf("86Box connected.\n");
    sClientFd = clientFd;

    // Associate socket with COM0 for the secLink stack
    sockShimSetFd(0, clientFd);

    // Seed RNG from /dev/urandom and open secLink
    seedRng();
    link = secLinkOpen(0, 115200, 8, 'N', 1, 0, onRecvFromDos, &bbsFd);
    if (!link) {
        fprintf(stderr, "Failed to open secLink.\n");
        close(clientFd);
        return 1;
    }

    // DH key exchange (blocks until both sides complete)
    printf("Waiting for secLink handshake...\n");
    rc = secLinkHandshake(link);
    if (rc != SECLINK_SUCCESS) {
        fprintf(stderr, "Handshake failed: %d\n", rc);
        secLinkClose(link);
        return 1;
    }
    printf("Handshake complete.\n");

    // Wait for ENTER from terminal before connecting to BBS
    printf("Waiting for terminal to send ENTER...\n");
    while (sRunning) {
        secLinkPoll(link);
        if (sGotEnter) {
            break;
        }
        usleep(10000);
    }
    if (!sRunning) {
        secLinkClose(link);
        return 0;
    }

    // Send test string to verify data path
    const char *testMsg = "SecLink proxy ready.\r\n";
    secLinkSend(link, (const uint8_t *)testMsg, (int)strlen(testMsg), CHANNEL_TERMINAL, true, false);
    printf("Sent test message to terminal.\n");

    // Now connect to BBS
    printf("Connecting to %s:%d...\n", bbsHost, bbsPort);
    bbsFd = connectToBbs(bbsHost, bbsPort);
    if (bbsFd < 0) {
        fprintf(stderr, "Failed to connect to BBS: %s\n", strerror(errno));
        secLinkClose(link);
        return 1;
    }
    printf("BBS connected. Proxying traffic.\n");

    // Set BBS socket non-blocking for the main loop
    int flags = fcntl(bbsFd, F_GETFL, 0);
    fcntl(bbsFd, F_SETFL, flags | O_NONBLOCK);

    // Main proxy loop
    fds[0].fd     = clientFd;
    fds[0].events = POLLIN;
    fds[1].fd     = bbsFd;
    fds[1].events = POLLIN;

    while (sRunning) {
        poll(fds, 2, POLL_TIMEOUT_MS);

        // Process incoming secLink packets from 86Box
        // (callback forwards decrypted data to BBS)
        secLinkPoll(link);

        // Read from BBS, filter telnet, send clean data to 86Box
        if (fds[1].revents & POLLIN) {
            uint8_t raw[64];
            uint8_t clean[64];
            ssize_t n = read(bbsFd, raw, sizeof(raw));
            if (n <= 0) {
                printf("BBS disconnected.\n");
                break;
            }

            int cleanLen = telnetFilter(bbsFd, raw, (int)n, clean);
            if (cleanLen > 0) {
                // Retry with ACK processing if the send window is full
                int rc = secLinkSend(link, clean, cleanLen, CHANNEL_TERMINAL, true, false);
                while (rc != SECLINK_SUCCESS && sRunning) {
                    secLinkPoll(link);
                    usleep(1000);
                    rc = secLinkSend(link, clean, cleanLen, CHANNEL_TERMINAL, true, false);
                }
            }
        }

        // Check for disconnects
        if (fds[0].revents & (POLLERR | POLLHUP)) {
            printf("86Box disconnected.\n");
            break;
        }
        if (fds[1].revents & (POLLERR | POLLHUP)) {
            printf("BBS disconnected.\n");
            break;
        }
    }

    printf("Shutting down.\n");
    secLinkClose(link);
    close(bbsFd);
    close(clientFd);

    return 0;
}