From 1a7a4894b1e19f07a4c3dd0db2c9df9c4db162cd Mon Sep 17 00:00:00 2001
From: Zach Leslie <xaque208@gmail.com>
Date: Sun, 8 May 2016 12:43:16 -0700
Subject: [PATCH] Add support for scoped queries

This work adds support for setting the scope of a given query.  This
allows retrieval of a specific objects, or searching at a specific level
of the tree for the desired results.
---
 lib/puppet/parser/functions/ldapquery.rb | 15 ++++++++--
 lib/puppet_x/ldapquery.rb                | 37 +++++++++++++++++-------
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/lib/puppet/parser/functions/ldapquery.rb b/lib/puppet/parser/functions/ldapquery.rb
index 5be8370..7518fd9 100644
--- a/lib/puppet/parser/functions/ldapquery.rb
+++ b/lib/puppet/parser/functions/ldapquery.rb
@@ -1,3 +1,11 @@
+# Provides a query interface to an LDAP server
+#
+# @example simple query
+#   ldapquery("(objectClass=dnsDomain)", ['dc'])
+#
+# @example more complex query for ssh public keys
+#   ldapquery('(&(objectClass=ldapPublicKey)(sshPublicKey=*)(objectClass=posixAccount))', ['uid', 'sshPublicKey'])
+#
 require_relative '../../../puppet_x/ldapquery'
 
 begin
@@ -9,14 +17,15 @@ end
 Puppet::Parser::Functions.newfunction(:ldapquery,
                                       :type => :rvalue) do |args|
 
-  if args.size > 3
+  if args.size > 4
     raise Puppet::ParseError, "Too many arguments received in ldapquery()"
   end
 
-  filter, attributes, base = args
+  filter, attributes, base, scope = args
 
   attributes ||= []
   base ||= Puppet[:ldapbase]
+  scope ||= 'sub'
 
-  return PuppetX::LDAPquery.new(filter, attributes, base).results
+  return PuppetX::LDAPquery.new(filter, attributes, base, scope).results
 end
diff --git a/lib/puppet_x/ldapquery.rb b/lib/puppet_x/ldapquery.rb
index 6f02408..c27b7b4 100644
--- a/lib/puppet_x/ldapquery.rb
+++ b/lib/puppet_x/ldapquery.rb
@@ -8,11 +8,24 @@ module PuppetX
     def initialize(
       filter,
       attributes=[],
-      base=Puppet[:ldapbase]
+      base=Puppet[:ldapbase],
+      scope='sub'
     )
       @filter = filter
       @attributes = attributes
       @base = base
+
+      if scope
+        if scope == 'sub'
+          @scope = Net::LDAP::SearchScope_WholeSubtree
+        elsif scope == 'base'
+          @scope = Net::LDAP::SearchScope_BaseObject
+        elsif scope == 'single'
+          @scope = Net::LDAP::SearchScope_SingleLevel
+        else
+          raise Puppet::ParseError, 'Received param "scope" not one of ["sub","base","single"]'
+        end
+      end
     end
 
     def get_config
@@ -66,23 +79,27 @@ module PuppetX
       # Query the LDAP server for attributes using the filter
       #
       # Returns: An array of Net::LDAP::Entry objects
-      ldapfilter = @filter
-      attributes = @attributes
-      base = @base
-
       conf = self.get_config()
 
       start_time = Time.now
       ldap = Net::LDAP.new(conf)
-      ldapfilter = Net::LDAP::Filter.construct(@filter)
+
+      search_args = {
+          :base => @base,
+          :attributes => @attributes,
+          :scope => @scope,
+          :time => 10,
+      }
+
+      if @filter and @filter.length > 0
+        ldapfilter = Net::LDAP::Filter.construct(@filter)
+        search_args[:filter] = ldapfilter
+      end
 
       entries = []
 
       begin
-        ldap.search(:base => base,
-                    :filter => ldapfilter,
-                    :attributes => attributes,
-                    :time => 10) do |entry|
+        ldap.search(search_args) do |entry|
           entries << entry
         end
         end_time = Time.now