Land rapid7#2087, @egypt's random ID generator

Author: Tod Beardsley

lib/rex/random_identifier_generator.rb

@@ -0,0 +1,168 @@
+
+# A quick way to produce unique random strings that follow the rules of
+# identifiers, i.e., begin with a letter and contain only alphanumeric
+# characters and underscore.
+#
+# The advantage of using this class over, say, {Rex::Text.rand_text_alpha}
+# each time you need a new identifier is that it ensures you don't have
+# collisions.
+#
+# @example
+#   vars = Rex::RandomIdentifierGenerator.new
+#   asp_code = <<-END_CODE
+#     Sub #{vars[:func]}()
+#       Dim #{vars[:fso]}
+#       Set #{vars[:fso]} = CreateObject("Scripting.FileSystemObject")
+#       ...
+#     End Sub
+#     #{vars[:func]}
+#   END_CODE
+#
+class Rex::RandomIdentifierGenerator
+
+	# Raised when a RandomIdentifierGenerator cannot create any more
+	# identifiers without collisions.
+	class ExhaustedSpaceError < StandardError; end
+
+	# Default options
+	DefaultOpts = {
+		# Arbitrary
+		:max_length => 12,
+		:min_length => 3,
+		# This should be pretty universal for identifier rules
+		:char_set => Rex::Text::AlphaNumeric+"_",
+		:first_char_set => Rex::Text::LowerAlpha
+	}
+
+	# @param opts [Hash] Options, see {DefaultOpts} for default values
+	# @option opts :max_length [Fixnum]
+	# @option opts :min_length [Fixnum]
+	# @option opts :char_set [String]
+	def initialize(opts={})
+		# Holds all identifiers.
+		@value_by_name = {}
+		# Inverse of value_by_name so we can ensure uniqueness without
+		# having to search through the whole list of values
+		@name_by_value = {}
+
+		@opts = DefaultOpts.merge(opts)
+		if @opts[:min_length] < 1 || @opts[:max_length] < 1 || @opts[:max_length] < @opts[:min_length]
+			raise ArgumentError, "Invalid length options"
+		end
+
+		# This is really just the maximum number of shortest names. This
+		# will still be a pretty big number most of the time, so don't
+		# bother calculating the real one, which will potentially be
+		# expensive, since we're talking about a 36-digit decimal number to
+		# represent the total possibilities for the range of 10- to
+		# 20-character identifiers.
+		#
+		# 26 because the first char is lowercase alpha, (min_length - 1) and
+		# not just min_length because it includes that first alpha char.
+		@max_permutations = 26 * (@opts[:char_set].length ** (@opts[:min_length]-1))
+		# The real number of permutations could be calculated thusly:
+		#((@opts[:min_length]-1) .. (@opts[:max_length]-1)).reduce(0) { |a, e|
+		#	a + (26 * @opts[:char_set].length ** e)
+		#}
+	end
+
+	# Return a unique random identifier for +name+, generating a new one
+	# if necessary.
+	#
+	# @param name [Symbol] A descriptive, intention-revealing name for an
+	#   identifier. This is what you would normally call the variable if
+	#   you weren't generating it.
+	# @return [String]
+	def get(name)
+		return @value_by_name[name] if @value_by_name[name]
+
+		@value_by_name[name] = generate
+		@name_by_value[@value_by_name[name]] = name
+
+		@value_by_name[name]
+	end
+	alias [] get
+
+	# Add a new identifier. Its name will be checked for uniqueness among
+	# previously-generated names.
+	#
+	# @note This should be called *before* any calls to {#get} to avoid
+	#   potential collisions. If you do hit a collision, this method will
+	#   raise.
+	#
+	# @param name (see #get)
+	# @param value [String] The identifier that will be returned by
+	#   subsequent calls to {#get} with the sane +name+.
+	# @raise RuntimeError if +value+ already exists
+	# @return [void]
+	def store(name, value)
+
+		case @name_by_value[value]
+		when name
+			# we already have this value and it is associated with this name
+			# nothing to do here
+		when nil
+			# don't have this value yet, so go ahead and just insert
+			@value_by_name[name] = value
+			@name_by_value[value] = name
+		else
+			# then the caller is trying to insert a duplicate
+			raise RuntimeError, "Value is not unique!"
+		end
+
+		self
+	end
+
+	# Create a random string that satisfies most languages' requirements
+	# for identifiers. In particular, with a default configuration, the
+	# first character will always be lowercase alpha (unless modified by a
+	# block), and the whole thing will contain only a-zA-Z0-9_ characters.
+	#
+	# If called with a block, the block will be given the identifier before
+	# uniqueness checks. The block's return value will be the new
+	# identifier. Note that the block may be called multiple times if it
+	# returns a non-unique value.
+	#
+	# @note Calling this method with a block that returns only values that
+	#   this generator already contains will result in an infinite loop.
+	#
+	# @example
+	#   rig = Rex::RandomIdentifierGenerator.new
+	#   const = rig.generate { |val| val.capitalize }
+	#   rig.insert(:SOME_CONSTANT, const)
+	#   ruby_code = <<-EOC
+	#     #{rig[:SOME_CONSTANT]} = %q^generated ruby constant^
+	#     def #{rig[:my_method]}; ...; end
+	#   EOC
+	#
+	# @param len [Fixnum] Avoid setting this unless a specific size is
+	#   necessary. Default is random within range of min .. max
+	# @return [String] A string that matches <tt>[a-z][a-zA-Z0-9_]*</tt>
+	# @yield [String] The identifier before uniqueness checks. This allows
+	#   you to modify the value and still avoid collisions.
+	def generate(len=nil)
+		raise ArgumentError, "len must be positive integer" if len && len < 1
+		raise ExhaustedSpaceError if @value_by_name.length >= @max_permutations
+
+		# pick a random length within the limits
+		len ||= rand(@opts[:min_length] .. (@opts[:max_length]))
+
+		ident = ""
+
+		# XXX: Infinite loop if block returns only values we've already
+		# generated.
+		loop do
+			ident  = Rex::Text.rand_base(1, "", @opts[:first_char_set])
+			ident << Rex::Text.rand_base(len-1, "", @opts[:char_set])
+			if block_given?
+				ident = yield ident
+			end
+			# Try to make another one if it collides with a previously
+			# generated one.
+			break unless @name_by_value.key?(ident)
+		end
+
+		ident
+	end
+
+end

