Land #2087, @egypt's random ID generator
Tod Beardsley
commit
bb16683415
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Reference in New Issue