class PayloadsTests:
def __init__(self):
self.badges = Badges()
self.importer = Importer()
self.local_storage = LocalStorage()
self.payloads = Payloads()
def perform_test(self):
fail = False
all_payloads = self.local_storage.get("payloads")
if all_payloads:
for database in all_payloads.keys():
self.badges.output_process("Testing payloads from " +
database + " database...")
payloads = all_payloads[database]
for platform in payloads.keys():
for architecture in payloads[platform].keys():
for payload in payloads[platform][architecture].keys():
try:
_ = self.importer.import_payload(
payloads[platform][architecture][payload]
['Path'])
self.badges.output_success(
self.payloads.get_full_name(
platform, architecture, payload) +
': OK')
except Exception:
self.badges.output_error(
self.payloads.get_full_name(
platform, architecture, payload) +
': FAIL')
fail = True
return fail
class Payloads:
def __init__(self):
self.importer = Importer()
self.local_storage = LocalStorage()
self.badges = Badges()
def check_exist(self, name):
if self.check_style(name):
all_payloads = self.local_storage.get("payloads")
if all_payloads:
for database in all_payloads.keys():
payloads = all_payloads[database]
platform = self.get_platform(name)
architecture = self.get_architecture(name)
if platform in payloads.keys():
if architecture in payloads[platform].keys():
payload = self.get_name(name)
if payload in payloads[platform][
architecture].keys():
return True
return False
@staticmethod
def check_style(name):
if len(name.split('/')) >= 3:
return True
return False
def get_platform(self, name):
if self.check_style(name):
return name.split('/')[0]
return None
def get_architecture(self, name):
if self.check_style(name):
return name.split('/')[1]
return None
def get_name(self, name):
if self.check_style(name):
return os.path.join(*(name.split(os.path.sep)[2:]))
return None
def get_payload_object(self, platform, architecture, name):
payload_full_name = self.get_full_name(platform, architecture, name)
if self.check_exist(payload_full_name):
database = self.get_database(payload_full_name)
return self.local_storage.get(
"payloads")[database][platform][architecture][name]
return None
def get_database(self, name):
if self.check_style(name):
all_payloads = self.local_storage.get("payloads")
if all_payloads:
for database in all_payloads.keys():
payloads = all_payloads[database]
platform = self.get_platform(name)
architecture = self.get_architecture(name)
if platform in payloads.keys():
if architecture in payloads[platform].keys():
payload = self.get_name(name)
if payload in payloads[platform][
architecture].keys():
return database
return None
@staticmethod
def get_full_name(platform, architecture, name):
return platform + '/' + architecture + '/' + name
def check_current_module(self):
if self.local_storage.get("current_module"):
if len(self.local_storage.get("current_module")) > 0:
return True
return False
def get_current_module_object(self):
if self.check_current_module():
return self.local_storage.get_array(
"current_module",
self.local_storage.get("current_module_number"))
return None
def get_current_module_name(self):
if self.check_current_module():
return self.local_storage.get_array(
"current_module",
self.local_storage.get(
"current_module_number")).details['Module']
return None
def import_payload(self, module_name, platform, architecture, name):
payloads = self.get_payload_object(platform, architecture, name)
try:
payload_object = self.importer.import_payload(payloads['Path'])
current_module_name = module_name
imported_payloads = self.local_storage.get("imported_payloads")
if imported_payloads:
if current_module_name in imported_payloads.keys():
imported_payloads[current_module_name].update({
payload_object.details['Payload']:
copy.copy(payload_object)
})
else:
imported_payloads.update({
current_module_name: {
payload_object.details['Payload']:
copy.copy(payload_object)
}
})
else:
imported_payloads = {
current_module_name: {
payload_object.details['Payload']:
copy.copy(payload_object)
}
}
self.local_storage.set("imported_payloads", imported_payloads)
except Exception:
return None
return payload_object
def check_imported(self, module_name, name):
imported_payloads = self.local_storage.get("imported_payloads")
current_module_name = module_name
if imported_payloads:
if current_module_name in imported_payloads.keys():
if name in imported_payloads[current_module_name].keys():
return True
return False
def get_current_payload(self):
imported_payloads = self.local_storage.get("imported_payloads")
current_module_object = self.get_current_module_object()
current_module_name = current_module_object.details['Module']
if hasattr(current_module_object, "payload"):
name = current_module_object.payload['Value']
if current_module_name in imported_payloads.keys():
if name in imported_payloads[current_module_name].keys():
return imported_payloads[current_module_name][name]
return None
def add_payload(self, module_name, platform, architecture, name):
payloads = self.get_payload_object(platform, architecture, name)
not_installed = list()
for dependence in payloads['Dependencies']:
if not self.importer.import_check(dependence):
not_installed.append(dependence)
if not not_installed:
full_name = self.get_full_name(platform, architecture, name)
if not self.check_imported(module_name, full_name):
payload_object = self.import_payload(module_name, platform,
architecture, name)
if not payload_object:
self.badges.output_error(
"Failed to select payload from database!")
return False
else:
self.badges.output_error(
"Payload depends this dependencies which is not installed:")
for dependence in not_installed:
self.badges.output_empty(" * " + dependence)
return False
return True