//go:build windows

package safeio

import (
	"path/filepath"
	"runtime"
	"strings"
	"unsafe"

	"golang.org/x/sys/windows"
)

// ProtectPrivateDirectory sets a protected DACL containing only the current owner.
func ProtectPrivateDirectory(path string) error {
	parents, target, err := openCanonicalWindowsParent(path)
	if err != nil {
		return ErrUnsafeFile
	}
	defer parents.Close()
	handle, err := openWindowsComponentWithAccess(filepath.Join(parents.directory, target), true, windows.GENERIC_READ|windows.WRITE_DAC|windows.WRITE_OWNER)
	if err != nil {
		return ErrUnsafeFile
	}
	defer windows.CloseHandle(handle)
	if err := setOwnerOnlyDACL(handle); err != nil {
		return ErrUnsafeFile
	}
	return validateOwnerOnlyDACL(handle)
}

// ValidatePrivateDirectory requires a canonical directory protected for its current owner only.
func ValidatePrivateDirectory(path string) error {
	parents, target, err := openCanonicalWindowsParent(path)
	if err != nil {
		return ErrUnsafeFile
	}
	defer parents.Close()
	handle, err := openWindowsComponent(filepath.Join(parents.directory, target), true)
	if err != nil {
		return ErrUnsafeFile
	}
	defer windows.CloseHandle(handle)
	if err := validateOwnerOnlyDACL(handle); err != nil {
		return ErrUnsafeFile
	}
	return nil
}

// ProtectPrivateRegular sets a protected DACL containing only the current owner.
func ProtectPrivateRegular(path string) error {
	parents, target, err := openCanonicalWindowsParent(path)
	if err != nil {
		return ErrUnsafeFile
	}
	defer parents.Close()
	handle, err := openWindowsComponentWithAccess(filepath.Join(parents.directory, target), false, windows.GENERIC_READ|windows.WRITE_DAC|windows.WRITE_OWNER)
	if err != nil {
		return ErrUnsafeFile
	}
	defer windows.CloseHandle(handle)
	if err := setOwnerOnlyDACL(handle); err != nil {
		return ErrUnsafeFile
	}
	return validateOwnerOnlyDACL(handle)
}

// ValidatePrivateRegular requires a canonical, single-link file protected for its current owner only.
func ValidatePrivateRegular(path string) error {
	parents, target, err := openCanonicalWindowsParent(path)
	if err != nil {
		return ErrUnsafeFile
	}
	defer parents.Close()
	handle, err := openWindowsComponent(filepath.Join(parents.directory, target), false)
	if err != nil {
		return ErrUnsafeFile
	}
	defer windows.CloseHandle(handle)
	if err := validateOwnerOnlyDACL(handle); err != nil {
		return ErrUnsafeFile
	}
	return nil
}

type windowsParentHandles struct {
	directory string
	handles   []windows.Handle
}

func (parents *windowsParentHandles) Close() {
	for index := len(parents.handles) - 1; index >= 0; index-- {
		_ = windows.CloseHandle(parents.handles[index])
	}
}

// openCanonicalWindowsParent retains every directory handle from the volume root through the
// target parent without FILE_SHARE_DELETE. The resulting parent cannot be renamed or replaced by
// a reparse point while an operation uses its absolute child paths.
func openCanonicalWindowsParent(path string) (*windowsParentHandles, string, error) {
	if err := ValidateCanonicalPath(path); err != nil {
		return nil, "", err
	}
	volume := filepath.VolumeName(path)
	root := volume + `\`
	components := strings.Split(strings.TrimPrefix(path, root), `\`)
	if volume == "" || len(components) == 0 || components[0] == "" {
		return nil, "", ErrUnsafeFile
	}
	parents := &windowsParentHandles{directory: root}
	rootHandle, err := openWindowsComponent(root, true)
	if err != nil {
		return nil, "", err
	}
	parents.handles = append(parents.handles, rootHandle)
	for _, component := range components[:len(components)-1] {
		parents.directory = filepath.Join(parents.directory, component)
		handle, err := openWindowsComponent(parents.directory, true)
		if err != nil {
			parents.Close()
			return nil, "", err
		}
		parents.handles = append(parents.handles, handle)
	}
	return parents, components[len(components)-1], nil
}

func setOwnerOnlyDACL(handle windows.Handle) error {
	sid, err := currentOwnerSID()
	if err != nil {
		return err
	}
	var pinner runtime.Pinner
	pinner.Pin(sid)
	defer pinner.Unpin()
	acl, err := windows.ACLFromEntries([]windows.EXPLICIT_ACCESS{{
		AccessPermissions: windows.GENERIC_ALL,
		AccessMode:        windows.GRANT_ACCESS,
		Trustee: windows.TRUSTEE{
			TrusteeForm:  windows.TRUSTEE_IS_SID,
			TrusteeType:  windows.TRUSTEE_IS_USER,
			TrusteeValue: windows.TrusteeValueFromSID(sid),
		},
	}}, nil)
	if err != nil {
		return err
	}
	return windows.SetSecurityInfo(handle, windows.SE_FILE_OBJECT,
		windows.OWNER_SECURITY_INFORMATION|windows.DACL_SECURITY_INFORMATION|windows.PROTECTED_DACL_SECURITY_INFORMATION,
		sid, nil, acl, nil)
}

func validateOwnerOnlyDACL(handle windows.Handle) error {
	ownerSID, err := currentOwnerSID()
	if err != nil {
		return err
	}
	descriptor, err := windows.GetSecurityInfo(handle, windows.SE_FILE_OBJECT,
		windows.OWNER_SECURITY_INFORMATION|windows.DACL_SECURITY_INFORMATION)
	if err != nil || descriptor == nil {
		return ErrUnsafeFile
	}
	owner, _, err := descriptor.Owner()
	if err != nil || owner == nil || !windows.EqualSid(owner, ownerSID) {
		return ErrUnsafeFile
	}
	control, _, err := descriptor.Control()
	if err != nil || control&windows.SE_DACL_PROTECTED == 0 {
		return ErrUnsafeFile
	}
	dacl, defaulted, err := descriptor.DACL()
	if err != nil || defaulted || dacl == nil || dacl.AceCount != 1 {
		return ErrUnsafeFile
	}
	var ace *windows.ACCESS_ALLOWED_ACE
	if err := windows.GetAce(dacl, 0, &ace); err != nil || ace == nil || ace.Header.AceType != windows.ACCESS_ALLOWED_ACE_TYPE || ace.Header.AceFlags != 0 || ace.Mask != windows.GENERIC_ALL {
		return ErrUnsafeFile
	}
	aceSID := (*windows.SID)(unsafe.Pointer(&ace.SidStart))
	if !windows.EqualSid(aceSID, ownerSID) {
		return ErrUnsafeFile
	}
	return nil
}

func currentOwnerSID() (*windows.SID, error) {
	user, err := windows.GetCurrentProcessToken().GetTokenUser()
	if err != nil || user == nil || user.User.Sid == nil {
		return nil, ErrUnsafeFile
	}
	return user.User.Sid, nil
}