lib/ronin/chars/char_set.rb in ronin-0.1.1 vs lib/ronin/chars/char_set.rb in ronin-0.1.2

@@ -19,96 +19,179 @@
 # along with this program; if not, write to the Free Software
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 #++
 #
 
+require 'set'
+
 module Ronin
   module Chars
-    class CharSet < Array
+    class CharSet < SortedSet
 
       #
       # Creates a new CharSet object with the given _chars_.
       #
       def initialize(*chars)
-        format_char = lambda { |char|
-          if (char.kind_of?(Array) || char.kind_of?(Range))
-            char.map(&format_char)
-          elsif char.kind_of?(Integer)
-            char.chr
-          else
-            char.to_s
+        super()
+
+        merge_chars = lambda { |element|
+          if element.kind_of?(String)
+            element.each_byte(&merge_chars)
+          elsif element.kind_of?(Integer)
+            self << element
+          elsif element.kind_of?(Enumerable)
+            element.each(&merge_chars)
           end
         }
 
-        super(chars.map(&format_char).flatten.uniq)
+        merge_chars.call(chars)
       end
 
+      alias include_byte? include?
+      alias bytes to_a
+      alias each_byte each
+      alias select_bytes select
+      alias map_bytes map
+
       #
-      # Create a new CharSet object containing the characters that match
+      # Returns +true+ if the character set includes the specified _char_,
+      # returns +false+ otherwise.
+      #
+      def include_char?(char)
+        char.each_byte do |b|
+          return include?(b)
+        end
+      end
+
+      #
+      # Returns all the characters within the character set as Strings.
+      #
+      def chars
+        map { |b| b.chr }
+      end
+
+      #
+      # Iterates over every character within the character set, passing
+      # each to the given _block_.
+      #
+      def each_char(&block)
+        each { |b| block.call(b.chr) } if block
+      end
+
+      #
+      # Selects an Array of characters from the character set that match
       # the given _block_.
       #
-      def select(&block)
-        CharSet.new(super(&block))
+      def select_chars(&block)
+        chars.select(&block)
       end
 
       #
-      # Creates a new CharSet object by passing each character to the
-      # specified _block_.
+      # Maps the characters of the character set using the given _block_.
       #
-      def map(&block)
-        CharSet.new(super(&block))
+      def map_chars(&block)
+        chars.map(&block)
       end
 
       #
+      # Returns a random byte from the character set.
+      #
+      def random_byte
+        self.entries[rand(self.length)]
+      end
+
+      #
       # Returns a random char from the character set.
       #
       def random_char
-        self[rand(self.length)]
+        random_byte.chr
       end
 
       #
+      # Pass a random byte to the specified _block_, _n_ times.
+      #
+      def each_random_byte(n,&block)
+        n.times { block.call(random_byte) }
+      end
+
+      #
       # Pass a random character to the specified _block_, _n_ times.
       #
       def each_random_char(n,&block)
-        n.to_i.times { block.call(random_char) }
+        each_random_byte(n) { |b| block.call(b.chr) }
       end
 
       #
       # Returns an Array of the specified _length_ containing
-      # random characters from the character set.
+      # random bytes from the character set. The specified _length_ may
+      # be an Integer, Array or a Range of lengths.
       #
-      def random_array(length)
-        if length.kind_of?(Range)
-          return Array.new(length.sort_by { rand }.first) { random_char }
+      def random_bytes(length)
+        if (length.kind_of?(Array) || length.kind_of?(Range))
+          return Array.new(length.sort_by { rand }.first) { random_byte }
         else
-          return Array.new(length.to_i) { random_char }
+          return Array.new(length) { random_byte }
         end
       end
 
       #
+      # Returns an Array of the specified _length_ containing
+      # random characters from the character set. The specified _length_
+      # may be an Integer, Array or a Range of lengths.
+      #
+      def random_chars(length)
+        random_bytes(length).map { |b| b.chr }
+      end
+
+      #
       # Returns a String of the specified _length_ containing
       # random characters from the character set.
       #
       def random_string(length)
-        random_array(length).join
+        random_chars(length).join
       end
 
       #
-      # Return a new CharSet that is the union of the character set and the
-      # specified _other_set_.
+      # Creates a new CharSet object containing the both the characters
+      # of the character set and the specified _other_set_.
       #
       def |(other_set)
-        CharSet.new(self,other_set)
+        super(CharSet.new(other_set))
       end
 
       alias + |
 
       #
-      # Returns a new CharSet that is the intersection of the character set
-      # and the specified _other_set_.
+      # Returns +true+ if all of the bytes within the specified _string_
+      # are included in the character set, returns +false+ otherwise.
       #
-      def -(other_set)
-        CharSet.new(super(other_set.to_a))
+      #   Chars.alpha =~ "hello"
+      #   # => true
+      #
+      def =~(string)
+        string.each_byte do |b|
+          return false unless include?(b)
+        end
+
+        return true
+      end
+
+      #
+      # Inspects the character set.
+      #
+      def inspect
+        "#<#{self.class.name}: {" + map { |b|
+          case b
+          when (0x07..0x0d), (0x20..0x7e)
+            b.chr.dump
+          when 0x00
+            # sly hack to make char-sets more friendly
+            # to us C programmers
+            '"\0"'
+          else
+            "0x%02x" % b
+          end
+        }.join(', ') + "}>"
       end
 
     end
   end
 end