def _get_schema_class(self):
assert self.schema_class_dict is not None, "You must define schema_class"
if request.method == 'POST':
requested_type = request.json.get('type', None)
if not requested_type:
raise InvalidUsage('Type is required.')
if requested_type not in self.schema_class_dict:
raise InvalidUsage('Invalid vulnerability type.')
return self.schema_class_dict[requested_type]
# We use web since it has all the fields
return self.schema_class_dict['VulnerabilityWeb']
def _perform_create(self, data, workspace_name):
model = {'host': Host, 'service': Service, 'comment': Comment}
obj = db.session.query(model[data['object_type']]).get(
data['object_id'])
workspace = self._get_workspace(workspace_name)
if not obj:
raise InvalidUsage('Can\'t comment inexistent object')
if obj.workspace != workspace:
raise InvalidUsage('Can\'t comment object of another workspace')
return super(CommentCreateMixing,
self)._perform_create(data, workspace_name)
def set_parent(self, data, **kwargs):
parent_type = data.pop('parent_type', None)
parent_id = data.pop('parent', None)
if parent_type == 'Host':
parent_class = Host
parent_field = 'host_id'
not_parent_field = 'service_id'
elif parent_type == 'Service':
parent_class = Service
parent_field = 'service_id'
not_parent_field = 'host_id'
elif 'partial' in kwargs and kwargs['partial']:
return data
else:
raise ValidationError(
f'Unknown parent type: {parent_type}')
try:
parent = db.session.query(parent_class).join(Workspace).filter(
Workspace.name == self.context['workspace_name'],
parent_class.id == parent_id).one()
except NoResultFound:
raise InvalidUsage(f'Parent id not found: {parent_id}')
data[parent_field] = parent.id
data[not_parent_field] = None
return data
