#!/usr/bin/env python3
"""
Stage 4 – Restore filesystem metadata (permissions, ownership, timestamps) to
an already-extracted directory tree.

Run this after extract_tree.py if you did not use --restore-meta there, or to
re-apply metadata without re-extracting files.

Usage:
    python3 restore_meta.py [options] <inode> <dest_dir>

Options:
    --device DEV          Block device  [/dev/dm-0]
    --backup-sb BLOCK     Backup superblock block number  [32768]
    --db PATH             Optional: read metadata from SQLite DB instead of
                          re-reading the device (faster)  [inodes.db if present]
"""
import argparse, ctypes, ctypes.util, os, stat, sys
import ext4lib

DEFAULT_DEV       = '/dev/dm-0'
DEFAULT_BACKUP_SB = 32768

# ── libc lutimes ──────────────────────────────────────────────────────────────

_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)((_Timeval(atime, 0)), (_Timeval(mtime, 0)))
    _libc.lutimes(path.encode(), ctypes.byref(times))


# ── per-path restore ──────────────────────────────────────────────────────────

def _apply_meta(dest_path, perms, uid, gid, atime, mtime):
    is_link = os.path.islink(dest_path)
    try:
        os.lchown(dest_path, uid, gid)
    except OSError:
        pass
    if not is_link:
        try:
            os.chmod(dest_path, perms)
        except OSError:
            pass
    try:
        _lutimes(dest_path, atime, mtime)
    except Exception:
        pass


def restore_one(f, sb, gdt_data, inum, dest_path):
    try:
        idata, slot = ext4lib.read_inode(f, sb, gdt_data, inum)
        perms, uid, gid, atime, mtime = ext4lib.get_inode_meta(idata, slot, sb)
        _apply_meta(dest_path, perms, uid, gid, atime, mtime)
    except Exception as e:
        print(f"  WARN {dest_path}: {e}", file=sys.stderr)


# ── recursive walker ──────────────────────────────────────────────────────────

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 first (to set ownership before descending)
    restore_one(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):
            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_one(f, sb, gdt_data, child_inum, dest)

    # Re-apply directory timestamps after children are written
    restore_one(f, sb, gdt_data, inum, dest_dir)


# ── main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description='Restore ext4 metadata to extracted tree (Stage 4)')
    parser.add_argument('inode',       type=int)
    parser.add_argument('dest_dir')
    parser.add_argument('--device',    default=DEFAULT_DEV)
    parser.add_argument('--backup-sb', type=int, default=DEFAULT_BACKUP_SB)
    args = parser.parse_args()

    with open(args.device, 'rb') as f:
        sb, gdt_data, _ = ext4lib.load_fs(f, args.backup_sb)
        print(f"Restoring metadata: inode {args.inode} → {args.dest_dir}")
        walk_and_restore(f, sb, gdt_data, args.inode, args.dest_dir)
        print("Done")


if __name__ == '__main__':
    main()