module Puppet::Util::Windows::Security

# File lib/puppet/util/windows/security.rb, line 416
def add_access_allowed_ace(acl, mask, sid, inherit = NO_INHERITANCE)
  string_to_sid_ptr(sid) do |sid_ptr|
    raise Puppet::Util::Windows::Error.new("Invalid SID") unless IsValidSid(sid_ptr)

    unless AddAccessAllowedAceEx(acl, ACL_REVISION, inherit, mask, sid_ptr)
      raise Puppet::Util::Windows::Error.new("Failed to add access control entry")
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 426
def add_access_denied_ace(acl, mask, sid)
  string_to_sid_ptr(sid) do |sid_ptr|
    raise Puppet::Util::Windows::Error.new("Invalid SID") unless IsValidSid(sid_ptr)

    unless AddAccessDeniedAce(acl, ACL_REVISION, mask, sid_ptr)
      raise Puppet::Util::Windows::Error.new("Failed to add access control entry")
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 198
def add_attributes(path, flags)
  oldattrs = get_attributes(path)

  if (oldattrs | flags) != oldattrs
    set_attributes(path, oldattrs | flags)
  end
end

# File lib/puppet/util/windows/security.rb, line 165
def change_sid(old_sid, new_sid, info, path)
  if old_sid != new_sid
    mode = get_mode(path)

    string_to_sid_ptr(new_sid) do |psid|
      with_privilege(SE_RESTORE_NAME) do
        open_file(path, WRITE_OWNER) do |handle|
          set_security_info(handle, info, psid)
        end
      end
    end

    # rebuild dacl now that sid has changed
    set_mode(mode, path)
  end
end

# File lib/puppet/util/windows/security.rb, line 190
def get_attributes(path)
  attributes = GetFileAttributes(path)

  raise Puppet::Util::Windows::Error.new("Failed to get file attributes") if attributes == INVALID_FILE_ATTRIBUTES

  attributes
end

# File lib/puppet/util/windows/security.rb, line 436
def get_dacl(handle)
  get_dacl_ptr(handle) do |dacl_ptr|
    # REMIND: need to handle NULL DACL
    raise Puppet::Util::Windows::Error.new("Invalid DACL") unless IsValidAcl(dacl_ptr)

    # ACL structure, size and count are the important parts. The
    # size includes both the ACL structure and all the ACEs.
    #
    # BYTE AclRevision
    # BYTE Padding1
    # WORD AclSize
    # WORD AceCount
    # WORD Padding2
    acl_buf = 0.chr * 8
    memcpy(acl_buf, dacl_ptr, acl_buf.size)
    ace_count = acl_buf.unpack('CCSSS')[3]

    dacl = []

    # deny all
    return dacl if ace_count == 0

    0.upto(ace_count - 1) do |i|
      ace_ptr = [0].pack('L')

      next unless GetAce(dacl_ptr, i, ace_ptr)

      # ACE structures vary depending on the type. All structures
      # begin with an ACE header, which specifies the type, flags
      # and size of what follows. We are only concerned with
      # ACCESS_ALLOWED_ACE and ACCESS_DENIED_ACEs, which have the
      # same structure:
      #
      # BYTE  C AceType
      # BYTE  C AceFlags
      # WORD  S AceSize
      # DWORD L ACCESS_MASK
      # DWORD L Sid
      # ..      ...
      # DWORD L Sid

      ace_buf = 0.chr * 8
      memcpy(ace_buf, ace_ptr.unpack('L')[0], ace_buf.size)

      ace_type, ace_flags, size, mask = ace_buf.unpack('CCSL')

      # skip aces that only serve to propagate inheritance
      next if (ace_flags & INHERIT_ONLY_ACE).nonzero?

      case ace_type
      when ACCESS_ALLOWED_ACE_TYPE
        sid_ptr = ace_ptr.unpack('L')[0] + 8 # address of ace_ptr->SidStart
        raise Puppet::Util::Windows::Error.new("Failed to read DACL, invalid SID") unless IsValidSid(sid_ptr)
        sid = sid_ptr_to_string(sid_ptr)
        dacl << {:sid => sid, :type => ace_type, :mask => mask}
      else
        Puppet.warning "Unsupported access control entry type: 0x#{ace_type.to_s(16)}"
      end
    end

    dacl
  end
