150 lines
3.9 KiB
Go
150 lines
3.9 KiB
Go
package winio
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"runtime"
|
|
"syscall"
|
|
"unicode/utf16"
|
|
)
|
|
|
|
//sys adjustTokenPrivileges(token syscall.Handle, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (success bool, err error) [true] = advapi32.AdjustTokenPrivileges
|
|
//sys impersonateSelf(level uint32) (err error) = advapi32.ImpersonateSelf
|
|
//sys revertToSelf() (err error) = advapi32.RevertToSelf
|
|
//sys openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *syscall.Handle) (err error) = advapi32.OpenThreadToken
|
|
//sys getCurrentThread() (h syscall.Handle) = GetCurrentThread
|
|
//sys lookupPrivilegeValue(systemName string, name string, luid *uint64) (err error) = advapi32.LookupPrivilegeValueW
|
|
//sys lookupPrivilegeName(systemName string, luid *uint64, buffer *uint16, size *uint32) (err error) = advapi32.LookupPrivilegeNameW
|
|
//sys lookupPrivilegeDisplayName(systemName string, name *uint16, buffer *uint16, size *uint32, languageId *uint32) (err error) = advapi32.LookupPrivilegeDisplayNameW
|
|
|
|
const (
|
|
SE_PRIVILEGE_ENABLED = 2
|
|
|
|
ERROR_NOT_ALL_ASSIGNED syscall.Errno = 1300
|
|
|
|
SeBackupPrivilege = "SeBackupPrivilege"
|
|
SeRestorePrivilege = "SeRestorePrivilege"
|
|
)
|
|
|
|
const (
|
|
securityAnonymous = iota
|
|
securityIdentification
|
|
securityImpersonation
|
|
securityDelegation
|
|
)
|
|
|
|
type PrivilegeError struct {
|
|
privileges []uint64
|
|
}
|
|
|
|
func (e *PrivilegeError) Error() string {
|
|
s := ""
|
|
if len(e.privileges) > 1 {
|
|
s = "Could not enable privileges "
|
|
} else {
|
|
s = "Could not enable privilege "
|
|
}
|
|
for i, p := range e.privileges {
|
|
if i != 0 {
|
|
s += ", "
|
|
}
|
|
s += `"`
|
|
s += getPrivilegeName(p)
|
|
s += `"`
|
|
}
|
|
return s
|
|
}
|
|
|
|
func RunWithPrivilege(name string, fn func() error) error {
|
|
return RunWithPrivileges([]string{name}, fn)
|
|
}
|
|
|
|
func RunWithPrivileges(names []string, fn func() error) error {
|
|
var privileges []uint64
|
|
for _, name := range names {
|
|
p := uint64(0)
|
|
err := lookupPrivilegeValue("", name, &p)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
privileges = append(privileges, p)
|
|
}
|
|
runtime.LockOSThread()
|
|
defer runtime.UnlockOSThread()
|
|
token, err := newThreadToken()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer releaseThreadToken(token)
|
|
err = adjustPrivileges(token, privileges)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return fn()
|
|
}
|
|
|
|
func adjustPrivileges(token syscall.Handle, privileges []uint64) error {
|
|
var b bytes.Buffer
|
|
binary.Write(&b, binary.LittleEndian, uint32(len(privileges)))
|
|
for _, p := range privileges {
|
|
binary.Write(&b, binary.LittleEndian, p)
|
|
binary.Write(&b, binary.LittleEndian, uint32(SE_PRIVILEGE_ENABLED))
|
|
}
|
|
prevState := make([]byte, b.Len())
|
|
reqSize := uint32(0)
|
|
success, err := adjustTokenPrivileges(token, false, &b.Bytes()[0], uint32(len(prevState)), &prevState[0], &reqSize)
|
|
if !success {
|
|
return err
|
|
}
|
|
if err == ERROR_NOT_ALL_ASSIGNED {
|
|
return &PrivilegeError{privileges}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func getPrivilegeName(luid uint64) string {
|
|
var nameBuffer [256]uint16
|
|
bufSize := uint32(len(nameBuffer))
|
|
err := lookupPrivilegeName("", &luid, &nameBuffer[0], &bufSize)
|
|
if err != nil {
|
|
return fmt.Sprintf("<unknown privilege %d>", luid)
|
|
}
|
|
|
|
var displayNameBuffer [256]uint16
|
|
displayBufSize := uint32(len(displayNameBuffer))
|
|
var langId uint32
|
|
err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langId)
|
|
if err != nil {
|
|
return fmt.Sprintf("<unknown privilege %s>", utf16.Decode(nameBuffer[:bufSize]))
|
|
}
|
|
|
|
return string(utf16.Decode(displayNameBuffer[:displayBufSize]))
|
|
}
|
|
|
|
func newThreadToken() (syscall.Handle, error) {
|
|
err := impersonateSelf(securityImpersonation)
|
|
if err != nil {
|
|
panic(err)
|
|
return 0, err
|
|
}
|
|
|
|
var token syscall.Handle
|
|
err = openThreadToken(getCurrentThread(), syscall.TOKEN_ADJUST_PRIVILEGES|syscall.TOKEN_QUERY, false, &token)
|
|
if err != nil {
|
|
rerr := revertToSelf()
|
|
if rerr != nil {
|
|
panic(rerr)
|
|
}
|
|
return 0, err
|
|
}
|
|
return token, nil
|
|
}
|
|
|
|
func releaseThreadToken(h syscall.Handle) {
|
|
err := revertToSelf()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
syscall.Close(h)
|
|
}
|