Skip to content

Commit

Permalink
Split recalculate out into multiple functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@dwelch-r7
dwelch-r7 committed May 12, 2023
1 parent 20ac531 commit 8e3b8af
Showing 1 changed file with 195 additions and 111 deletions.
306 changes: 195 additions & 111 deletions lib/msf/core/payload_set.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,140 +70,225 @@ def recalculate
new_keys = []

# Recalculate single payloads
_singles.each_pair { |name, op|
mod, handler = op

# if the payload has a dependency, check
# if it is supported on the system
payload_dependencies = op[4].dependencies
unless payload_dependencies.empty?
supported = payload_dependencies.all?(&:available?)
elog("Dependency for #{name} is not supported") unless supported
next unless supported
_singles.each_pair do |single_name, single_info|
single_payload = calculate_single(name: single_name, single_info: single_info)
next unless single_payload

new_keys.push single_name

_adapters.each_pair do |adapter_name, adapter_info|
adapted_single = calculate_adapted_single(adapter_name: adapter_name,
adapter_info: adapter_info,
single_info: single_info,
single_payload: single_payload)
next unless adapted_single

new_keys.push adapted_single.refname
end
end

# Build the payload dupe using the determined handler
# and module
p = build_payload(handler, mod)
# Recalculate staged payloads
_stagers.each_pair do |stager_name, stager_info|
_stager_mod, _handler, _stager_platform, _stager_arch, stager_inst, _stager_modinfo = stager_info

# Add it to the set
add_single(p, name, op[5])
new_keys.push name
next unless stager_dependencies_available?(stager_name: stager_name, stager_dependencies: stager_inst.dependencies)

_adapters.each_pair { |adapter_name, ep|
adapter_mod, _, adapter_platform, adapter_arch, adapter_inst = ep
next unless adapter_inst.compatible?(p.new)
# Walk the array of stages
_stages.each_pair do |stage_name, stage_info|

ap = build_payload(handler, mod, adapter_mod)
staged_payload = calculate_staged(stage_name: stage_name,
stager_name: stager_name,
stage_info: stage_info,
stager_info: stager_info)
next unless staged_payload

combined = build_adapted_name(adapter_name, p.refname)
add_single(ap, combined, ep[5])
new_keys.push combined
}
}
new_keys.push staged_payload.refname