end

# File lib/puppet/util/windows/security.rb, line 500
def get_dacl_ptr(handle)
  dacl = [0].pack('L')
  sd = [0].pack('L')

  rv = GetSecurityInfo(
       handle,
       SE_FILE_OBJECT,
       DACL_SECURITY_INFORMATION,
       nil,
       nil,
       dacl, #dacl
       nil, #sacl
       sd) #sec desc
  raise Puppet::Util::Windows::Error.new("Failed to get DACL") unless rv == ERROR_SUCCESS
  begin
    yield dacl.unpack('L')[0]
  ensure
    LocalFree(sd.unpack('L')[0])
  end
end

# File lib/puppet/util/windows/security.rb, line 146
def get_group(path)
  return unless supports_acl?(path)

  get_sid(GROUP_SECURITY_INFORMATION, path)
end

# File lib/puppet/util/windows/security.rb, line 230
def get_mode(path)
  return unless supports_acl?(path)

  owner_sid = get_owner(path)
  group_sid = get_group(path)
  well_known_world_sid = Win32::Security::SID::Everyone
  well_known_nobody_sid = Win32::Security::SID::Nobody

  with_privilege(SE_BACKUP_NAME) do
    open_file(path, READ_CONTROL) do |handle|
      mode = 0

      get_dacl(handle).each do |ace|
        case ace[:sid]
        when owner_sid
          MASK_TO_MODE.each_pair do |k,v|
            if (ace[:mask] & k) == k
              mode |= (v << 6)
            end
          end
        when group_sid
          MASK_TO_MODE.each_pair do |k,v|
            if (ace[:mask] & k) == k
              mode |= (v << 3)
            end
          end
        when well_known_world_sid
          MASK_TO_MODE.each_pair do |k,v|
            if (ace[:mask] & k) == k
              mode |= (v << 6) | (v << 3) | v
            end
          end
          if File.directory?(path) and (ace[:mask] & (FILE_WRITE_DATA | FILE_EXECUTE | FILE_DELETE_CHILD)) == (FILE_WRITE_DATA | FILE_EXECUTE)
            mode |= S_ISVTX;
          end
        when well_known_nobody_sid
          if (ace[:mask] & FILE_APPEND_DATA).nonzero?
            mode |= S_ISVTX
          end
        else
          #puts "Warning, unable to map SID into POSIX mode: #{ace[:sid]}"
          mode |= S_IEXTRA
        end

        # if owner and group the same, then user and group modes are the OR of both
        if owner_sid == group_sid
          mode |= ((mode & S_IRWXG) << 3) | ((mode & S_IRWXU) >> 3)
          #puts "owner: #{group_sid}, 0x#{ace[:mask].to_s(16)}, #{mode.to_s(8)}"
        end
      end

      #puts "get_mode: #{mode.to_s(8)}"
      mode
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 124
def get_owner(path)
  return unless supports_acl?(path)

  get_sid(OWNER_SECURITY_INFORMATION, path)
end

# File lib/puppet/util/windows/security.rb, line 536
def get_security_info(handle, info)
  sid = [0].pack('L')
  sd = [0].pack('L')

  rv = GetSecurityInfo(
       handle,
       SE_FILE_OBJECT,
       info, # security info
       info == OWNER_SECURITY_INFORMATION ? sid : nil,
       info == GROUP_SECURITY_INFORMATION ? sid : nil,
       nil, #dacl
       nil, #sacl
       sd) #sec desc
  raise Puppet::Util::Windows::Error.new("Failed to get security information") unless rv == ERROR_SUCCESS

  begin
    return sid_ptr_to_string(sid.unpack('L')[0])
  ensure
    LocalFree(sd.unpack('L')[0])
  end
end

# File lib/puppet/util/windows/security.rb, line 182
def get_sid(info, path)
  with_privilege(SE_BACKUP_NAME) do
    open_file(path, READ_CONTROL) do |handle|
      get_security_info(handle, info)
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 560
def open_file(path, access)
  handle = CreateFile(
           path,
           access,
           FILE_SHARE_READ | FILE_SHARE_WRITE,
           0, # security_attributes
           OPEN_EXISTING,
           FILE_FLAG_BACKUP_SEMANTICS,
           0) # template
  raise Puppet::Util::Windows::Error.new("Failed to open '#{path}'") if handle == INVALID_HANDLE_VALUE
  begin
    yield handle
  ensure
    CloseHandle(handle)
  end
