Windowsの特権を操作

久しぶりに特権関連を書いてみた。

stdafx.h

#pragma once

#include "targetver.h"

#define _CRTDBG_MAP_ALLOC
#include <stdlib.h>
#include <crtdbg.h>

#include <tchar.h>

#include <string>
#include <iostream>
#include <boost/format.hpp>
#include <boost/scoped_array.hpp>

#include <Windows.h>
#include <NTSecAPI.h>

sample.cpp

#include "stdafx.h"

#if defined(UNICODE) || defined(_UNICODE)
#define tstring std::wstring
#define tcout std::wcout
#define tformat boost::wformat
#else
#define tstring std::string
#define tcout std::cout
#define tformat boost::format
#endif

class Privilege
{
private:
    LUID_AND_ATTRIBUTES luidAndAttrs;
    tstring name;
    tstring displayName;

public:
    static const int PRIV_NAME_MAX_LENGTH = 64;
    static const int PRIV_DISPNAME_MAX_LENGTH = 128;
    typedef std::vector<boost::shared_ptr<Privilege>> PrivilegeContainer;

    Privilege(LUID_AND_ATTRIBUTES luidAndAttrs)
    {
        this->luidAndAttrs = luidAndAttrs;
    }

    static PrivilegeContainer GetPrivileges()
    {
        PrivilegeContainer result;
        HANDLE hToken;

        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
        {
            tcout << "OpenProcessToken 失敗" << std::endl;
            return result;
        }

        DWORD length;
        GetTokenInformation(hToken, TokenPrivileges, NULL, 0, &length);
        boost::scoped_array<TOKEN_PRIVILEGES> privs(new TOKEN_PRIVILEGES[length]);
        GetTokenInformation(hToken, TokenPrivileges, privs.get(), length, &length);


        for (DWORD i=0; i<privs.get()->PrivilegeCount; i++)
        {
            auto luidAndAttrs = privs.get()->Privileges[i];
            boost::shared_ptr<Privilege> priv(new Privilege(luidAndAttrs));
            result.push_back(priv);
        }

        CloseHandle(hToken);

        return result;
    }

    tstring GetName()
    {
        if (this->name.empty())
        {
            DWORD cchName = PRIV_NAME_MAX_LENGTH;
            TCHAR name[PRIV_NAME_MAX_LENGTH+1] = {0};

            LookupPrivilegeName(NULL, &this->luidAndAttrs.Luid, name, &cchName);
            this->name = tstring(name);
        }
        return this->name;
    }

    tstring GetDisplayName()
    {
        if (this->displayName.empty())
        {
            DWORD cchDispName = PRIV_DISPNAME_MAX_LENGTH;
            TCHAR dispname[PRIV_DISPNAME_MAX_LENGTH+1] = {0};
            DWORD langID = 0;

            LookupPrivilegeDisplayName(NULL, GetName().c_str(), dispname, &cchDispName, &langID);
            this->displayName = tstring(dispname);
        }
        return this->displayName;
    }

    DWORD GetStatus() const
    {
        return this->luidAndAttrs.Attributes;
    }

    tstring GetStatusText() const
    {
        tstring text;
        if (this->luidAndAttrs.Attributes & SE_PRIVILEGE_ENABLED)
        {
            text.append(_T("有効"));
        }
        else
        {
            text.append(_T("無効"));
        }
        return text;
    }

    void SetStatus(DWORD status)
    {
        HANDLE hToken;

        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
        {
            tcout << "OpenProcessToken 失敗" << std::endl;
            return;
        }
        TOKEN_PRIVILEGES privs;
        privs.PrivilegeCount = 1;
        privs.Privileges[0].Attributes = status;
        privs.Privileges[0].Luid = this->luidAndAttrs.Luid;
        if (!AdjustTokenPrivileges(hToken, FALSE, &privs, 0, NULL, NULL))
        {
            tcout << "AdjustTokenPrivileges 失敗" << std::endl;
            CloseHandle(hToken);
            return;
        }

        this->luidAndAttrs.Attributes = status;
        CloseHandle(hToken);
    }
};

int _tmain(int argc, _TCHAR* argv[])
{
    _CrtSetDbgFlag(_CRTDBG_ALLOC_MEM_DF | _CRTDBG_LEAK_CHECK_DF);
    _tsetlocale(LC_ALL, _T(""));

    tcout << _T("----- before -----") << std::endl;

    auto privs = Privilege::GetPrivileges();
    for each(auto priv in privs)
    {
        auto ptr = priv.get();
        tcout
            << tformat(L"%-4s %-64s %-128s") 
                % ptr->GetStatusText()
                % ptr->GetName()
                % ptr->GetDisplayName()
            << std::endl;
    }

    tcout << _T("----- 特権操作 -----") << std::endl;

    auto pred = 
        [](const boost::shared_ptr<Privilege> &p) {
            return (p->GetName() == SE_SHUTDOWN_NAME);
        };
    auto targetPriv = std::find_if(privs.begin(), privs.end(), pred);
    tcout << targetPriv->get()->GetName() << std::endl;
    targetPriv->get()->SetStatus(SE_PRIVILEGE_ENABLED);

    tcout << _T("----- after -----") << std::endl;

    privs = Privilege::GetPrivileges();
    for each(auto priv in privs)
    {
        auto ptr = priv.get();
        tcout
            << tformat(L"%-4s %-64s %-128s") 
                % ptr->GetStatusText()
                % ptr->GetName()
                % ptr->GetDisplayName()
            << std::endl;
    }

    return EXIT_SUCCESS;
}