Bring some sanity to the datastore

Before, the datastore would store options case-sensitive, but would
access them case-insensitive, resulting is a number of string compares.
This commit stores options in their downcase form to reduce
update/lookup time. This adds up to reducing msfconsole boot time by
about 10% and rspec time by about 45 sec. (!) on my box.

One tricky part of this conversion is that there are several places (in
pro and framework) where we export or otherwise access the datastore as
a plain hash (case-sensitive). I believe I have caught all the ways we
access the datastore that are case-sensitive and substituted the
original key capitalization in those cases.

Author: acammack-r7

lib/msf/core/data_store.rb

@@ -1,4 +1,7 @@
 # -*- coding: binary -*-
+
+require 'set'
+
 module Msf
 
 ###
@@ -13,17 +16,19 @@ class DataStore < Hash
   # Initializes the data store's internal state.
   #
   def initialize()
-    @options     = Hash.new
-    @imported    = Hash.new
-    @imported_by = Hash.new
+    @options       = Hash.new
+    @imported      = Hash.new
+    @imported_by   = Hash.new
+    @original_keys = Set.new
   end
 
   #
   # Clears the imported flag for the supplied key since it's being set
   # directly.
   #
   def []=(k, v)
-    k = find_key_case(k)
+    add_key(k)
+    k = k.downcase
     @imported[k] = false
     @imported_by[k] = nil
 
@@ -44,31 +49,32 @@ def []=(k, v)
   # Case-insensitive wrapper around hash lookup
   #
   def [](k)
-    super(find_key_case(k))
+    super(k.downcase)
   end
 
   #
   # Case-insensitive wrapper around store
   #
   def store(k,v)
-    super(find_key_case(k), v)
+    add_key(k)
+    super(k.downcase, v)
   end
 
   #
   # Case-insensitive wrapper around delete
   #
   def delete(k)
-    super(find_key_case(k))
+    super(k.downcase)
   end
 
-
-  #
-  # Updates a value in the datastore with the specified name, k, to the
-  # specified value, v.  This update does not alter the imported status of
-  # the value.
-  #
-  def update_value(k, v)
-    self.store(k, v)
+  # Override Hash's to_h method so we can include the original case of each key
+  # (failing to do this breaks a number of places in framework and pro that use
+  # serialized datastores)
+  def to_h
+    @original_keys.reduce({}) do |acc, key|
+      acc[key] = self[key]
+      acc
+    end
   end
 
   #
@@ -128,15 +134,16 @@ def import_options_from_s(option_str, delim = nil)
   # Imports options from a hash and stores them in the datastore.
   #
   def import_options_from_hash(option_hash, imported = true, imported_by = nil)
-    option_hash.each_pair { |key, val|
+    option_hash.each_pair do |key, val|
       import_option(key, val, imported, imported_by)
-    }
+    end
   end
 
   def import_option(key, val, imported=true, imported_by=nil, option=nil)
     self.store(key, val)
 
-    @options[key] = option
+    key = key.downcase
+    @options[key]     = option
     @imported[key]    = imported
     @imported_by[key] = imported_by
   end
@@ -145,21 +152,9 @@ def import_option(key, val, imported=true, imported_by=nil, option=nil)
   # Serializes the options in the datastore to a string.
   #
   def to_s(delim = ' ')
-    str = ''
-
-    keys.sort.each { |key|
-      str << "#{key}=#{self[key]}" + ((str.length) ? delim : '')
-    }
-
-    return str
-  end
-
-  def to_h
-    datastore_hash = {}
-    self.keys.each do |k|
-      datastore_hash[k.to_s] = self[k].to_s
+    @original_keys.reduce('') do |acc, key|
+      acc << "#{key}=#{self[key]}#{delim}"
     end
-    datastore_hash
   end
 
   #
@@ -199,9 +194,10 @@ def from_file(path, name = 'global')
   # not include default option values.
   #
   def user_defined
-    reject { |k, v|
-      @imported[k] == true
-    }
+    @original_keys.reduce({}) do |acc, k|
+      acc[k] = self[k] unless @imported[k.downcase]
+      acc
+    end
   end
 
   #
@@ -222,40 +218,26 @@ def clear_non_user_defined
   # Completely clear all values in the hash
   #
   def clear
-    self.keys.each {|k| self.delete(k) }
-    self
+    @options.clear
+    @imported.clear
+    @imported_by.clear
+    @original_keys.clear
+    super
   end
 
-  #
-  # Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
-  #    "can't add a new key into hash during iteration"
-  #
+  # Yield the original-cased key
   def each(&block)
-    list = []
-    self.keys.sort.each do |sidx|
-      list << [sidx, self[sidx]]
+    @original_keys.each do |key|
+      block.call(key, self[key])
     end
-    list.each(&block)
   end
 
-protected
-
-  #
-  # Case-insensitive key lookup
-  #
-  def find_key_case(k)
+  protected
 
-    # Scan each key looking for a match
-    self.each_key do |rk|
-      if (rk.downcase == k.downcase)
-        return rk
-      end
-    end
-
-    # Fall through to the non-existent value
-    return k
+  # Keep track of the original, case-sensitive key
+  def add_key(k)
+    @original_keys.add(k) unless include? k.downcase
   end
-
 end
 
 ###
@@ -278,7 +260,7 @@ def initialize(m)
   # if we can't directly find it
   #
   def fetch(key)
-    key = find_key_case(key)
+    key = key.downcase
     val = nil
     val = super if(@imported_by[key] != 'self')
     if (val.nil? and @_module and @_module.framework)
@@ -292,7 +274,7 @@ def fetch(key)
   # Same as fetch
   #
   def [](key)
-    key = find_key_case(key)
+    key = key.downcase
     val = nil
     val = super if(@imported_by[key] != 'self')
     if (val.nil? and @_module and @_module.framework)
@@ -315,11 +297,10 @@ def default?(key)
   def copy
     clone = self.class.new(@_module)
     self.keys.each do |k|
-      clone.import_option(k, self[k].kind_of?(String) ? self[k].dup : self[k], @imported[k], @imported_by[k])
+      clone.import_option(k, self[k].kind_of?(String) ? self[k].dup : self[k], @imported[k.downcase], @imported_by[k.downcase])
     end
     clone
   end
 end
 
 end
-

lib/msf/core/exploit/powershell.rb

@@ -58,7 +58,7 @@ def make_subs(script, subs)
   # @return [String] Encoded script
   def encode_script(script_in, eof = nil)
     opts = {}
-    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |k|
+    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/i && v }.keys.map do |k|
       mod_method = k.split('::').last.intern
       opts[mod_method.to_sym] = true
     end
@@ -76,7 +76,7 @@ def encode_script(script_in, eof = nil)
   # @return [String] Compressed script with decompression stub
   def compress_script(script_in, eof=nil)
     opts = {}
-    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |k|
+    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/i && v }.keys.map do |k|
       mod_method = k.split('::').last.intern
       opts[mod_method.to_sym] = true
     end