// Security-context smoke test. This validates authority attenuation, assignment, and ambient privilege enforcement.

#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <ntlite/syscall.h>
#include <ntlite/path.h>
#include <ntlite/process_loader.h>

static void print(const char* text) {
    SYS_WRITE_CONSOLE_REQUEST request = { sizeof(request), SysSvFromCstr(text) };
    (void)SysWriteConsole(&request);
}

static void print_u32(uint32_t value) {
    char digits[16];
    uint32_t count = 0;
    if (value == 0) {
        print("0");
        return;
    }
    while (value && count < sizeof(digits)) {
        digits[count++] = (char)('0' + value % 10u);
        value /= 10u;
    }
    while (count) {
        char one[2] = { digits[--count], 0 };
        print(one);
    }
}

// Exits without relying on a C runtime return path.
static void exit_now(int32_t code) {
    while (true) (void)SysExitProcess(code);
}

// Fails the privileged parent with a visible diagnostic.
static void fail_status(const char* operation, KSTATUS status) {
    print("security_smoke: ");
    print(operation);
    print(" status=");
    print_u32(status);
    print("\n");
    exit_now(1);
}

// Parses one unsigned decimal value without requiring formatted libc helpers.
static bool parse_u32(SYS_STRING_VIEW text, size_t start, uint32_t* valueOut) {
    if (!valueOut || start >= text.Length) return false;
    uint32_t value = 0;
    bool any = false;
    for (size_t i = start; i < text.Length; i++) {
        char c = text.Data[i];
        if (c < '0' || c > '9') return false;
        uint32_t digit = (uint32_t)(c - '0');
        if (value > (UINT32_MAX - digit) / 10u) return false;
        value = value * 10u + digit;
        any = true;
    }
    if (!any) return false;
    *valueOut = value;
    return true;
}

// Appends one unsigned decimal value to a bounded command string.
static bool append_u32(char* buffer, size_t capacity, size_t* lengthInOut, uint32_t value) {
    if (!buffer || !lengthInOut || *lengthInOut >= capacity) return false;
    char reverse[16];
    size_t count = 0;
    do {
        reverse[count++] = (char)('0' + value % 10u);
        value /= 10u;
    } while (value && count < sizeof(reverse));
    if (*lengthInOut + count + 1 > capacity) return false;
    while (count) buffer[(*lengthInOut)++] = reverse[--count];
    buffer[*lengthInOut] = '\0';
    return true;
}

