def test_different_object_type(self, vulnerability, workspace,
empty_command_factory, session):
command = empty_command_factory.create(workspace=workspace)
session.add(vulnerability)
session.flush()
invalid_co = CommandObject.create(vulnerability, command)
invalid_co.object_type = 'host'
session.add(invalid_co)
session.commit()
command = empty_command_factory.create(workspace=workspace)
session.flush()
CommandObject.create(vulnerability, command)
for _ in range(5):
new_command = empty_command_factory.create(workspace=workspace)
session.flush()
session.add(
CommandObject(vulnerability,
command=new_command,
workspace=workspace,
created_persistent=False))
session.commit()
assert vulnerability.creator_command_id == command.id
assert vulnerability.creator_command_tool == command.tool
# I'm Py3
def _set_command_id(self, obj, created):
try:
# validates the data type from user input.
command_id = int(flask.request.args.get('command_id', None))
except TypeError:
command_id = None
if command_id:
command = db.session.query(Command).filter(
Command.id == command_id,
Command.workspace == obj.workspace).first()
if command is None:
raise InvalidUsage('Command not found.')
# if the object is created and updated in the same command
# the command object already exists
# we skip the creation.
object_type = obj.__class__.__table__.name
command_object = CommandObject.query.filter_by(
object_id=obj.id,
object_type=object_type,
command=command,
workspace=obj.workspace,
).first()
if created or not command_object:
command_object = CommandObject(object_id=obj.id,
object_type=object_type,
command=command,
workspace=obj.workspace,
created_persistent=created)
db.session.add(command)
db.session.add(command_object)
def _create_command_object_for(ws, created, obj, command):
assert command is not None
db.session.add(
CommandObject(obj,
command=command,
created_persistent=created,
workspace=ws))
db.session.commit()
def test_creator_command(self, workspace, vulnerability,
empty_command_factory, session):
command = empty_command_factory.create(workspace=workspace)
session.add(vulnerability)
session.flush()
assert vulnerability.id is not None
assert command.id is not None
assert vulnerability.workspace is workspace
assert command.workspace is workspace
CommandObject.create(vulnerability, command)
for _ in range(5):
new_command = empty_command_factory.create(workspace=workspace)
session.flush()
session.add(
CommandObject(vulnerability,
command=new_command,
workspace=workspace,
created_persistent=False))
session.commit()
assert vulnerability.creator_command_id == command.id
assert vulnerability.creator_command_tool == command.tool
def test_creates_command_object_on_duplicates(session, command, service,
vulnerability_factory,
vulnerability_web_factory,
credential_factory):
vuln_host = vulnerability_factory.create(workspace=service.workspace,
host=service.host,
service=None)
vuln_service = vulnerability_factory.create(workspace=service.workspace,
service=service,
host=None)
vuln_web = vulnerability_web_factory.create(workspace=service.workspace,
service=service)
host_cred = credential_factory.create(workspace=service.workspace,
host=service.host,
service=None)
session.add(command)
session.add(service)
session.add(vuln_host)
session.add(vuln_service)
session.add(vuln_web)
session.add(host_cred)
session.commit()
assert command.workspace == service.workspace
assert len(command.workspace.command_objects) == 0
objects_with_command_object = [
('host', service.host),
('service', service),
('vulnerability', vuln_host),
('vulnerability', vuln_service),
('vulnerability', vuln_web),
# ('credential', host_cred), # Commented because unique constraint of credential is not working
]
for (table_name, obj) in objects_with_command_object:
assert obj.id is not None and command.id is not None
db.session.add(
CommandObject(
object_type=table_name,
object_id=obj.id,
command=command,
created_persistent=True,
workspace=command.workspace,
))
session.commit()
data = {
'hosts': [{
'ip':
service.host.ip,
'description':
service.host.description,
'vulnerabilities': [{
'name': vuln_host.name,
'severity': 'high',
'desc': vuln_host.description,
'type': 'Vulnerability',
}],
'credentials': [{
'name': host_cred.name,
'username': host_cred.username,
}],
'services': [{
'name':
service.name,
'protocol':
service.protocol,
'port':
service.port,
'vulnerabilities': [
{
'name': vuln_service.name,
'severity': 'high',
'desc': vuln_service.description,
'type': 'Vulnerability',
},
{
'name': vuln_web.name,
'severity': 'high',
'desc': vuln_web.description,
'type': 'VulnerabilityWeb',
'method': vuln_web.method,
'pname': vuln_web.parameter_name,
'path': vuln_web.path,
'website': vuln_web.website,
},
]
}]
}]
}
data['command'] = command_data.copy()
command2 = new_empty_command(command.workspace)
bc.bulk_create(command.workspace, command2, data)
assert count(Command, command.workspace) == 2
new_command = Command.query.filter_by(tool='pytest').one()
for (table_name, obj) in objects_with_command_object:
assert obj.id is not None and new_command.id is not None
CommandObject.query.filter(
CommandObject.workspace == command.workspace,
CommandObject.command == new_command,
CommandObject.object_type == table_name,
CommandObject.object_id == obj.id,
CommandObject.created_persistent == false(),
).one()
def populate(self, workspace, service, session, user,
vulnerability_factory, credential_factory,
empty_command_factory):
session.commit()
self.session = session
assert service.workspace_id == workspace.id
workspace.set_scope(['*.infobytesec.com', '192.168.1.0/24'])
self.user = user
self.workspace = workspace
self.permission = WorkspacePermission(user=user, workspace=workspace)
session.add(self.permission)
self.host = service.host
self.host.set_hostnames(['a.com', 'b.com'])
self.service = service
self.host_cred = credential_factory.create(
host=self.host,
service=None,
workspace=workspace,
creator=user,
)
self.service_cred = credential_factory.create(
host=None,
service=service,
workspace=workspace,
creator=user,
)
self.host_vuln = vulnerability_factory.create(
host=self.host,
service=None,
workspace=workspace,
creator=user,
)
self.service_vuln = vulnerability_factory.create(
host=None,
service=service,
workspace=workspace,
creator=user,
)
session.flush()
for vuln in [self.host_vuln, self.service_vuln]:
vuln.references = ['CVE-1234', 'CVE-4331']
vuln.policy_violations = ["PCI-DSS"]
self.attachment = File(
name='test.png',
filename='test.png',
content=b'test',
object_type='vulnerability',
object_id=self.service_vuln.id,
creator=user,
)
self.session.add(self.attachment)
self.host_attachment = File(
name='test.png',
filename='test.png',
content=b'test',
object_type='host',
object_id=self.host.id,
creator=user,
)
self.session.add(self.host_attachment)
self.comment = Comment(
text="test",
object_type='host',
object_id=self.host.id,
workspace=self.workspace,
creator=user,
)
self.session.add(self.comment)
self.reply_comment = Comment(
text="ok",
object_type='host',
object_id=self.host.id,
workspace=self.workspace,
reply_to=self.comment,
creator=user,
)
self.command = empty_command_factory.create(workspace=workspace,
creator=user)
CommandObject.create(self.host_vuln, self.command)
CommandObject.create(self.service_vuln, self.command)
self.methodology_template = MethodologyTemplate(name="test", )
session.add(self.methodology_template)
self.methodology_template_task = TaskTemplate(
name="aaaa", template=self.methodology_template)
session.add(self.methodology_template)
self.methodology = Methodology(name="test",
template=self.methodology_template,
workspace=self.workspace)
session.add(self.methodology)
self.methodology_task = Task(name="aaaa",
workspace=self.workspace,
template=self.methodology_template_task,
methodology=self.methodology)
session.add(self.methodology_template_task)
self.methodology_task_assigned = TaskAssignedTo(
task=self.methodology_task,
user=self.user,
)
session.add(self.methodology_task_assigned)
session.commit()