Add custom charsets to fqdn_rand_string()

Author: elyscape

README.markdown

@@ -195,30 +195,17 @@ also appear in the second array. For example, `difference(["a","b","c"],["b","c"
 * `floor`: Returns the largest integer less than or equal to the argument.
 Takes a single numeric value as an argument. *Type*: rvalue
 
-* `fqdn_rand_base64`: Generates a random base64 string, combining the `$fqdn` fact and an optional seed for repeatable randomness.
+* `fqdn_rand_string`: Generates a random string using an optionally-specified character set (default is alphanumeric), combining the `$fqdn` fact and an optional seed for repeatable randomness.
 
   *Usage:*
   ```
-  fqdn_rand_base64(LENGTH, [SEED])
-  ```
-  *Examples:*
-  ```
-  fqdn_rand_base64(10)
-  fqdn_rand_base64(10, 'custom seed')
-  ```
-
-  *Type*: rvalue
-
-* `fqdn_rand_string`: Generates a random alphanumeric string, combining the `$fqdn` fact and an optional seed for repeatable randomness.
-
-  *Usage:*
-  ```
-  fqdn_rand_string(LENGTH, [SEED])
+  fqdn_rand_string(LENGTH, [CHARSET], [SEED])
   ```
   *Examples:*
   ```
   fqdn_rand_string(10)
-  fqdn_rand_string(10, 'custom seed')
+  fqdn_rand_string(10, 'ABCDEF!@#$%^')
+  fqdn_rand_string(10, '', 'custom seed')
   ```
 
   *Type*: rvalue

lib/puppet/parser/functions/fqdn_rand_base64.rb

@@ -2,23 +2,28 @@
   :fqdn_rand_string,
   :arity => -2,
   :type => :rvalue,
-  :doc => "Usage: `fqdn_rand_string(LENGTH, [SEED])`. LENGTH is required and
-  must be a positive integer; SEED is optional and may be any number or string.
+  :doc => "Usage: `fqdn_rand_string(LENGTH, [CHARSET], [SEED])`. LENGTH is
+  required and must be a positive integer. CHARSET is optional and may be
+  `undef` or a string. SEED is optional and may be any number or string.
 
-  Generates a random alphanumeric string LENGTH characters long, combining the
-  `$fqdn` fact and the value of SEED for repeatable randomness.  (That is, each
-  node will get a different random string from this function, but a given
-  node's result will be the same every time unless its hostname changes.)
-  Adding a SEED can be useful if you need more than one unrelated string.") do |args|
+  Generates a random string LENGTH characters long using the character set
+  provided by CHARSET, combining the `$fqdn` fact and the value of SEED for
+  repeatable randomness. (That is, each node will get a different random
+  string from this function, but a given node's result will be the same every
+  time unless its hostname changes.) Adding a SEED can be useful if you need
+  more than one unrelated string. CHARSET will default to alphanumeric if
+  `undef` or an empty string.") do |args|
     raise(ArgumentError, "fqdn_rand_string(): wrong number of arguments (0 for 1)") if args.size == 0
     Puppet::Parser::Functions.function('is_integer')
     raise(ArgumentError, "fqdn_rand_base64(): first argument must be a positive integer") unless function_is_integer([args[0]]) and args[0].to_i > 0
+    raise(ArgumentError, "fqdn_rand_base64(): second argument must be undef or a string") unless args[1].nil? or args[1].is_a? String
 
     Puppet::Parser::Functions.function('fqdn_rand')
 
     length = args.shift.to_i
+    charset = args.shift.to_s.chars.to_a
 
-    charset = (0..9).map { |i| i.to_s } + ('A'..'Z').to_a + ('a'..'z').to_a
+    charset = (0..9).map { |i| i.to_s } + ('A'..'Z').to_a + ('a'..'z').to_a if charset.empty?
 
     rand_string = ''
     for current in 1..length

lib/puppet/parser/functions/fqdn_rand_string.rb

@@ -18,9 +18,29 @@
     expect { fqdn_rand_string(0.5) }.to( raise_error(ArgumentError, /first argument must be a positive integer/))
   end
 
-  it "provides a valid alphanumeric string" do
-    string = %r{\A[a-zA-Z0-9]*\z}
-    expect(fqdn_rand_string(100).match(string)).not_to eq(nil)
+  it "provides a valid alphanumeric string when no character set is provided" do
+    length = 100
+    string = %r{\A[a-zA-Z0-9]{#{length}}\z}
+    expect(fqdn_rand_string(length).match(string)).not_to eq(nil)
+  end
+
+  it "provides a valid alphanumeric string when an :undef character set is provided" do
+    length = 100
+    string = %r{\A[a-zA-Z0-9]{#{length}}\z}
+    expect(fqdn_rand_string(length, :charset => nil).match(string)).not_to eq(nil)
+  end
+
+  it "provides a valid alphanumeric string when an empty character set is provided" do
+    length = 100
+    string = %r{\A[a-zA-Z0-9]{#{length}}\z}
+    expect(fqdn_rand_string(length, :charset => '').match(string)).not_to eq(nil)
+  end
+
+  it "uses a provided character set" do
+    length = 100
+    charset = '!@#$%^&*()-_=+'
+    string = %r{\A[#{charset}]{#{length}}\z}
+    expect(fqdn_rand_string(length, :charset => charset).match(string)).not_to eq(nil)
   end
 
   it "provides a random string exactly as long as the given length" do
@@ -54,12 +74,18 @@
 
   def fqdn_rand_string(max, args = {})
     host = args[:host] || '127.0.0.1'
+    charset = args[:charset]
     extra = args[:extra_identifier] || []
 
     scope = PuppetlabsSpec::PuppetInternals.scope
     scope.stubs(:[]).with("::fqdn").returns(host)
     scope.stubs(:lookupvar).with("::fqdn").returns(host)
 
-    scope.function_fqdn_rand_string([max] + extra)
+    function_args = [max]
+    if args.has_key?(:charset) or !extra.empty?
+      function_args << charset
+    end
+    function_args += extra
+    scope.function_fqdn_rand_string(function_args)
   end
 end