import gdb

import pwndbg.gdblib.memory
import pwndbg.heap
import pwndbg.symbol
import pwndbg.vmmap
import tests

BINARY = tests.binaries.get("heap_bins.out")


def test_heap_bins(start_binary):
    """
    Tests pwndbg heap bins commands
    """
    start_binary(BINARY)
    gdb.execute("set context-output /dev/null")
    gdb.execute("b breakpoint", to_string=True)

    # check if all bins are empty at first
    gdb.execute("continue")
    allocator = pwndbg.heap.current

    addr = pwndbg.symbol.address("tcache_size")
    tcache_size = allocator._request2size(pwndbg.gdblib.memory.u64(addr))
    addr = pwndbg.symbol.address("tcache_count")
    tcache_count = pwndbg.gdblib.memory.u64(addr)
    addr = pwndbg.symbol.address("fastbin_size")
    fastbin_size = allocator._request2size(pwndbg.gdblib.memory.u64(addr))
    addr = pwndbg.symbol.address("fastbin_count")
    fastbin_count = pwndbg.gdblib.memory.u64(addr)
    addr = pwndbg.symbol.address("smallbin_size")
    smallbin_size = allocator._request2size(pwndbg.gdblib.memory.u64(addr))
    addr = pwndbg.symbol.address("smallbin_count")
    smallbin_count = pwndbg.gdblib.memory.u64(addr)
    addr = pwndbg.symbol.address("largebin_size")
    largebin_size = allocator._request2size(pwndbg.gdblib.memory.u64(addr))
    addr = pwndbg.symbol.address("largebin_count")
    largebin_count = pwndbg.gdblib.memory.u64(addr)

    result = allocator.tcachebins()
    assert result["type"] == "tcachebins"
    assert tcache_size in result
    assert result[tcache_size][1] == 0 and len(result[tcache_size][0]) == 1

    result = allocator.fastbins()
    assert result["type"] == "fastbins"
    assert fastbin_size in result
    assert len(result[fastbin_size]) == 1

    result = allocator.unsortedbin()
    assert result["type"] == "unsortedbin"
    assert len(result["all"][0]) == 1
    assert not result["all"][2]

    result = allocator.smallbins()
    assert result["type"] == "smallbins"
    assert smallbin_size in result
    assert len(result[smallbin_size][0]) == 1 and len(result[smallbin_size][1]) == 1
    assert not result[smallbin_size][2]

    result = allocator.largebins()
    assert result["type"] == "largebins"
    largebin_size = list(result.items())[allocator.largebin_index(largebin_size) - 64][0]
    assert largebin_size in result
    assert len(result[largebin_size][0]) == 1 and len(result[largebin_size][1]) == 1
    assert not result[largebin_size][2]

    # check tcache
    gdb.execute("continue")

    result = allocator.tcachebins()
    assert result["type"] == "tcachebins"
    assert tcache_size in result
    assert (
        result[tcache_size][1] == tcache_count and len(result[tcache_size][0]) == tcache_count + 1
    )
    for addr in result[tcache_size][0][:-1]:
        assert pwndbg.vmmap.find(addr)

    # check fastbin
    gdb.execute("continue")

    result = allocator.fastbins()
    assert result["type"] == "fastbins"
    assert (fastbin_size in result) and (len(result[fastbin_size]) == fastbin_count + 1)
    for addr in result[fastbin_size][:-1]:
        assert pwndbg.vmmap.find(addr)

    # check unsortedbin
    gdb.execute("continue")

    result = allocator.unsortedbin()
    assert result["type"] == "unsortedbin"
    assert (
        len(result["all"][0]) == smallbin_count + 2 and len(result["all"][1]) == smallbin_count + 2
    )
    assert not result["all"][2]
    for addr in result["all"][0][:-1]:
        assert pwndbg.vmmap.find(addr)
    for addr in result["all"][1][:-1]:
        assert pwndbg.vmmap.find(addr)

    # check smallbins
    gdb.execute("continue")

    result = allocator.smallbins()
    assert result["type"] == "smallbins"
    assert (
        len(result[smallbin_size][0]) == smallbin_count + 2
        and len(result[smallbin_size][1]) == smallbin_count + 2
    )
    assert not result[smallbin_size][2]
    for addr in result[smallbin_size][0][:-1]:
        assert pwndbg.vmmap.find(addr)
    for addr in result[smallbin_size][1][:-1]:
        assert pwndbg.vmmap.find(addr)

    # check largebins
    gdb.execute("continue")

    result = allocator.largebins()
    assert result["type"] == "largebins"
    assert (
        len(result[largebin_size][0]) == largebin_count + 2
        and len(result[largebin_size][1]) == largebin_count + 2
    )
    assert not result[largebin_size][2]
    for addr in result[largebin_size][0][:-1]:
        assert pwndbg.vmmap.find(addr)
    for addr in result[largebin_size][1][:-1]:
        assert pwndbg.vmmap.find(addr)

    # check corrupted
    gdb.execute("continue")
    result = allocator.smallbins()
    assert result["type"] == "smallbins"
    assert result[smallbin_size][2]

    result = allocator.largebins()
    assert result["type"] == "largebins"
    assert result[largebin_size][2]