end

# File lib/puppet/util/windows/security.rb, line 206
def remove_attributes(path, flags)
  oldattrs = get_attributes(path)

  if (oldattrs & ~flags) != oldattrs
    set_attributes(path, oldattrs & ~flags)
  end
end

# File lib/puppet/util/windows/security.rb, line 391
def set_acl(path, protected = true)
  with_privilege(SE_BACKUP_NAME) do
    with_privilege(SE_RESTORE_NAME) do
      open_file(path, READ_CONTROL | WRITE_DAC) do |handle|
        acl = 0.chr * 1024 # This can be increased later as needed

        unless InitializeAcl(acl, acl.size, ACL_REVISION)
          raise Puppet::Util::Windows::Error.new("Failed to initialize ACL")
        end

        raise Puppet::Util::Windows::Error.new("Invalid DACL") unless IsValidAcl(acl)

        yield acl

        # protected means the object does not inherit aces from its parent
        info = DACL_SECURITY_INFORMATION
        info |= protected ? PROTECTED_DACL_SECURITY_INFORMATION : UNPROTECTED_DACL_SECURITY_INFORMATION

        # set the DACL
        set_security_info(handle, info, acl)
      end
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 214
def set_attributes(path, flags)
  raise Puppet::Util::Windows::Error.new("Failed to set file attributes") unless SetFileAttributes(path, flags)
end

# File lib/puppet/util/windows/security.rb, line 135
def set_group(group_sid, path)
  old_sid = get_group(path)

  change_sid(old_sid, group_sid, GROUP_SECURITY_INFORMATION, path)
end

# File lib/puppet/util/windows/security.rb, line 306
def set_mode(mode, path, protected = true)
  owner_sid = get_owner(path)
  group_sid = get_group(path)
  well_known_world_sid = Win32::Security::SID::Everyone
  well_known_nobody_sid = Win32::Security::SID::Nobody

  owner_allow = STANDARD_RIGHTS_ALL  | FILE_READ_ATTRIBUTES | FILE_WRITE_ATTRIBUTES
  group_allow = STANDARD_RIGHTS_READ | FILE_READ_ATTRIBUTES | SYNCHRONIZE
  other_allow = STANDARD_RIGHTS_READ | FILE_READ_ATTRIBUTES | SYNCHRONIZE
  nobody_allow = 0

  MODE_TO_MASK.each do |k,v|
    if ((mode >> 6) & k) == k
      owner_allow |= v
    end
    if ((mode >> 3) & k) == k
      group_allow |= v
    end
    if (mode & k) == k
      other_allow |= v
    end
  end

  if (mode & S_ISVTX).nonzero?
    nobody_allow |= FILE_APPEND_DATA;
  end

  isdir = File.directory?(path)

  if isdir
    if (mode & (S_IWUSR | S_IXUSR)) == (S_IWUSR | S_IXUSR)
      owner_allow |= FILE_DELETE_CHILD
    end
    if (mode & (S_IWGRP | S_IXGRP)) == (S_IWGRP | S_IXGRP) and (mode & S_ISVTX) == 0
      group_allow |= FILE_DELETE_CHILD
    end
    if (mode & (S_IWOTH | S_IXOTH)) == (S_IWOTH | S_IXOTH) and (mode & S_ISVTX) == 0
      other_allow |= FILE_DELETE_CHILD
    end
  end

  # if owner and group the same, then map group permissions to the one owner ACE
  isownergroup = owner_sid == group_sid
  if isownergroup
    owner_allow |= group_allow
  end

  # if any ACE allows write, then clear readonly bit, but do this before we overwrite
  # the DACl and lose our ability to set the attribute
  if ((owner_allow | group_allow | other_allow ) & FILE_WRITE_DATA) == FILE_WRITE_DATA
    remove_attributes(path, FILE_ATTRIBUTE_READONLY)
  end

  set_acl(path, protected) do |acl|
    #puts "ace: owner #{owner_sid}, mask 0x#{owner_allow.to_s(16)}"
    add_access_allowed_ace(acl, owner_allow, owner_sid)

    unless isownergroup
      #puts "ace: group #{group_sid}, mask 0x#{group_allow.to_s(16)}"
      add_access_allowed_ace(acl, group_allow, group_sid)
    end

    #puts "ace: other #{well_known_world_sid}, mask 0x#{other_allow.to_s(16)}"
    add_access_allowed_ace(acl, other_allow, well_known_world_sid)

    #puts "ace: nobody #{well_known_nobody_sid}, mask 0x#{nobody_allow.to_s(16)}"
    add_access_allowed_ace(acl, nobody_allow, well_known_nobody_sid)

    # add inherit-only aces for child dirs and files that are created within the dir
    if isdir
      inherit = INHERIT_ONLY_ACE | CONTAINER_INHERIT_ACE
      add_access_allowed_ace(acl, owner_allow, Win32::Security::SID::CreatorOwner, inherit)
      add_access_allowed_ace(acl, group_allow, Win32::Security::SID::CreatorGroup, inherit)

      inherit = INHERIT_ONLY_ACE |  OBJECT_INHERIT_ACE
      add_access_allowed_ace(acl, owner_allow & ~FILE_EXECUTE, Win32::Security::SID::CreatorOwner, inherit)
      add_access_allowed_ace(acl, group_allow & ~FILE_EXECUTE, Win32::Security::SID::CreatorGroup, inherit)
    end
  end

  nil