spec/lib/rex/random_identifier_generator_spec.rb

@@ -0,0 +1,123 @@
+require 'spec_helper'
+require 'rex/random_identifier_generator'
+
+describe Rex::RandomIdentifierGenerator do
+	let(:options) do
+		{ :min_length => 10, :max_length => 20 }
+	end
+
+	subject(:rig) { described_class.new(options) }
+
+	it { should respond_to(:generate) }
+	it { should respond_to(:[]) }
+	it { should respond_to(:get) }
+
+	describe "#generate" do
+		it "should respect :min_length" do
+			1000.times do
+				rig.generate.length.should >= options[:min_length]
+			end
+		end
+
+		it "should respect :max_length" do
+			1000.times do
+				rig.generate.length.should <= options[:max_length]
+			end
+		end
+
+		it "should allow mangling in a block" do
+			ident = rig.generate { |identifier| identifier.upcase }
+			ident.should match(/\A[A-Z0-9_]*\Z/)
+
+			ident = subject.generate { |identifier| identifier.downcase }
+			ident.should match(/\A[a-z0-9_]*\Z/)
+
+			ident = subject.generate { |identifier| identifier.gsub("A","B") }
+			ident.should_not include("A")
+		end
+	end
+
+	describe "#get" do
+		let(:options) do
+			{ :min_length=>3, :max_length=>3 }
+		end
+		it "should return the same thing for subsequent calls" do
+			rig.get(:rspec).should == rig.get(:rspec)
+		end
+		it "should not return the same for different names" do
+			# Statistically...
+			count = 1000
+			a = Set.new
+			count.times do |n|
+				a.add rig.get(n)
+			end
+			a.size.should == count
+		end
+
+		context "with an exhausted set" do
+			let(:options) do
+				{ :char_set => "abcd", :min_length=>2, :max_length=>2 }
+			end
+			let(:max_permutations) do
+				# 26 because first char is hardcoded to be lowercase alpha
+				26 * (options[:char_set].length ** options[:min_length])
+			end
+
+			it "doesn't infinite loop" do
+				Timeout.timeout(1) do
+					expect {
+						(max_permutations + 1).times { |i| rig.get(i) }
+					}.to raise_error(Rex::RandomIdentifierGenerator::ExhaustedSpaceError)
+				# don't rescue TimeoutError here because we want that to be a
+				# failure case
+				end
+			end
+
+		end
+
+	end
+
+	describe "#store" do
+			let(:options) do
+				{ :char_set => "abcd", :min_length=>8, :max_length=>20 }
+			end
+
+		it "should allow smaller than minimum length" do
+			value = "a"*(options[:min_length]-1)
+			rig.store(:spec, value)
+			rig.get(:spec).should == value
+		end
+
+		it "should allow bigger than maximum length" do
+			value = "a"*(options[:max_length]+1)
+			rig.store(:spec, value)
+			rig.get(:spec).should == value
+		end
+
+		it "should raise if value is not unique" do
+			value = "a"*(options[:max_length]+1)
+			rig.store(:spec0, value)
+			rig.get(:spec0).should == value
+			expect { rig.store(:spec1, value) }.to raise_error
+		end
+
+		it "should overwrite a previously stored value" do
+			orig_value = "a"*(options[:max_length])
+			rig.store(:spec, orig_value)
+			rig.get(:spec).should == orig_value
+
+			new_value = "b"*(options[:max_length])
+			rig.store(:spec, new_value)
+			rig.get(:spec).should == new_value
+		end
+
+		it "should overwrite a previously generated value" do
+			rig.get(:spec)
+
+			new_value = "a"*(options[:max_length])
+			rig.store(:spec, new_value)
+			rig.get(:spec).should == new_value
+		end
+
+	end
+end