# Recalculate staged payloads
_stagers.each_pair { |stager_name, op|
stager_mod, handler, stager_platform, stager_arch, stager_inst = op

# Pass if the stager has a dependency
# and doesn't have the dependency installed
stager_dependencies = stager_inst.dependencies
unless stager_dependencies.empty?
supported = stager_dependencies.all?(&:available?)
elog("Dependency for #{stager_name} is not supported") unless supported
next unless supported
end
_adapters.each_pair do |adapter_name, adapter_info|
adapted_staged_payload = calculate_adapted_staged(staged_payload: staged_payload,
adapter_name: adapter_name,
stage_info: stage_info,
stager_info: stager_info,
adapter_info: adapter_info)
next unless adapted_staged_payload

# Walk the array of stages
_stages.each_pair { |stage_name, ip|
stage_mod, _, stage_platform, stage_arch, stage_inst = ip

#
# if the stager or stage has a dependency, check
# if they are compatible
#
unless stager_dependencies.empty? && stage_inst.dependencies.empty?
next unless stager_dependencies == stage_inst.dependencies
new_keys.push adapted_staged_payload.refname
end
end
end

# No intersection between platforms on the payloads?
if ((stager_platform) and
(stage_platform) and
(stager_platform & stage_platform).empty?)
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible platforms: #{stager_platform.names} - #{stage_platform.names}", 'core', LEV_2)
next
end
# Blow away anything that was cached but didn't exist during the
# recalculation
self.delete_if do |k, v|
next if v == SymbolicModule
!!(old_keys.include?(k) and not new_keys.include?(k))
end

# No intersection between architectures on the payloads?
if ((stager_arch) and
(stage_arch) and
((stager_arch & stage_arch).empty?))
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible architectures: #{stager_arch.join} - #{stage_arch.join}", 'core', LEV_2)
next
end
flush_blob_cache
end

# If the stage has a convention, make sure it's compatible with
# the stager's
if ((stage_inst) and (stage_inst.compatible?(stager_inst) == false))
dlog("Stager #{stager_name} and stage #{stage_name} are incompatible.", 'core', LEV_2)
next
end
def calculate_single(name:, single_info:)
mod, handler, _single_platform, _single_arch, single_inst, single_modinfo = single_info

# Build the payload dupe using the handler, stager,
# and stage
p = build_payload(handler, stager_mod, stage_mod)

# If the stager has an alias for the handler type (such as is the
# case for ordinal based stagers), use it in preference of the
# handler's actual type.
if (stager_mod.respond_to?('handler_type_alias') == true)
handler_type = stager_mod.handler_type_alias
else
handler_type = handler.handler_type
end
# if the payload has a dependency, check
# if it is supported on the system
payload_dependencies = single_inst.dependencies
unless payload_dependencies.empty?
supported = payload_dependencies.all?(&:available?)
elog("Dependency for #{name} is not supported") unless supported
return nil unless supported
end

# Associate the name as a combination of the stager and stage
combined = stage_name
# Build the payload dupe using the determined handler
# and module
payload = build_payload(handler, mod)

# If a valid handler exists for this stager, then combine it
combined += '/' + handler_type
# Add it to the set
add_single(payload, name, single_modinfo)

# Sets the modules derived name
p.refname = combined
payload
end

# Add the stage
add_stage(p, combined, stage_name, handler_type, {
'files' => op[5]['files'] + ip[5]['files'],
'paths' => op[5]['paths'] + ip[5]['paths'],
'type' => op[5]['type']})
new_keys.push combined
def calculate_adapted_single(adapter_name:, adapter_info:, single_info:, single_payload:)
adapter_mod, _, _adapter_platform, _adapter_arch, adapter_inst, adapted_modinfo = adapter_info
single_mod, handler, _single_platform, _single_arch, _single_inst, _single_modinfo = single_info

_adapters.each_pair { |adapter_name, ep|
adapter_mod, _, adapter_platform, adapter_arch, adapter_inst = ep
next unless adapter_inst.compatible?(p.new)
return nil unless adapter_inst.compatible?(single_payload.new)

ap = build_payload(handler, stager_mod, stage_mod, adapter_mod)
payload = build_payload(handler, single_mod, adapter_mod)

combined = build_adapted_name(adapter_name, p.refname)
ap.refname = combined
ap.framework = framework
ap.file_path = ep[5]['files'][0]
adapted_name = build_adapted_name(adapter_name, single_payload.refname)
add_single(payload, adapted_name, adapted_modinfo)

self[combined] = ap
new_keys.push combined
}
}
}
payload
end

# Blow away anything that was cached but didn't exist during the
# recalculation
self.delete_if do |k, v|
next if v == SymbolicModule
!!(old_keys.include?(k) and not new_keys.include?(k))
def stager_dependencies_available?(stager_name:, stager_dependencies:)
# Pass if the stager has a dependency
# and doesn't have the dependency installed
supported = true # Default to true for stagers with no dependencies
unless stager_dependencies.empty?
supported = stager_dependencies.all?(&:available?)
elog("Dependency for #{stager_name} is not supported") unless supported
end
supported
end

flush_blob_cache
def stage_and_stager_compatible?(stager_info:, stage_info:, stager_name:, stage_name:)
_stager_mod, _handler, stager_platform, stager_arch, stager_inst = stager_info
_stage_mod, _, stage_platform, stage_arch, stage_inst = stage_info

stager_dependencies = stager_inst.dependencies
stage_dependencies = stage_inst.dependencies

unless stager_dependencies.empty? && stage_dependencies.empty?
return false unless stager_dependencies == stage_dependencies
end

