From 273046d8f0ba07eb46268b011081e8676be24a07 Mon Sep 17 00:00:00 2001 From: James Lee Date: Tue, 9 Jul 2013 02:06:44 -0500 Subject: [PATCH 1/3] Add a class for generating random identifiers Will be useful for all kinds of things, but brought about in discussions specifically for Util::EXE in #2037. --- lib/rex/random_identifier_generator.rb | 155 ++++++++++++++++++ .../rex/random_identifier_generator_spec.rb | 123 ++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 lib/rex/random_identifier_generator.rb create mode 100644 spec/lib/rex/random_identifier_generator_spec.rb diff --git a/lib/rex/random_identifier_generator.rb b/lib/rex/random_identifier_generator.rb new file mode 100644 index 0000000000..1a2c1df23b --- /dev/null +++ b/lib/rex/random_identifier_generator.rb @@ -0,0 +1,155 @@ + +# 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+"_", + } + + # @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. + @identifiers = {} + + @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. Since + # this will still be a pretty big number most of the time, 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 @identifiers[name] if @identifiers[name] + + @identifiers[name] = generate + + @identifiers[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 @identifiers.key(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 + @identifiers[name] = value + 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. + # + # Note that the first character will always be lowercase alpha. + # + # @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 [a-z][a-zA-Z0-9_]* + # @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 @identifiers.length >= @max_permutations + + # pick a random length within the limits + len ||= rand(@opts[:min_length] .. (@opts[:max_length])) + + ident = "" + + # XXX: infinite loop if we've exhausted the space. Mitigated by the + # fact that you'd have to call generate at least 26*62 times (in the + # case of 2-character names) to hit it with the default :char_set. + loop do + ident = Rex::Text.rand_text_alpha_lower(1) + 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 @identifiers.value?(ident) + end + + ident + end + +end diff --git a/spec/lib/rex/random_identifier_generator_spec.rb b/spec/lib/rex/random_identifier_generator_spec.rb new file mode 100644 index 0000000000..1660f4898c --- /dev/null +++ b/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 From afa6a36df3c0608e8d891d2840234a038adb9983 Mon Sep 17 00:00:00 2001 From: James Lee Date: Tue, 9 Jul 2013 02:49:19 -0500 Subject: [PATCH 2/3] Make first char's character class configurable --- lib/rex/random_identifier_generator.rb | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/lib/rex/random_identifier_generator.rb b/lib/rex/random_identifier_generator.rb index 1a2c1df23b..77c1e95f37 100644 --- a/lib/rex/random_identifier_generator.rb +++ b/lib/rex/random_identifier_generator.rb @@ -31,6 +31,7 @@ class Rex::RandomIdentifierGenerator :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 @@ -46,8 +47,8 @@ class Rex::RandomIdentifierGenerator raise ArgumentError, "Invalid length options" end - # This is really just the maximum number of shortest names. Since - # this will still be a pretty big number most of the time, don't + # 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 @@ -108,9 +109,17 @@ class Rex::RandomIdentifierGenerator end # Create a random string that satisfies most languages' requirements - # for identifiers. + # 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. # - # Note that the first character will always be lowercase alpha. + # 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 @@ -139,7 +148,7 @@ class Rex::RandomIdentifierGenerator # fact that you'd have to call generate at least 26*62 times (in the # case of 2-character names) to hit it with the default :char_set. loop do - ident = Rex::Text.rand_text_alpha_lower(1) + 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 From 4cc179a24cfb3d662bca5f174ad7181fc0060bfd Mon Sep 17 00:00:00 2001 From: James Lee Date: Wed, 10 Jul 2013 12:38:42 -0500 Subject: [PATCH 3/3] Store inverted hash for better lookups Also clarifies comment about infinite loops --- lib/rex/random_identifier_generator.rb | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/lib/rex/random_identifier_generator.rb b/lib/rex/random_identifier_generator.rb index 77c1e95f37..b85ae55d50 100644 --- a/lib/rex/random_identifier_generator.rb +++ b/lib/rex/random_identifier_generator.rb @@ -40,7 +40,10 @@ class Rex::RandomIdentifierGenerator # @option opts :char_set [String] def initialize(opts={}) # Holds all identifiers. - @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] @@ -71,11 +74,12 @@ class Rex::RandomIdentifierGenerator # you weren't generating it. # @return [String] def get(name) - return @identifiers[name] if @identifiers[name] + return @value_by_name[name] if @value_by_name[name] - @identifiers[name] = generate + @value_by_name[name] = generate + @name_by_value[@value_by_name[name]] = name - @identifiers[name] + @value_by_name[name] end alias [] get @@ -93,13 +97,14 @@ class Rex::RandomIdentifierGenerator # @return [void] def store(name, value) - case @identifiers.key(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 - @identifiers[name] = value + @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!" @@ -137,16 +142,15 @@ class Rex::RandomIdentifierGenerator # 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 @identifiers.length >= @max_permutations + 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 we've exhausted the space. Mitigated by the - # fact that you'd have to call generate at least 26*62 times (in the - # case of 2-character names) to hit it with the default :char_set. + # 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]) @@ -155,7 +159,7 @@ class Rex::RandomIdentifierGenerator end # Try to make another one if it collides with a previously # generated one. - break unless @identifiers.value?(ident) + break unless @name_by_value.key?(ident) end ident