#!/usr/bin/env python3
"""Minimal BTF parser: print struct layouts (with anonymous member expansion)."""
import struct, sys

ELF = "/home/pwn/ctf/pwn3/ubuntu-kernel/vmlinux-39.elf"

def get_section(path, name):
    data = open(path, "rb").read()
    e_shoff = struct.unpack_from("<Q", data, 0x28)[0]
    e_shentsize = struct.unpack_from("<H", data, 0x3A)[0]
    e_shnum = struct.unpack_from("<H", data, 0x3C)[0]
    e_shstrndx = struct.unpack_from("<H", data, 0x3E)[0]
    shstr_off = struct.unpack_from("<Q", data, e_shoff + e_shstrndx * e_shentsize + 0x18)[0]
    for i in range(e_shnum):
        off = e_shoff + i * e_shentsize
        sh_name = struct.unpack_from("<I", data, off)[0]
        sh_offset = struct.unpack_from("<Q", data, off + 0x18)[0]
        sh_size = struct.unpack_from("<Q", data, off + 0x20)[0]
        nm = data[shstr_off + sh_name:data.find(b"\0", shstr_off + sh_name)].decode()
        if nm == name:
            return data[sh_offset:sh_offset + sh_size]
    raise KeyError(name)

btf = get_section(ELF, ".BTF")
magic, version, flags, hdr_len = struct.unpack_from("<HBBI", btf, 0)
type_off, type_len, str_off, str_len = struct.unpack_from("<IIII", btf, 8)
assert magic == 0xEB9F
td = btf[hdr_len + type_off: hdr_len + type_off + type_len]
strs = btf[hdr_len + str_off: hdr_len + str_off + str_len]

def s(off):
    if off == 0:
        return ""
    end = strs.find(b"\0", off)
    return strs[off:end].decode(errors="replace")

KIND = {1:"INT",2:"PTR",3:"ARRAY",4:"STRUCT",5:"UNION",6:"ENUM",7:"FWD",8:"TYPEDEF",
        9:"VOLATILE",10:"CONST",11:"RESTRICT",12:"FUNC",13:"FUNC_PROTO",14:"VAR",
        15:"DATASEC",16:"FLOAT",17:"DECL_TAG",18:"TYPE_TAG",19:"ENUM64"}

types = [None]
pos = 0
while pos < len(td):
    name_off, info, size_type = struct.unpack_from("<III", td, pos); pos += 12
    vlen = info & 0xffff; kind = (info >> 24) & 0x1f; kflag = (info >> 31) & 1
    extra = None
    if kind in (4, 5):
        extra = []
        for _ in range(vlen):
            mn, mt, mo = struct.unpack_from("<III", td, pos); pos += 12
            extra.append((s(mn), mt, mo))
    elif kind == 1:
        extra = struct.unpack_from("<I", td, pos)[0]; pos += 4
    elif kind == 3:
        extra = struct.unpack_from("<III", td, pos); pos += 12
    elif kind == 6:
        pos += vlen * 8
    elif kind == 19:
        pos += vlen * 12
    elif kind == 13:
        pos += vlen * 8
    elif kind == 14:
        pos += 4
    elif kind == 15:
        pos += vlen * 12
    elif kind == 17:
        pos += 8
    types.append({"name": s(name_off), "kind": kind, "vlen": vlen,
                  "size_type": size_type, "kflag": kflag, "extra": extra})

def resolve(tid):
    for _ in range(30):
        t = types[tid]
        if t["kind"] in (8, 9, 10, 11, 18):
            tid = t["size_type"]
        else:
            return tid
    return tid

def find(name):
    for i, t in enumerate(types):
        if i and t and t["kind"] in (4, 5) and t["name"] == name:
            return i
    return None

def walk(tid, base=0, depth=0):
    tid = resolve(tid); t = types[tid]
    print(f"{'  '*depth}{t['name'] or '<anon>'} [{KIND[t['kind']]}] size={t['size_type']:#x} base={base:#x}")
    if t["kind"] not in (4, 5):
        return
    for mn, mt, mo in t["extra"]:
        bit = mo & 0xffffff if t["kflag"] else mo
        off = base + bit // 8
        rt = resolve(mt); tt = types[rt]
        print(f"{'  '*depth}  {mn or '<anon>':30s} off={off:#06x} type={rt} {KIND[tt['kind']]} {tt['name']}")
        if (not mn) and tt["kind"] in (4, 5) and depth < 3:
            walk(rt, off, depth + 1)
        elif mn in ("it", "w") and tt["kind"] in (4, 5) and depth < 3:
            walk(rt, off, depth + 1)

for nm in ("k_itimer", "cpu_timer", "timerqueue_node", "timerqueue_head",
           "posix_cputimers", "posix_cputimer_base", "cred"):
    i = find(nm)
    print(f"\n===== {nm} id={i}")
    if i:
        walk(i)

print("\n===== task_struct fields")
i = find("task_struct")
t = types[i]
want = {"posix_cputimers", "cred", "real_cred", "splice_pipe", "fs", "files", "comm", "pid", "tasks", "thread_info"}
for mn, mt, mo in t["extra"]:
    if mn in want:
        rt = resolve(mt); tt = types[rt]
        print(f"  {mn:24s} off={mo//8:#x} type={rt} {KIND[tt['kind']]} {tt['name']} size={tt['size_type']:#x}")
print("task_struct size:", hex(t["size_type"]))