// Runs inside the stripped child and returns a distinct failure code for each leaked ambient authority.
static int run_restricted_child(uint32_t parentPid) {
    SYS_SECURITY_CONTEXT_INFO info;
    memset(&info, 0, sizeof(info));
    info.Size = sizeof(info);
    KSTATUS status = SysQuerySecurityContext(SYS_SECURITY_CONTEXT_CURRENT, &info);
    if (!KSUCCESS(status)) return 10;
    if (info.EffectivePrivileges || info.DelegablePrivileges)
        return 11;

    SYS_HANDLE process = SYS_HANDLE_INVALID;
    status = SysOpenProcess(parentPid, SYS_RIGHT_QUERY, &process);
    if (status != NTLSTATUS_ACCESS_DENIED) return 12;

    size_t processCount = 0;
    status = SysGetProcessList((SYS_PROCESS_INFO*)0, 0, &processCount);
    if (status != NTLSTATUS_ACCESS_DENIED) return 13;

    size_t required = 0;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_SCHEDULER, (void*)0, 0, &required);
    if (status != NTLSTATUS_ACCESS_DENIED) return 14;

    SYS_DISPLAY_INFO displayInfo;
    memset(&displayInfo, 0, sizeof(displayInfo));
    status = SysQueryDisplayInfo(&displayInfo);
    if (status != NTLSTATUS_ACCESS_DENIED) return 16;

    SYS_HANDLE socketHandle = SYS_HANDLE_INVALID;
    status = SysCreateSocket(0, 0, 0, &socketHandle);
    if (status != NTLSTATUS_ACCESS_DENIED) return 17;

    SYS_PRINT_DEBUG_REQUEST debugRequest = { sizeof(debugRequest), SysSvFromParts((const char*)0, 0) };
    status = SysPrintDebug(&debugRequest);
    if (status != NTLSTATUS_ACCESS_DENIED) return 18;

    SYS_WRITE_CONSOLE_REQUEST consoleRequest = { sizeof(consoleRequest), SysSvFromParts((const char*)0, 0) };
    status = SysWriteConsole(&consoleRequest);
    if (status != NTLSTATUS_ACCESS_DENIED) return 19;

    SYS_PEB* peb = SysGetPeb();
    if (!peb || !peb->ImagePath.Data || peb->ImagePath.Length == 0) return 20;
    SYS_HANDLE file = SYS_HANDLE_INVALID;
    status = NtlOpenFile(peb->ImagePath,
                                                            SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY,
                                                            SYS_FILE_OPEN_EXISTING,
                                                            SYS_FILE_FLAG_NONE,
                             &file);
    if (!KSUCCESS(status)) return 21;
    (void)SysCloseHandle(file);

    uint32_t endpointCount = 0;
    SYS_AUDIO_ENUMERATE_ENDPOINTS_REQUEST enumerateAudio;
    memset(&enumerateAudio, 0, sizeof(enumerateAudio));
    enumerateAudio.Size = sizeof(enumerateAudio);
    enumerateAudio.CountOut = &endpointCount;
    status = SysAudioEnumerateEndpoints(&enumerateAudio);
    if (status != NTLSTATUS_ACCESS_DENIED) return 22;

    SYS_HANDLE listener = SYS_HANDLE_INVALID;
    SYS_CREATE_PORT_REQUEST createPort;
    memset(&createPort, 0, sizeof(createPort));
    createPort.Size = sizeof(createPort);
    createPort.HandleOut = &listener;
    status = SysCreatePort(&createPort);
    if (!KSUCCESS(status)) return 25;
    status = SysRegisterSystemService(SYS_HANDLE_INVALID, listener);
    (void)SysCloseHandle(listener);
    if (status != NTLSTATUS_INVALID_HANDLE) return 26;

    uint32_t selfPid = SysGetCurrentPid();
    status = SysOpenProcess(selfPid, SYS_RIGHT_QUERY, &process);
    if (!KSUCCESS(status)) return 27;
    status = SysCloseHandle(process);
    if (!KSUCCESS(status)) return 28;

    process = SYS_HANDLE_INVALID;
    status = SysOpenProcess(selfPid, SYS_RIGHT_QUERY | SYS_PROCESS_RIGHT_SECURITY_CONTEXT, &process);
    if (status != NTLSTATUS_ACCESS_DENIED) return 29;

    return 0;
}

// Creates one zero-authority security context from the caller's current delegable root.
static SYS_HANDLE derive_stripped_context(void) {
    SYS_HANDLE handle = SYS_HANDLE_INVALID;
    SYS_DERIVE_SECURITY_CONTEXT_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.SourceContext = SYS_SECURITY_CONTEXT_CURRENT;
    request.Flags = SYS_SECURITY_CONTEXT_FLAG_NONE;
    request.ContextHandleOut = &handle;
    KSTATUS status = SysDeriveSecurityContext(&request);
    if (!KSUCCESS(status)) fail_status("derive stripped context", status);
    return handle;
}

