Browse Source

Revert "Land #6729, Speed up the datastore"

This reverts commit c6b1955a5a, reversing
changes made to 4fb7472391.
Louis Sato 3 years ago
parent
commit
d5d0b9e9b8

+ 65
- 46
lib/msf/core/data_store.rb View File

@@ -1,7 +1,4 @@
1 1
 # -*- coding: binary -*-
2
-
3
-require 'set'
4
-
5 2
 module Msf
6 3
 
7 4
 ###
@@ -16,10 +13,9 @@ class DataStore < Hash
16 13
   # Initializes the data store's internal state.
17 14
   #
18 15
   def initialize()
19
-    @options       = Hash.new
20
-    @imported      = Hash.new
21
-    @imported_by   = Hash.new
22
-    @original_keys = Set.new
16
+    @options     = Hash.new
17
+    @imported    = Hash.new
18
+    @imported_by = Hash.new
23 19
   end
24 20
 
25 21
   #
@@ -27,8 +23,7 @@ class DataStore < Hash
27 23
   # directly.
28 24
   #
29 25
   def []=(k, v)
30
-    add_key(k)
31
-    k = k.downcase
26
+    k = find_key_case(k)
32 27
     @imported[k] = false
33 28
     @imported_by[k] = nil
34 29
 
@@ -49,32 +44,31 @@ class DataStore < Hash
49 44
   # Case-insensitive wrapper around hash lookup
50 45
   #
51 46
   def [](k)
52
-    super(k.downcase)
47
+    super(find_key_case(k))
53 48
   end
54 49
 
55 50
   #
56 51
   # Case-insensitive wrapper around store
57 52
   #
58 53
   def store(k,v)
59
-    add_key(k)
60
-    super(k.downcase, v)
54
+    super(find_key_case(k), v)
61 55
   end
62 56
 
63 57
   #
64 58
   # Case-insensitive wrapper around delete
65 59
   #
66 60
   def delete(k)
67
-    super(k.downcase)
61
+    super(find_key_case(k))
68 62
   end
69 63
 
70
-  # Override Hash's to_h method so we can include the original case of each key
71
-  # (failing to do this breaks a number of places in framework and pro that use
72
-  # serialized datastores)
73
-  def to_h
74
-    @original_keys.reduce({}) do |acc, key|
75
-      acc[key] = self[key]
76
-      acc
77
-    end
64
+
65
+  #
66
+  # Updates a value in the datastore with the specified name, k, to the
67
+  # specified value, v.  This update does not alter the imported status of
68
+  # the value.
69
+  #
70
+  def update_value(k, v)
71
+    self.store(k, v)
78 72
   end
79 73
 
80 74
   #
@@ -134,16 +128,15 @@ class DataStore < Hash
134 128
   # Imports options from a hash and stores them in the datastore.
135 129
   #
136 130
   def import_options_from_hash(option_hash, imported = true, imported_by = nil)
137
-    option_hash.each_pair do |key, val|
131
+    option_hash.each_pair { |key, val|
138 132
       import_option(key, val, imported, imported_by)
139
-    end
133
+    }
140 134
   end
141 135
 
142 136
   def import_option(key, val, imported=true, imported_by=nil, option=nil)
143 137
     self.store(key, val)
144 138
 
145
-    key = key.downcase
146
-    @options[key]     = option
139
+    @options[key] = option
147 140
     @imported[key]    = imported
148 141
     @imported_by[key] = imported_by
149 142
   end
@@ -152,9 +145,21 @@ class DataStore < Hash
152 145
   # Serializes the options in the datastore to a string.
153 146
   #
154 147
   def to_s(delim = ' ')
155
-    @original_keys.reduce('') do |acc, key|
156
-      acc << "#{key}=#{self[key]}#{delim}"
148
+    str = ''
149
+
150
+    keys.sort.each { |key|
151
+      str << "#{key}=#{self[key]}" + ((str.length) ? delim : '')
152
+    }
153
+
154
+    return str
155
+  end
156
+
157
+  def to_h
158
+    datastore_hash = {}
159
+    self.keys.each do |k|
160
+      datastore_hash[k.to_s] = self[k].to_s
157 161
     end
162
+    datastore_hash
158 163
   end
159 164
 
160 165
   #
@@ -194,10 +199,9 @@ class DataStore < Hash
194 199
   # not include default option values.
195 200
   #
196 201
   def user_defined
