Skip to content

Commit 4020b7e

Browse files
committed
Improve module extraction time
1 parent 77f46f7 commit 4020b7e

File tree

1 file changed

+77
-61
lines changed

1 file changed

+77
-61
lines changed

resources/extract_module_info.rc

Lines changed: 77 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -43,68 +43,84 @@ common_mixins = (
4343

4444
_, total_time = record_time do
4545
STDERR.puts 'Extracting module paths'
46-
module_paths, module_paths_timing = record_time { framework.modules.to_a }
47-
#STDERR.puts "Extracted module paths in #{module_paths_timing}"
48-
49-
modules_metadata = module_paths.map do |module_path, _module_class|
50-
index += 1
51-
# STDERR.puts "#{index} - #{module_path}"
52-
mod, elapsed = record_time { framework.modules.create(module_path) }
53-
# STDERR.puts "#{index} - #{module_path}, time taken: #{elapsed}, #{mod.type}"
54-
55-
options = mod.options.values.sort_by(&:name).map do |option|
56-
{
57-
type: option.type,
58-
name: option.name,
59-
required: option.required,
60-
default: option.default.to_s,
61-
aliases: option.aliases,
62-
advanced: option.advanced,
63-
evasion: option.evasion,
64-
description: option.desc
65-
}
46+
module_sets, module_paths_timing = record_time { framework.modules.send(:module_set_by_type) }
47+
STDERR.puts "Extracted module paths in #{module_paths_timing}"
48+
49+
modules_metadata = module_sets.map do |(module_set_name, modules)|
50+
module_paths = modules.keys
51+
module_paths.each do |module_path|
52+
module_folder = Msf::Modules::Loader::Base::DIRECTORY_BY_TYPE[module_set_name]
53+
full_module_path = "#{module_folder}/#{module_path}"
54+
55+
index += 1
56+
STDERR.puts "#{index} - #{full_module_path}"
57+
mod, module_creation_time = record_time { framework.modules.create(full_module_path) }
58+
59+
if mod.nil?
60+
STDERR.puts "#{index} - #{full_module_path}, failed to load"
61+
next
62+
end
63+
# STDERR.puts "#{index} - #{full_module_path}, module creation time: #{module_creation_time}, #{mod.type}"
64+
65+
data, module_data_extraction_time = record_time do
66+
options = mod.options.values.sort_by(&:name).map do |option|
67+
{
68+
type: option.type,
69+
name: option.name,
70+
required: option.required,
71+
default: option.default.to_s,
72+
aliases: option.aliases,
73+
advanced: option.advanced,
74+
evasion: option.evasion,
75+
description: option.desc
76+
}
77+
end
78+
79+
data = {
80+
name: mod.name.dup.force_encoding("UTF-8"),
81+
fullname: mod.realname.dup.force_encoding("UTF-8"),
82+
aliases: mod.aliases,
83+
rank: mod.rank,
84+
session_types: mod.respond_to?(:session_types) ? mod.session_types : nil,
85+
disclosure_date: mod.disclosure_date.nil? ? nil : mod.disclosure_date.to_s,
86+
type: mod.type,
87+
author: mod.author.map { |x| x.to_s.force_encoding("UTF-8") },
88+
description: mod.description.to_s.strip.force_encoding("UTF-8"),
89+
references: mod.references.map { |ref| { 'type' => ref.ctx_id, 'value' => ref.ctx_val, 'site' => ref.site } },
90+
platform: sort_platform_string(mod.platform_to_s),
91+
arch: mod.arch_to_s,
92+
mixins: mod.class.ancestors.map(&:to_s) - common_mixins - [mod.class.to_s],
93+
autofilter_ports: (
94+
mod.respond_to?(:autofilter_ports) ? mod.autofilter_ports : nil
95+
),
96+
autofilter_services: (
97+
mod.respond_to?(:autofilter_services) ? mod.autofilter_services : nil
98+
),
99+
targets: (
100+
(mod.respond_to?(:targets) && mod.targets) ? mod.targets.map(&:name) : nil
101+
),
102+
path: mod.file_path.sub(/^#{Msf::Config.install_root}/, ''),
103+
ref_name: mod.class.refname,
104+
check: mod.has_check?,
105+
postAuth: mod.post_auth?,
106+
default_credential: mod.default_cred?,
107+
notes: mod.notes,
108+
needs_cleanup: (
109+
mod.respond_to?(:needs_cleanup?) ? mod.needs_cleanup? : nil
110+
),
111+
options: options
112+
}
113+
114+
if module_set_name == 'exploit'
115+
data[:compatible_payloads] = mod.compatible_payloads.map { |payload_name, payload_class| payload_name }
116+
end
117+
118+
data
119+
end
120+
# STDERR.puts "#{index} - #{full_module_path}, module data extraction time: #{module_data_extraction_time}, #{mod.type}"
121+
122+
data
66123
end
67-
68-
data = {
69-
name: mod.name.dup.force_encoding("UTF-8"),
70-
fullname: mod.realname.dup.force_encoding("UTF-8"),
71-
aliases: mod.aliases,
72-
rank: mod.rank,
73-
session_types: mod.respond_to?(:session_types) ? mod.session_types : nil,
74-
disclosure_date: mod.disclosure_date.nil? ? nil : mod.disclosure_date.to_s,
75-
type: mod.type,
76-
author: mod.author.map { |x| x.to_s.force_encoding("UTF-8") },
77-
description: mod.description.to_s.strip.force_encoding("UTF-8"),
78-
references: mod.references.map { |ref| { 'type' => ref.ctx_id, 'value' => ref.ctx_val, 'site' => ref.site } },
79-
platform: sort_platform_string(mod.platform_to_s),
80-
arch: mod.arch_to_s,
81-
mixins: mod.class.ancestors.map(&:to_s) - common_mixins - [mod.class.to_s],
82-
autofilter_ports: (
83-
mod.respond_to?(:autofilter_ports) ? mod.autofilter_ports : nil
84-
),
85-
autofilter_services: (
86-
mod.respond_to?(:autofilter_services) ? mod.autofilter_services : nil
87-
),
88-
targets: (
89-
(mod.respond_to?(:targets) && mod.targets) ? mod.targets.map(&:name) : nil
90-
),
91-
path: mod.file_path.sub(/^#{Msf::Config.install_root}/, ''),
92-
ref_name: mod.class.refname,
93-
check: mod.has_check?,
94-
postAuth: mod.post_auth?,
95-
default_credential: mod.default_cred?,
96-
notes: mod.notes,
97-
needs_cleanup: (
98-
mod.respond_to?(:needs_cleanup?) ? mod.needs_cleanup? : nil
99-
),
100-
options: options
101-
}
102-
103-
if mod.type == 'exploit'
104-
data[:compatible_payloads] = mod.compatible_payloads.map { |payload_name, payload_class| payload_name }
105-
end
106-
107-
data
108124
end
109125

110126
File.open(output_path, mode: 'w') do |output_file|

0 commit comments

Comments
 (0)