# No intersection between platforms on the payloads?
if ((stager_platform) and
(stage_platform) and
(stager_platform & stage_platform).empty?)
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible platforms: #{stager_platform.names} - #{stage_platform.names}", 'core', LEV_2)
return false
end

# No intersection between architectures on the payloads?
if ((stager_arch) and
(stage_arch) and
((stager_arch & stage_arch).empty?))
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible architectures: #{stager_arch.join} - #{stage_arch.join}", 'core', LEV_2)
return false
end

# If the stage has a convention, make sure it's compatible with
# the stager's
if ((stage_inst) and (stage_inst.compatible?(stager_inst) == false))
dlog("Stager #{stager_name} and stage #{stage_name} are incompatible.", 'core', LEV_2)
return false
end

# No intersection between platforms on the payloads?
if ((stager_platform) and
(stage_platform) and
(stager_platform & stage_platform).empty?)
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible platforms: #{stager_platform.names} - #{stage_platform.names}", 'core', LEV_2)
return false
end

# No intersection between architectures on the payloads?
if ((stager_arch) and
(stage_arch) and
((stager_arch & stage_arch).empty?))
dlog("Stager #{stager_name} and stage #{stage_name} have incompatible architectures: #{stager_arch.join} - #{stage_arch.join}", 'core', LEV_2)
return false
end

# If the stage has a convention, make sure it's compatible with
# the stager's
if ((stage_inst) and (stage_inst.compatible?(stager_inst) == false))
dlog("Stager #{stager_name} and stage #{stage_name} are incompatible.", 'core', LEV_2)
return false
end
true
end

def calculate_staged(stage_name:, stager_name:, stage_info:, stager_info:)
# if the stager or stage has a dependency, check
# if they are compatible
return nil unless stage_and_stager_compatible?(stager_info: stager_info,
stage_info: stage_info,
stager_name: stager_name,
stage_name: stage_name)

stager_mod, handler, _stager_platform, _stager_arch, _stager_inst, stager_modinfo = stager_info
stage_mod, _, _stage_platform, _stage_arch, _stage_inst, stage_modinfo = stage_info
# Build the payload dupe using the handler, stager,
# and stage
payload = build_payload(handler, stager_mod, stage_mod)

# If the stager has an alias for the handler type (such as is the
# case for ordinal based stagers), use it in preference of the
# handler's actual type.
if (stager_mod.respond_to?('handler_type_alias') == true)
handler_type = stager_mod.handler_type_alias
else
handler_type = handler.handler_type
end

# Associate the name as a combination of the stager and stage
staged_refname = stage_name

# If a valid handler exists for this stager, then combine it
staged_refname += '/' + handler_type

# Sets the modules derived name
payload.refname = staged_refname

# Add the stage
add_stage(payload, staged_refname, stage_name, handler_type, {
'files' => stager_modinfo['files'] + stage_modinfo['files'],
'paths' => stager_modinfo['paths'] + stage_modinfo['paths'],
'type' => stager_modinfo['type'] })

payload
end

def calculate_adapted_staged(staged_payload:, adapter_name:, stage_info:, stager_info:, adapter_info:)
stage_mod, _, _stage_platform, _stage_arch, _stage_inst = stage_info
stager_mod, handler, _stager_platform, _stager_arch, _stager_inst, _stager_modinfo = stager_info
adapter_mod, _, _adapter_platform, _adapter_arch, adapter_inst, adapter_modinfo = adapter_info

return nil unless adapter_inst.compatible?(staged_payload.new)

payload = build_payload(handler, stager_mod, stage_mod, adapter_mod)

adapted_refname = build_adapted_name(adapter_name, staged_payload.refname)
payload.refname = adapted_refname
payload.framework = framework
payload.file_path = adapter_modinfo['files'][0]
self[payload.refname] = payload
payload
end

# This method is called when a new payload module class is loaded up. For
Expand Down Expand Up @@ -474,4 +559,3 @@ def build_adapted_name(aname, pname)
end

end

0 comments on commit 8e3b8af

Please sign in to comment.