def __init__(self, scenario_name):
# Variables
self.scenario_name = scenario_name
self.scenario_id = get_unique_id(length=8)
now = datetime.now()
self.creation_date = now.strftime("%d/%m/%Y %H:%M:%S")
self.last_accessed = self.creation_date[:]
self.description = ""
self.exploit = Exploit()
self.vulnerability = Vulnerability()
self.machines = dict()
def objectFromDictionary(self, dict):
"""
Creates a Scenario object from a dictionary.
:param dict: A dictionary containing the Scenario's data
:return: A Scenario object
"""
self.scenario_id = dict["scenario_id"]
self.creation_date = dict["creation_date"]
self.last_accessed = dict["last_accessed"]
self.description = dict["description"]
self.exploit = Exploit().objectFromDictionary(dict["exploit"])
self.vulnerability = Vulnerability().objectFromDictionary(
dict["vulnerability"])
for m in dict["machines"]:
self.machines[m] = VirtualMachine().objectFromDictionary(
dict["machines"][m])
return self
def newEmpty(self, vulnerability_name):
"""
Creates a new vulnerability which includes the folders and the vulnerability JSON file
:param vulnerability_name: String with the vulnerability name
:return: Response object containing the status of the request
"""
response = Response()
if vulnerability_name not in self.vulnerabilities_dict:
vulnerability = Vulnerability(vulnerability_name)
self.vulnerabilities_dict[vulnerability_name] = vulnerability
self.db_manager.insertVulnerability(
vulnerability.dictionary().copy())
response.setResponse(True)
response.setBody(vulnerability.dictionary())
else:
response.setResponse(False)
response.setReason('Vulnerability already exist')
response.setBody(dict())
return response.dictionary()
def addVulnerabilitiesToDB(self):
"""
Pre-populates the database with vulnerabilities.
:return: None
"""
vulnerabilities_to_add = ['rConfig_3_9_searchColumn_SQL_Injection']
currentVulnerabilities = self.getVulnerabilities()
vulnerabilities_list = [
vulnerability['name'] for vulnerability in currentVulnerabilities
]
vulnerabilities_set = set(vulnerabilities_list)
for vulnerability_name in vulnerabilities_to_add:
if vulnerability_name not in vulnerabilities_set:
json_name = ''.join([vulnerability_name, ".json"])
with open(
self.file_manager.getVulnerabilityJSONPath(
vulnerability_name) / json_name) as outfile:
vulnerability_dict = json.load(outfile)
vulnerability = Vulnerability().objectFromDictionary(
vulnerability_dict)
self.insertVulnerability(vulnerability.dictionary().copy())
return
def editOne(self, vulnerability_json):
"""
Edits a current scenario with a JSON file
:param vulnerability_json: JSON file new vulnerability
:return: Response object containing the status of the request
"""
response = Response()
print(vulnerability_json)
vulnerability_name = vulnerability_json["name"]
if vulnerability_name in self.vulnerabilities_dict:
vulnerability_json = Vulnerability().objectFromDictionary(
vulnerability_json)
self.vulnerabilities_dict[vulnerability_name] = vulnerability_json
self.db_manager.editVulnerability(
vulnerability_json.dictionary().copy())
response.setResponse(True)
response.setBody(
self.vulnerabilities_dict[vulnerability_name].dictionary())
else:
response.setReason('Vulnerability doesn\'t exist')
response.setResponse(False)
response.setBody(dict())
return response.dictionary()
def _initializeFromDatabase(self):
"""
Pre-populates the database with vulnerabilities.
:return: Dictionary containing scenario's data
"""
# Variables
vulnerabilities_dict = dict()
vulnerabilities = self.db_manager.getVulnerabilities()
for raw_vulnerability in vulnerabilities:
del raw_vulnerability["_id"]
vulnerability_name = raw_vulnerability["name"]
vulnerability = Vulnerability().objectFromDictionary(
raw_vulnerability)
vulnerabilities_dict[vulnerability_name] = vulnerability
return vulnerabilities_dict
class Scenario(Entity):
def __init__(self, scenario_name):
# Variables
self.scenario_name = scenario_name
self.scenario_id = get_unique_id(length=8)
now = datetime.now()
self.creation_date = now.strftime("%d/%m/%Y %H:%M:%S")
self.last_accessed = self.creation_date[:]
self.description = ""
self.exploit = Exploit()
self.vulnerability = Vulnerability()
self.machines = dict()
def setExploit(self, exploit):
"""
Sets the exploit info for this scenario
:param exploit: Object which carries the exploit info
"""
self.exploit = exploit
def setVulnerability(self, vulnerability):
"""
Sets the vulnerability info for this scenario
:param vulnerability: Object which carries the vulnerability info
"""
self.vulnerability = vulnerability
def addVM(self, vm):
"""
Adds a new virtual machine to this scenario
:param vm: Object which carries the virtual machine data
"""
self.machines[vm.name] = vm
def dictionary(self):
"""
Generates a dictionary for the Scenario object
:return: A dictionary with Scenario's data
"""
scenario_dict = dict()
scenario_dict["scenario_name"] = self.scenario_name
scenario_dict["scenario_id"] = self.scenario_id
scenario_dict["creation_date"] = self.creation_date
scenario_dict["last_accessed"] = self.last_accessed
scenario_dict["description"] = self.description
scenario_dict["exploit"] = self.exploit.dictionary()
scenario_dict["vulnerability"] = self.vulnerability.dictionary()
scenario_dict["machines"] = dict()
for name in self.machines:
scenario_dict["machines"][name] = self.machines[name].dictionary()
return scenario_dict
def objectFromDictionary(self, dict):
"""
Creates a Scenario object from a dictionary.
:param dict: A dictionary containing the Scenario's data
:return: A Scenario object
"""
self.scenario_id = dict["scenario_id"]
self.creation_date = dict["creation_date"]
self.last_accessed = dict["last_accessed"]
self.description = dict["description"]
self.exploit = Exploit().objectFromDictionary(dict["exploit"])
self.vulnerability = Vulnerability().objectFromDictionary(
dict["vulnerability"])
for m in dict["machines"]:
self.machines[m] = VirtualMachine().objectFromDictionary(
dict["machines"][m])
return self