// Runs the privileged parent half of the security-context smoke test.
static int run_parent(void) {
    SYS_PEB* peb = SysGetPeb();
    if (!peb || !peb->ImagePath.Data || peb->ImagePath.Length == 0)
        fail_status("missing image path", NTLSTATUS_FAILURE);

    SYS_SECURITY_CONTEXT_INFO rootInfo;
    memset(&rootInfo, 0, sizeof(rootInfo));
    rootInfo.Size = sizeof(rootInfo);
    KSTATUS status = SysQuerySecurityContext(SYS_SECURITY_CONTEXT_CURRENT, &rootInfo);
    if (!KSUCCESS(status)) fail_status("query current context", status);
    if (rootInfo.EffectivePrivileges != SYS_PRIV_ALL ||
        rootInfo.DelegablePrivileges != SYS_PRIV_ALL)
        fail_status("bootstrap context authority", NTLSTATUS_FAILURE);

    SYS_HANDLE selfProcess = SYS_HANDLE_INVALID;
    status = SysOpenProcess(SysGetCurrentPid(),
                            SYS_RIGHT_QUERY | SYS_PROCESS_RIGHT_SECURITY_CONTEXT,
                            &selfProcess);
    if (!KSUCCESS(status)) fail_status("open self security-context process", status);
    SYS_HANDLE exactContext = SYS_HANDLE_INVALID;
    status = SysOpenProcessSecurityContext(
        selfProcess, SYS_RIGHT_QUERY | SYS_SECURITY_CONTEXT_RIGHT_ASSIGN, &exactContext);
    if (!KSUCCESS(status)) fail_status("open exact process security context", status);
    SYS_SECURITY_CONTEXT_INFO exactInfo;
    memset(&exactInfo, 0, sizeof(exactInfo));
    exactInfo.Size = sizeof(exactInfo);
    status = SysQuerySecurityContext(exactContext, &exactInfo);
    if (!KSUCCESS(status)) fail_status("query exact process security context", status);
    if (exactInfo.SecurityId != rootInfo.SecurityId ||
        exactInfo.EffectivePrivileges != rootInfo.EffectivePrivileges ||
        exactInfo.DelegablePrivileges != rootInfo.DelegablePrivileges)
        fail_status("exact process security context identity", NTLSTATUS_FAILURE);
    (void)SysCloseHandle(exactContext);
    (void)SysCloseHandle(selfProcess);

    SYS_HANDLE queryOnlyProcess = SYS_HANDLE_INVALID;
    status = SysOpenProcess(SysGetCurrentPid(), SYS_RIGHT_QUERY, &queryOnlyProcess);
    if (!KSUCCESS(status)) fail_status("open query-only self process", status);
    exactContext = SYS_HANDLE_INVALID;
    status = SysOpenProcessSecurityContext(queryOnlyProcess, SYS_RIGHT_QUERY, &exactContext);
    (void)SysCloseHandle(queryOnlyProcess);
    if (status != NTLSTATUS_ACCESS_DENIED)
        fail_status("process security-context right", status);

    SYS_HANDLE restricted = derive_stripped_context();
    SYS_SECURITY_CONTEXT_INFO restrictedInfo;
    memset(&restrictedInfo, 0, sizeof(restrictedInfo));
    restrictedInfo.Size = sizeof(restrictedInfo);
    status = SysQuerySecurityContext(restricted, &restrictedInfo);
    if (!KSUCCESS(status)) fail_status("query stripped context", status);
    if (restrictedInfo.SecurityId <= rootInfo.SecurityId ||
        restrictedInfo.EffectivePrivileges || restrictedInfo.DelegablePrivileges)
        fail_status("stripped context contents", NTLSTATUS_FAILURE);

    SYS_HANDLE illegal = SYS_HANDLE_INVALID;
    SYS_DERIVE_SECURITY_CONTEXT_REQUEST amplify;
    memset(&amplify, 0, sizeof(amplify));
    amplify.Size = sizeof(amplify);
    amplify.SourceContext = restricted;
    amplify.EffectivePrivileges = SYS_PRIV_PROCESS_OPEN_ANY;
    amplify.ContextHandleOut = &illegal;
    status = SysDeriveSecurityContext(&amplify);
    if (status != NTLSTATUS_ACCESS_DENIED)
        fail_status("blocked privilege amplification", status);

    SYS_HANDLE queryOnly = SYS_HANDLE_INVALID;
    SYS_DERIVE_HANDLE_REQUEST deriveHandle = {
        sizeof(deriveHandle), restricted, SYS_RIGHT_QUERY, &queryOnly
    };
    status = SysDeriveHandle(&deriveHandle);
    if (!KSUCCESS(status)) fail_status("derive query-only context handle", status);

    SYS_HANDLE rejectedChild = SYS_HANDLE_INVALID;
    NTL_CREATE_SYSTEM_PROCESS_REQUEST rejectedCreate;
    memset(&rejectedCreate, 0, sizeof(rejectedCreate));
    rejectedCreate.Size = sizeof(rejectedCreate);
    rejectedCreate.Flags = NTL_PROCESS_LOAD_FLAG_NONE;
    rejectedCreate.Path = peb->ImagePath;
    rejectedCreate.CommandLine = SysSvFromCstr("restricted:1");
    rejectedCreate.SecurityContext = queryOnly;
    rejectedCreate.ResolveModuleName = NtlResolveSystemModuleName;
    rejectedCreate.ProcessHandleOut = &rejectedChild;
    status = NtlCreateSystemProcess(&rejectedCreate);
    if (status != NTLSTATUS_ACCESS_DENIED)
        fail_status("context assignment right", status);
    (void)SysCloseHandle(queryOnly);

    char command[48] = "restricted:";
    size_t commandLength = strlen(command);
    if (!append_u32(command, sizeof(command), &commandLength, SysGetCurrentPid()))
        fail_status("build child command", NTLSTATUS_BUFFER_TOO_SMALL);

    SYS_HANDLE child = SYS_HANDLE_INVALID;
    NTL_CREATE_SYSTEM_PROCESS_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Flags = NTL_PROCESS_LOAD_FLAG_NONE;
    create.Path = peb->ImagePath;
    create.CommandLine = SysSvFromCstr(command);
    create.SecurityContext = restricted;
    create.ResolveModuleName = NtlResolveSystemModuleName;
    create.ProcessHandleOut = &child;
    status = NtlCreateSystemProcess(&create);
    if (!KSUCCESS(status)) fail_status("create restricted child", status);

    status = SysWait(child, SYS_WAIT_FLAG_NONE);
    if (!KSUCCESS(status)) fail_status("wait restricted child", status);
    int32_t childExit = -1;
    status = SysGetProcessExitCode(child, &childExit);
    if (!KSUCCESS(status)) fail_status("query restricted child", status);
    (void)SysCloseHandle(child);
    if (childExit != 0) {
        print("security_smoke: restricted child failure=");
        print_u32((uint32_t)childExit);
        print("\n");
        return 1;
    }

    status = SysCloseHandle(restricted);
    if (!KSUCCESS(status)) fail_status("close stripped context", status);

    SYS_HANDLE later = derive_stripped_context();
    SYS_SECURITY_CONTEXT_INFO laterInfo;
    memset(&laterInfo, 0, sizeof(laterInfo));
    laterInfo.Size = sizeof(laterInfo);
    status = SysQuerySecurityContext(later, &laterInfo);
    if (!KSUCCESS(status)) fail_status("query later context", status);
    if (laterInfo.SecurityId <= restrictedInfo.SecurityId)
        fail_status("security id reuse", NTLSTATUS_FAILURE);
    (void)SysCloseHandle(later);

    print("security_smoke: PASS\n");
    return 0;
}

// Runs either the privileged parent or the stripped child based on the process command line.
int main(void) {
    SYS_PEB* peb = SysGetPeb();
    if (!peb) return 2;
    const char prefix[] = "restricted:";
    if (peb->CommandLine.Length > sizeof(prefix) - 1u &&
        memcmp(peb->CommandLine.Data, prefix, sizeof(prefix) - 1u) == 0) {
        uint32_t parentPid = 0;
        if (!parse_u32(peb->CommandLine, sizeof(prefix) - 1u, &parentPid)) return 3;
        return run_restricted_child(parentPid);
    }
    return run_parent();
}