end

# File lib/puppet/util/windows/security.rb, line 113
def set_owner(owner_sid, path)
  old_sid = get_owner(path)

  change_sid(old_sid, owner_sid, OWNER_SECURITY_INFORMATION, path)
end

# File lib/puppet/util/windows/security.rb, line 587
def set_privilege(privilege, enable)
  return unless Puppet.features.root?

  with_process_token(TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY) do |token|
    tmpLuid = 0.chr * 8

    # Get the LUID for specified privilege.
    unless LookupPrivilegeValue("", privilege, tmpLuid)
      raise Puppet::Util::Windows::Error.new("Failed to lookup privilege")
    end

    # DWORD + [LUID + DWORD]
    tkp = [1].pack('L') + tmpLuid + [enable ? SE_PRIVILEGE_ENABLED : 0].pack('L')

    unless AdjustTokenPrivileges(token, 0, tkp, tkp.length , nil, nil)
      raise Puppet::Util::Windows::Error.new("Failed to adjust process privileges")
    end
  end
end

# File lib/puppet/util/windows/security.rb, line 522
def set_security_info(handle, info, ptr)
  rv = SetSecurityInfo(
       handle,
       SE_FILE_OBJECT,
       info,
       (info & OWNER_SECURITY_INFORMATION) == OWNER_SECURITY_INFORMATION ? ptr : nil,
       (info & GROUP_SECURITY_INFORMATION) == GROUP_SECURITY_INFORMATION ? ptr : nil,
       (info & DACL_SECURITY_INFORMATION) == DACL_SECURITY_INFORMATION ? ptr : nil,
       nil)
  raise Puppet::Util::Windows::Error.new("Failed to set security information") unless rv == ERROR_SUCCESS
end

# File lib/puppet/util/windows/security.rb, line 152
def supports_acl?(path)
  flags = 0.chr * 4

  root = Pathname.new(path).enum_for(:ascend).to_a.last.to_s
  # 'A trailing backslash is required'
  root = "#{root}\\" unless root =~ /[\/\]$/
  unless GetVolumeInformation(root, nil, 0, nil, nil, flags, nil, 0)
    raise Puppet::Util::Windows::Error.new("Failed to get volume information")
  end

  (flags.unpack('L')[0] & Windows::File::FILE_PERSISTENT_ACLS) != 0
end

# File lib/puppet/util/windows/security.rb, line 578
def with_privilege(privilege)
  set_privilege(privilege, true)
  yield
ensure
  set_privilege(privilege, false)
end

# File lib/puppet/util/windows/security.rb, line 608
def with_process_token(access)
  token = 0.chr * 4

  unless OpenProcessToken(GetCurrentProcess(), access, token)
    raise Puppet::Util::Windows::Error.new("Failed to open process token")
  end
  begin
    token = token.unpack('L')[0]

    yield token
  ensure
    CloseHandle(token)
  end
end