Enabling and checking token privileges in VB.NET

' ***********************************************************************
' This is what I used for testing this Stack Overflow answer I wrote:
' http://stackoverflow.com/a/36862571/1871033
' Especially the way the NativeMethods are defined and imported again
' are not how it should be, but it didn't matter for the answer, I just
' tried to get things to work quickly. So I share this full code for
' reference, so others don't have to puzzle it together like I had to,
' but it's not supposed to be a perfect example.
' Note that I only wrote parts of this code, the other part was written
' by SO user ElektroStudios.
' ***********************************************************************

Imports System.Runtime.InteropServices
Imports WindowsApplication1.Form1.NativeMethods
Imports System.ComponentModel

Public Class Form1
    Public Class NativeMethods

        ' http://msdn.microsoft.com/en-us/library/windows/desktop/aa379295%28v=vs.85%29.aspx
        <DllImport("advapi32.dll", SetLastError:=True)>
        Public Shared Function OpenProcessToken(ByVal processHandle As IntPtr,
                                            ByVal desiredAccess As TokenAccess,
                                            ByRef tokenHandle As IntPtr
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' http://msdn.microsoft.com/en-us/library/windows/desktop/aa379180%28v=vs.85%29.aspx
        <DllImport("Advapi32.dll", SetLastError:=True, CharSet:=CharSet.Auto, BestFitMapping:=False, ThrowOnUnmappableChar:=True)>
        Public Shared Function LookupPrivilegeValue(ByVal lpSystemName As String,
                                                ByVal lpName As String,
                                                ByRef lpLuid As Luid
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' http://msdn.microsoft.com/es-es/library/windows/desktop/aa375202%28v=vs.85%29.aspx
        <DllImport("Advapi32.dll", SetLastError:=True)>
        Public Shared Function AdjustTokenPrivileges(ByVal tokenHandle As IntPtr,
                                                 ByVal disableAllPrivileges As Boolean,
                                                 ByRef newState As TokenPrivileges,
                                                 ByVal bufferLength As Integer,
                                                 ByRef refPreviousState As TokenPrivileges,
                                                 ByRef refReturnLength As IntPtr
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa379304%28v=vs.85%29.aspx
        <DllImport("Advapi32.dll", SetLastError:=True)>
        Public Shared Function PrivilegeCheck(ByVal token As IntPtr,
                              <[In], Out> ByRef privileges As PrivilegeSet,
                                          ByRef refResult As Boolean
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa446616%28v=vs.85%29.aspx
        <DllImport("advapi32.dll", SetLastError:=True)>
        Public Shared Function DuplicateToken(ByVal tokenHandle As IntPtr,
                                          ByVal impersonationLevel As SecurityImpersonationLevel,
                                          ByRef duplicateTokenHandle As IntPtr
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa446671%28v=vs.85%29.aspx
        <DllImport("Advapi32.dll", SetLastError:=True)>
        Public Shared Function GetTokenInformation(ByVal tokenHandle As IntPtr,
                                               ByVal tokenInformationClass As TokenInformationClass,
                                               ByVal tokenInformation As IntPtr,
                                               ByVal tokenInformationLength As Integer,
                                               ByRef refReturnLength As Integer
        ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function

        ' http://msdn.microsoft.com/en-us/library/windows/desktop/aa374905%28v=vs.85%29.aspx
        <Flags>
        Public Enum TokenAccess As UInteger
            ' THIS ENUMERATION IS PARTIALLY DEFINED.
            ' **************************************
            AdjustPrivileges = &H20UI
            Query = &H8UI
            Duplicate = &H2UI
        End Enum

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa379630%28v=vs.85%29.aspx
        <Flags>
        Public Enum TokenPrivilegeAttributes As UInteger
            PrivilegeDisabled = &H0UI
            PrivilegeEnabledByDefault = &H1UI
            PrivilegeEnabled = &H2UI
            PrivilegeRemoved = &H4UI
            PrivilegeUsedForAccess = &H80000000UI
        End Enum

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa379572(v=vs.85).aspx
        Public Enum SecurityImpersonationLevel As Integer
            SecurityAnonymous = 0
            SecurityIdentification = 1
            SecurityImpersonation = 2
            SecurityDelegation = 3
        End Enum

        ' http://msdn.microsoft.com/en-us/library/windows/desktop/aa379261%28v=vs.85%29.aspx
        <StructLayout(LayoutKind.Sequential)>
        Public Structure Luid
            Public LowPart As UInteger
            Public HighPart As Integer
        End Structure

        ' http://msdn.microsoft.com/en-us/library/windows/desktop/aa379263%28v=vs.85%29.aspx
        <StructLayout(LayoutKind.Sequential)>
        Public Structure LuidAndAttributes
            Public Luid As Luid
            Public Attributes As TokenPrivilegeAttributes
        End Structure

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa379630%28v=vs.85%29.aspx
        <StructLayout(LayoutKind.Sequential)>
        Public Structure TokenPrivileges
            Public PrivilegeCount As UInteger

            <MarshalAs(UnmanagedType.ByValArray, SizeConst:=1)>
            Public Privileges As LuidAndAttributes()
        End Structure

        ' https://msdn.microsoft.com/en-us/library/windows/desktop/aa379307%28v=vs.85%29.aspx
        <StructLayout(LayoutKind.Sequential)>
        Public Structure PrivilegeSet
            Public PrivilegeCount As UInteger
            Public Control As UInteger

            <MarshalAs(UnmanagedType.ByValArray, SizeConst:=1)>
            Public Privileges As LuidAndAttributes()
        End Structure

        Public Enum TokenInformationClass As UInteger
            TokenUser = 1
            TokenGroups
            TokenPrivileges
            TokenOwner
            TokenPrimaryGroup
            TokenDefaultDacl
            TokenSource
            TokenType
            TokenImpersonationLevel
            TokenStatistics
            TokenRestrictedSids
            TokenSessionId
            TokenGroupsAndPrivileges
            TokenSessionReference
            TokenSandBoxInert
            TokenAuditPolicy
            TokenOrigin

        End Enum

        <DllImport("kernel32.dll", SetLastError:=True)> _
        Public Shared Function CloseHandle(handle As IntPtr) As Boolean
        End Function

        <DllImport("Advapi32.dll", SetLastError:=True, CharSet:=CharSet.Auto, BestFitMapping:=False, ThrowOnUnmappableChar:=True)>
        Public Shared Function LookupPrivilegeName( _
                                                  lpSystemName As String, _
                                                  ByRef lpLuid As Luid, _
                                                  lpName As System.Text.StringBuilder, _
                                                  ByRef cchName As Integer _
                                                  ) As <MarshalAs(UnmanagedType.Bool)> Boolean
        End Function
    End Class

    Private Sub Button1_Click(sender As Object, e As EventArgs) Handles Button1.Click
        Dim pHandle As IntPtr = Process.GetCurrentProcess().Handle
        Dim privilegeName As String = "SeShutdownPrivilege"
        Dim tokenAccess As TokenAccess = (tokenAccess.AdjustPrivileges Or tokenAccess.Query Or tokenAccess.Duplicate)
        Dim hToken As IntPtr
        Dim hTokenDup As IntPtr

        Try
            ' ****************************************************************************
            ' 1st Step: Enable the "SeShutdownPrivilege" privilege in the current process.
            ' ****************************************************************************

            Dim win32Err As Integer

            ' Get the process token.
            NativeMethods.OpenProcessToken(pHandle, tokenAccess, hToken)

            ' Set up a LuidAndAttributes structure containing the privilege to enable,
            ' getting the LUID that corresponds to the privilege.
            Dim luAttr As New LuidAndAttributes
            luAttr.Attributes = TokenPrivilegeAttributes.PrivilegeEnabled
            NativeMethods.LookupPrivilegeValue(Nothing, privilegeName, luAttr.Luid)

            ' Set up a TokenPrivileges structure containing only the source privilege.
            Dim newState As New TokenPrivileges
            newState.PrivilegeCount = 1
            newState.Privileges = New LuidAndAttributes() {luAttr}

            ' Set up a TokenPrivileges structure for the previous (modified) privileges.
            Dim prevState As New TokenPrivileges
            prevState = New TokenPrivileges
            ReDim prevState.Privileges(CInt(newState.PrivilegeCount))

            ' Apply the TokenPrivileges structure to the source process token.
            Dim bufferLength As Integer = Marshal.SizeOf(prevState)
            Dim returnLength As IntPtr
            If Not NativeMethods.AdjustTokenPrivileges(hToken, False, newState, bufferLength, prevState, returnLength) Then
                win32Err = Marshal.GetLastWin32Error
                MessageBox.Show("AdjustTokenPrivileges failed.")
                Throw New Win32Exception(win32Err)
            End If

            ' *********************************************************************
            ' Everything OK at this point,
            ' as AdjustTokenPrivileges dididn't failed, I assume the privilege Is enabled in the process.
            '
            ' 2n Step: Check whether the privilege is enabled or not...
            ' *********************************************************************

            ' Set up a new one LuidAndAttributes structure containing the privilege to check,
            ' getting the LUID that corresponds to the privilege.
            luAttr = New LuidAndAttributes
            NativeMethods.LookupPrivilegeValue(Nothing, privilegeName, luAttr.Luid)

            ' *********************************************************************
            ' Trying PrivilegeCheck and Duplicatetoken methodology...
            ' *********************************************************************

            NativeMethods.DuplicateToken(hToken, SecurityImpersonationLevel.SecurityImpersonation, hTokenDup)
            win32Err = Marshal.GetLastWin32Error

            If (hTokenDup <> IntPtr.Zero) Then
                Dim result As Boolean
                Dim pSet As New PrivilegeSet
                pSet.Control = 0
                pSet.PrivilegeCount = 1
                pSet.Privileges = New LuidAndAttributes() {luAttr}

                If Not NativeMethods.PrivilegeCheck(hToken, pSet, result) Then
                    win32Err = Marshal.GetLastWin32Error
                    MessageBox.Show("PrivilegeCheck using original access-token failed.")
                    ' Ignore exception, to continue with the GetTokenInformation methodology.
                    ' Throw New Win32Exception(win32Err)

                Else
                    MessageBox.Show(String.Format("{0} (original token) state is: {1}",
                                                      privilegeName, pSet.Privileges(0).Attributes.ToString()))

                    If Not NativeMethods.PrivilegeCheck(hTokenDup, pSet, result) Then
                        win32Err = Marshal.GetLastWin32Error
                        MessageBox.Show("PrivilegeCheck using impersonated access-token failed.")
                        ' Ignore exception, to continue with the GetTokenInformation methodology.
                        ' Throw New Win32Exception(win32Err)
                    Else
                        MessageBox.Show(String.Format("{0} (impersonated token) state is: {1}",
                                                      privilegeName, pSet.Privileges(0).Attributes.ToString()))

                    End If

                End If

            Else
                MessageBox.Show("DuplicateToken failed.")
                ' Ignore exception, to continue with the GetTokenInformation methodology.
                ' Throw New Win32Exception(win32Err)

            End If

            ' *********************************************************************
            ' Trying GetTokenInformation methodology...
            ' *********************************************************************

            Dim tkp As New TokenPrivileges
            Dim tkpHandle As IntPtr = IntPtr.Zero ' <<< will be set later
            Dim tkInfoLength As Integer = 0

            ' Here we call GetTokenInformation the first time to receive the length of the data it would like to store.
            NativeMethods.GetTokenInformation(hToken, TokenInformationClass.TokenPrivileges, IntPtr.Zero, tkInfoLength, tkInfoLength)
            win32Err = Marshal.GetLastWin32Error

            ' Since the "current" length we pass is 0, we'll always get "error" 122, which is fine. We also get the required length returned.
            If (win32Err <> 122) Then
                MessageBox.Show("GetTokenInformation failed in the attempt to get the TokenPrivileges's size.")
                Throw New Win32Exception(win32Err)
            Else
                Try
                    ' Here we allocate memory for receiving the actual data. By now, tkInfoLength contains the size of the memory block we need to allocate.
                    tkpHandle = Marshal.AllocHGlobal(tkInfoLength)

                    ' This time, we shouldn't get an error 122, because this time we already have set the correct buffer size. GetTokenInformation should now write the data into the memory block we just allocated.
                    If Not NativeMethods.GetTokenInformation(hToken, TokenInformationClass.TokenPrivileges, tkpHandle, tkInfoLength, tkInfoLength) Then
                        win32Err = Marshal.GetLastWin32Error
                        MessageBox.Show("GetTokenInformation failed in the attempt to get the TokenPrivileges.")
                        Throw New Win32Exception(win32Err)

                    Else
                        ' We will now ask PtrToStructure to read the raw data out of the memory block and convert it to a managed structure of type TokenPrivileges which we can use in our code. That's it!
                        tkp = Marshal.PtrToStructure(tkpHandle, GetType(TokenPrivileges))

                        ' We have to iterate over all privileges listed in the TokenPrivileges structure to find the one we are looking for
                        Dim found As Boolean = False
                        For i As Integer = 0 To tkp.PrivilegeCount - 1
                            ' There is a problem: Marshal.PtrToStructure can't marshal variable-length structures, but the array TokenPrivileges::Privileges has
                            ' a variable length determined by the value of TokenPrivileges::PrivilegeCount! Since we don't know the size at compile time, the
                            ' size of the array was hardcoded to 1, which means that we would only be able to access the first element of the array.
                            ' To work around this, we calculate the raw memory offset pointing to the array element we need and load it separately into a
                            ' LuidAndAttributes variable.
                            ' The way this works is: The contents of the TokenPrivilege structure or stored in memory one after another, like this:
                            '   PrivilegeCount (type: UInteger)
                            '   Privileges(0) (type: LuidAndAttributes)
                            '   Privileges(1) (type: LuidAndAttributes) << these and all further we normally can't access
                            '   Privileges(2) (type: LuidAndAttributes)
                            ' ...and so on.
                            ' We are now calculating the offset into the structure for a specific array element. Let's use Privileges(2) as example:
                            ' To get to it, we need to take the pointer to the beginning of the structure and add the sizes of all previous elements,
                            ' which would be once the size of PrivilegeCount and then 2 times the size of a LuidAndAttributes structure.
                            Dim directPointer As New IntPtr(tkpHandle.ToInt64() + Len(tkp.PrivilegeCount) + i * Marshal.SizeOf(GetType(LuidAndAttributes)))
                            Dim luidAndAttributes As LuidAndAttributes = Marshal.PtrToStructure(directPointer, GetType(LuidAndAttributes))

                            ' Get the privilege name. We first call LookupPrivilegeName with a zero size to get the real size we need, then reserve space, then get the actual data
                            ' NOTE: The part below isn't actually necessary as commented Mark Hurd pointed out, because you already have the right privilege's LUID in luAttr.Luid.
                            ' But I'll leave it here anyway in case somebody uses this piece of code without the part which sets the privilege.
                            ' Another solution in this case would also be to run LookupPrivilegeValue to get the LUID to compare to (which is what led to luAttr.Luid as well).
                            Dim privNameLen As Integer = 0
                            Dim sb As New System.Text.StringBuilder()
                            NativeMethods.LookupPrivilegeName(Nothing, luidAndAttributes.Luid, sb, privNameLen)
                            sb.EnsureCapacity(privNameLen + 1)
                            If Not NativeMethods.LookupPrivilegeName(Nothing, luidAndAttributes.Luid, sb, privNameLen) Then
                                win32Err = Marshal.GetLastWin32Error
                                MessageBox.Show("LookupPrivilegeName failed.")
                                Throw New Win32Exception(win32Err)
                            End If

                            ' Now that we have the name, we can check if it's the one we are looking for.
                            ' NOTE: Refering to me comments above, this could be just: If luidAndAttributes.Luid = luAttr.Luid
                            If sb.ToString() = privilegeName Then
                                ' Found! So we can finally get the status of the privilege!
                                found = True
                                MessageBox.Show(String.Format("{0} state is: {1}", privilegeName, luidAndAttributes.Attributes.ToString()))
                                Exit For
                            End If
                        Next

                        If Not found Then MessageBox.Show(String.Format("{0} not found in list of privileges!", privilegeName))
                    End If
                Finally
                    ' Make sure the memory block is freed again (even when an error occured)
                    If tkpHandle Then Marshal.FreeHGlobal(tkpHandle)
                End Try

            End If
        Catch ex As Win32Exception
            MessageBox.Show(ex.NativeErrorCode & " " & ex.Message)

        Catch ex As Exception
            MessageBox.Show(ex.Message)

        Finally
            If (hTokenDup <> IntPtr.Zero) Then
                NativeMethods.CloseHandle(hTokenDup)
            End If

            If (hToken <> IntPtr.Zero) Then
                NativeMethods.CloseHandle(hToken)
            End If

        End Try
    End Sub
End Class