通过读源码win10驱动下实现3环的GetEnvironmentVariable

Posted 2021-01-30

效果图：

NTSTATUS NTAPI
RtlQueryEnvironmentVariable_U(PWSTR Environment,
    PCUNICODE_STRING Name,
    PUNICODE_STRING Value)
{
    NTSTATUS Status;
    PWSTR wcs;
    UNICODE_STRING var;
    PWSTR val;
    BOOLEAN SysEnvUsed = FALSE;

    DbgPrint("RtlQueryEnvironmentVariable_U Environment %p Variable %wZ Value %p
",
        Environment, Name, Value);

    if (Environment == NULL)
    {
        MPPEB Peb = RtlGetCurrentPeb();
        if (Peb) {
            //RtlAcquirePebLock();

            Environment = Peb->ProcessParameters->Environment;
            SysEnvUsed = TRUE;
        }
    }

    if (Environment == NULL)
    {
        //if (SysEnvUsed)
            //RtlReleasePebLock();
        return(STATUS_VARIABLE_NOT_FOUND);
    }

    Value->Length = 0;

    wcs = Environment;
    DbgPrint("Starting search at :%p
", wcs);
    while (*wcs)
    {
        var.Buffer = wcs++;
        wcs = wcschr(wcs, L‘=‘);
        if (wcs == NULL)
        {
            wcs = var.Buffer + wcslen(var.Buffer);
            DbgPrint("Search at :%S
", wcs);
        }
        if (*wcs)
        {
            var.Length = var.MaximumLength = (USHORT)(wcs - var.Buffer) * sizeof(WCHAR);
            val = ++wcs;
            wcs += wcslen(wcs);
            DbgPrint("Search at :%S
", wcs);

            if (RtlEqualUnicodeString(&var, Name, TRUE))
            {
                Value->Length = (USHORT)(wcs - val) * sizeof(WCHAR);
                if (Value->Length <= Value->MaximumLength)
                {
                    memcpy(Value->Buffer, val,
                        min(Value->Length + sizeof(WCHAR), Value->MaximumLength));
                    DbgPrint("Value %S
", val);
                    DbgPrint("Return STATUS_SUCCESS
");
                    Status = STATUS_SUCCESS;
                }
                else
                {
                    DbgPrint("Return STATUS_BUFFER_TOO_SMALL
");
                    Status = STATUS_BUFFER_TOO_SMALL;
                }

                //if (SysEnvUsed)
                    //RtlReleasePebLock();

                return(Status);
            }
        }
        wcs++;
    }

    /*if (SysEnvUsed)
        RtlReleasePebLock();
*/
    DbgPrint("Return STATUS_VARIABLE_NOT_FOUND: %wZ
", Name);
    return(STATUS_VARIABLE_NOT_FOUND);
}
//获取环境变量 
DWORD My_Get_Environment_Variable(IN LPCWSTR lpName,
    IN LPWSTR lpBuffer,
    IN DWORD nSize) {
    UNICODE_STRING VarName, VarValue;
    NTSTATUS Status;
    USHORT UniSize;

    if (nSize <= (UNICODE_STRING_MAX_CHARS - 1))
    {
        if (nSize)
        {
            UniSize = (USHORT)nSize * sizeof(WCHAR) - sizeof(UNICODE_NULL);
        }
        else
        {
            UniSize = 0;
        }
    }
    else
    {
        UniSize = UNICODE_STRING_MAX_BYTES - sizeof(UNICODE_NULL);
    }

    Status = RtlInitUnicodeStringEx(&VarName, lpName);
    if (!NT_SUCCESS(Status))
    {
        BaseSetLastNTError(Status);
        return 0;
    }

    RtlInitEmptyUnicodeString(&VarValue, lpBuffer, UniSize);

    Status = RtlQueryEnvironmentVariable_U(NULL, &VarName, &VarValue);
    if (!NT_SUCCESS(Status))
    {
    DbgPrint("RtlQueryEnvironmentVariable_U----------- %S
", VarValue);
        if (Status == STATUS_BUFFER_TOO_SMALL)
        {
            return (VarValue.Length / sizeof(WCHAR)) + sizeof(ANSI_NULL);
        }
        BaseSetLastNTError(Status);
        return 0;
    }

    lpBuffer[VarValue.Length / sizeof(WCHAR)] = UNICODE_NULL;

    return (VarValue.Length / sizeof(WCHAR));
}

用法：

wchar_t buffer[256];
    wchar_t pwcDevNameBuf = NULL;
    DWORD code = My_Get_Environment_Variable(L"TEMP", buffer, 256);
    DbgPrint("buffer----------- %S
", buffer);