#!/usr/bin/env python3
"""
Restore ext4 metadata (permissions, ownership, timestamps) to a recovered tree.
Run after dump_tree.py has extracted files.

Usage: python3 restore_meta.py <inode> <dest_dir>
"""
import struct, os, sys, stat, ctypes, ctypes.util

DEV             = '/dev/dm-0'
BLOCK           = 4096
BACKUP_SB_BLOCK = 32768

# ── reuse same low-level helpers from dump_tree.py ───────────────────────────
# (paste parse_superblock, parse_gdt_entry, read_at, read_inode,
#  read_extent_tree_blocks, read_dir_entries here)
# or factor them into a shared ext4lib.py and import from both scripts

import ext4lib 

libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)

class Timeval(ctypes.Structure):
    _fields_ = [('tv_sec', ctypes.c_long), ('tv_usec', ctypes.c_long)]

def lutimes(path, atime, mtime):
    times = (Timeval * 2)((atime, 0), (mtime, 0))
    libc.lutimes(path.encode(), ctypes.byref(times))

def get_inode_meta(idata, slot, sb):
    mode  = struct.unpack_from('<H', idata, slot +   0)[0]
    uid   = struct.unpack_from('<H', idata, slot +   2)[0]
    gid   = struct.unpack_from('<H', idata, slot +  24)[0]
    atime = struct.unpack_from('<I', idata, slot +   8)[0]
    mtime = struct.unpack_from('<I', idata, slot +  16)[0]

    if sb['inode_size'] >= 256:
        atime_extra = struct.unpack_from('<I', idata, slot + 132)[0]
        mtime_extra = struct.unpack_from('<I', idata, slot + 140)[0]
        atime |= (atime_extra & 0x3) << 32
        mtime |= (mtime_extra & 0x3) << 32

    uid_hi, gid_hi = struct.unpack_from('<HH', idata, slot + 120)
    uid |= uid_hi << 16
    gid |= gid_hi << 16

    return stat.S_IMODE(mode), uid, gid, atime, mtime

def restore_meta(f, sb, gdt_data, inum, dest_path):
    try:
        idata, slot = ext4lib.read_inode(f, sb, gdt_data, inum)
        mode, uid, gid, atime, mtime = get_inode_meta(idata, slot, sb)
        is_symlink = os.path.islink(dest_path)

        try:
            os.lchown(dest_path, uid, gid)
        except OSError:
            pass

        if not is_symlink:
            try:
                os.chmod(dest_path, mode)
            except OSError:
                pass

        try:
            lutimes(dest_path, atime, mtime)
        except Exception:
            pass

    except Exception as e:
        print(f"  WARN {dest_path}: {e}", file=sys.stderr)

def walk_and_restore(f, sb, gdt_data, inum, dest_dir, visited=None):
    if visited is None:
        visited = set()
    if inum in visited:
        return
    visited.add(inum)

    # Restore the directory itself
    restore_meta(f, sb, gdt_data, inum, dest_dir)

    try:
        entries = ext4lib.read_dir_entries(f, sb, gdt_data, inum)
    except Exception:
        return

    for name, (child_inum, ftype) in entries.items():
        if name in ('.', '..'):
            continue
        safe_name = name.replace('/', '_').replace('\x00', '')
        dest = os.path.join(dest_dir, safe_name)

        if not os.path.lexists(dest):
            # File wasn't recovered - skip
            continue

        if os.path.isdir(dest) and not os.path.islink(dest):
            walk_and_restore(f, sb, gdt_data, child_inum, dest, visited)
        else:
            restore_meta(f, sb, gdt_data, child_inum, dest)

    # Restore directory timestamps AFTER processing children
    # (writing children updates parent dir mtime/atime)
    restore_meta(f, sb, gdt_data, inum, dest_dir)

def main():
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <inode> <dest_dir>")
        sys.exit(1)

    root_inum = int(sys.argv[1])
    dest_dir  = sys.argv[2]

    with open(DEV, 'rb') as f:
        sb_data  = ext4lib.read_at(f, BACKUP_SB_BLOCK * BLOCK, 1024)
        sb       = ext4lib.parse_superblock(sb_data)
        assert sb['magic'] == 0xef53

        num_groups = (sb['blocks_count'] + sb['blocks_per_group'] - 1) \
                     // sb['blocks_per_group']
        gdt_data = ext4lib.read_at(f, (BACKUP_SB_BLOCK + 1) * BLOCK,
                           num_groups * sb['desc_size'])

        print(f"Restoring metadata: inode {root_inum} -> {dest_dir}")
        walk_and_restore(f, sb, gdt_data, root_inum, dest_dir)
        print("Done")

if __name__ == '__main__':
    main()