197
-    @original_keys.reduce({}) do |acc, k|
198
-      acc[k] = self[k] unless @imported[k.downcase]
199
-      acc
200
-    end
202
+    reject { |k, v|
203
+      @imported[k] == true
204
+    }
201 205
   end
202 206
 
203 207
   #
@@ -218,26 +222,40 @@ class DataStore < Hash
218 222
   # Completely clear all values in the hash
219 223
   #
220 224
   def clear
221
-    @options.clear
222
-    @imported.clear
223
-    @imported_by.clear
224
-    @original_keys.clear
225
-    super
225
+    self.keys.each {|k| self.delete(k) }
226
+    self
226 227
   end
227 228
 
228
-  # Yield the original-cased key
229
+  #
230
+  # Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
231
+  #    "can't add a new key into hash during iteration"
232
+  #
229 233
   def each(&block)
230
-    @original_keys.each do |key|
231
-      block.call(key, self[key])
234
+    list = []
235
+    self.keys.sort.each do |sidx|
236
+      list << [sidx, self[sidx]]
232 237
     end
238
+    list.each(&block)
233 239
   end
234 240
 
235
-  protected
241
+protected
242
+
243
+  #
244
+  # Case-insensitive key lookup
245
+  #
246
+  def find_key_case(k)
236 247
 
237
-  # Keep track of the original, case-sensitive key
238
-  def add_key(k)
239
-    @original_keys.add(k) unless include? k.downcase
248
+    # Scan each key looking for a match
249
+    self.each_key do |rk|
250
+      if (rk.downcase == k.downcase)
251
+        return rk
252
+      end
253
+    end
254
+
255
+    # Fall through to the non-existent value
256
+    return k
240 257
   end
258
+
241 259
 end
242 260
 
243 261
 ###
@@ -260,7 +278,7 @@ class ModuleDataStore < DataStore
260 278
   # if we can't directly find it
261 279
   #
262 280
   def fetch(key)
263
-    key = key.downcase
281
+    key = find_key_case(key)
264 282
     val = nil
265 283
     val = super if(@imported_by[key] != 'self')
266 284
     if (val.nil? and @_module and @_module.framework)
@@ -274,7 +292,7 @@ class ModuleDataStore < DataStore
274 292
   # Same as fetch
275 293
   #
276 294
   def [](key)
277
-    key = key.downcase
295
+    key = find_key_case(key)
278 296
     val = nil
279 297
     val = super if(@imported_by[key] != 'self')
280 298
     if (val.nil? and @_module and @_module.framework)
@@ -297,10 +315,11 @@ class ModuleDataStore < DataStore
297 315
   def copy
298 316
     clone = self.class.new(@_module)
299 317
     self.keys.each do |k|
300
-      clone.import_option(k, self[k].kind_of?(String) ? self[k].dup : self[k], @imported[k.downcase], @imported_by[k.downcase])
318
+      clone.import_option(k, self[k].kind_of?(String) ? self[k].dup : self[k], @imported[k], @imported_by[k])
301 319
     end
302 320
     clone
303 321
   end
304 322
 end
305 323
 
306 324
 end
325
+

+ 2
- 2
lib/msf/core/exploit/powershell.rb View File

@@ -58,7 +58,7 @@ module Exploit::Powershell
58 58
   # @return [String] Encoded script
59 59
   def encode_script(script_in, eof = nil)
60 60
     opts = {}
61
-    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/i && v }.keys.map do |k|
61
+    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |k|
62 62
       mod_method = k.split('::').last.intern
63 63
       opts[mod_method.to_sym] = true
64 64
     end
@@ -76,7 +76,7 @@ module Exploit::Powershell
76 76
   # @return [String] Compressed script with decompression stub
77 77
   def compress_script(script_in, eof=nil)
78 78
     opts = {}
79
-    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/i && v }.keys.map do |k|
79
+    datastore.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |k|
80 80
       mod_method = k.split('::').last.intern
81 81
       opts[mod_method.to_sym] = true
82 82
     end

+ 1
- 1
modules/post/multi/recon/local_exploit_suggester.rb View File

@@ -88,7 +88,7 @@ class MetasploitModule < Msf::Post
88 88
 
89 89
 
90 90
   def set_module_options(mod)
91
-    self.datastore.each do |k,v|
91
+    self.datastore.each_pair do |k,v|
92 92
       mod.datastore[k] = v
93 93
     end
94 94
     if !mod.datastore['SESSION'] && session.present?

Loading…
Cancel
Save