class ActivityFeedSchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
itime = fields.Method(serialize='get_itime', deserialize='load_itime', required=True, attribute='start_date')
sum_created_vulnerabilities = fields.Method(serialize='get_sum_created_vulnerabilities', allow_none=True)
sum_created_hosts = fields.Method(serialize='get_sum_created_hosts', allow_none=True)
sum_created_services = fields.Method(serialize='get_sum_created_services', allow_none=True)
sum_created_vulnerability_critical = fields.Method(serialize='get_sum_created_vulnerability_critical',
allow_none=True)
workspace = PrimaryKeyRelatedField('name', dump_only=True)
def load_itime(self, value):
return datetime.datetime.fromtimestamp(value)
def get_itime(self, obj):
return time.mktime(obj.start_date.utctimetuple()) * 1000
def get_sum_created_vulnerabilities(self, obj):
return obj.sum_created_vulnerabilities
def get_sum_created_hosts(self, obj):
return obj.sum_created_hosts
def get_sum_created_services(self, obj):
return obj.sum_created_services
def get_sum_created_vulnerability_critical(self, obj):
return obj.sum_created_vulnerability_critical
class Meta:
model = Command
fields = ('_id', 'command', 'ip', 'hostname',
'params', 'user', 'workspace', 'tool',
'import_source', 'itime', 'sum_created_vulnerabilities',
'sum_created_hosts', 'sum_created_services', 'sum_created_vulnerability_critical')
class WorkspaceSchema(AutoSchema):
name = fields.String(required=True,
validate=validate.Regexp(
r"^[a-z0-9][a-z0-9\_\$\(\)\+\-\/]*$", 0,
"ERORROROR"))
stats = SelfNestedField(WorkspaceSummarySchema())
duration = SelfNestedField(WorkspaceDurationSchema())
_id = fields.Integer(dump_only=True, attribute='id')
scope = MutableField(
PrimaryKeyRelatedField('name', many=True, dump_only=True),
fields.List(fields.String))
create_date = fields.DateTime(attribute='create_date', dump_only=True)
update_date = fields.DateTime(attribute='update_date', dump_only=True)
class Meta:
model = Workspace
fields = ('_id', 'id', 'customer', 'description', 'active', 'duration',
'name', 'public', 'scope', 'stats', 'create_date',
'update_date')
@post_load
def post_load_duration(self, data):
# Unflatten duration (move data[duration][*] to data[*])
duration = data.pop('duration', None)
if duration:
data.update(duration)
return data
class HostSchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
id = fields.Integer()
_rev = fields.String(default='')
ip = fields.String(default='')
description = fields.String(required=True) # Explicitly set required=True
default_gateway = NullToBlankString(attribute="default_gateway_ip",
required=False)
name = fields.String(dump_only=True, attribute='ip', default='')
os = fields.String(default='')
owned = fields.Boolean(default=False)
owner = PrimaryKeyRelatedField('username',
attribute='creator',
dump_only=True)
services = fields.Integer(attribute='open_service_count', dump_only=True)
vulns = fields.Integer(attribute='vulnerability_count', dump_only=True)
credentials = fields.Integer(attribute='credentials_count', dump_only=True)
hostnames = MutableField(
PrimaryKeyRelatedField('name',
many=True,
attribute="hostnames",
dump_only=True,
default=[]), fields.List(fields.String))
metadata = SelfNestedField(MetadataSchema())
type = fields.Function(lambda obj: 'Host', dump_only=True)
service_summaries = fields.Method('get_service_summaries', dump_only=True)
class Meta:
model = Host
fields = ('id', '_id', '_rev', 'ip', 'description', 'mac',
'credentials', 'default_gateway', 'metadata', 'name', 'os',
'owned', 'owner', 'services', 'vulns', 'hostnames', 'type',
'service_summaries')
def get_service_summaries(self, obj):
return [
service.summary for service in obj.services
if service.status == 'open'
]
class CommandSchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
itime = fields.Method(serialize='get_itime',
deserialize='load_itime',
required=True,
attribute='start_date')
duration = fields.Method(serialize='get_duration', allow_none=True)
workspace = PrimaryKeyRelatedField('name', dump_only=True)
def load_itime(self, value):
return datetime.datetime.fromtimestamp(value)
def get_itime(self, obj):
return time.mktime(obj.start_date.utctimetuple()) * 1000
def get_duration(self, obj):
if obj.end_date and obj.start_date:
return (obj.end_date - obj.start_date).seconds + (
(obj.end_date - obj.start_date).microseconds / 1000000.0)
if obj.start_date and not obj.end_date:
return 'In progress'
if not obj.start_date and not obj.end_date:
return 'Not started'
@post_load
def post_load_set_end_date_with_duration(self, data):
# there is a potential bug when updating, the start_date can be changed.
duration = data.pop('duration', None)
if duration:
data['end_date'] = data['start_date'] + datetime.timedelta(
seconds=duration)
class Meta:
model = Command
fields = ('_id', 'command', 'duration', 'itime', 'ip', 'hostname',
'params', 'user', 'workspace', 'tool', 'import_source')
class ServiceSchema(AutoSchema):
_id = fields.Integer(attribute='id', dump_only=True)
_rev = fields.String(default='', dump_only=True)
owned = fields.Boolean(default=False)
owner = PrimaryKeyRelatedField('username',
dump_only=True,
attribute='creator')
port = fields.Integer(dump_only=True) # Port is loaded via ports
ports = MutableField(fields.Integer(),
fields.Method(deserialize='load_ports'),
required=True,
attribute='port')
status = fields.String(default='open',
validate=OneOf(Service.STATUSES),
required=True,
allow_none=False)
parent = fields.Integer(
attribute='host_id') # parent is not required for updates
host_id = fields.Integer(attribute='host_id', dump_only=True)
vulns = fields.Integer(attribute='vulnerability_count', dump_only=True)
credentials = fields.Integer(attribute='credentials_count', dump_only=True)
metadata = SelfNestedField(MetadataSchema())
type = fields.Function(lambda obj: 'Service', dump_only=True)
def load_ports(self, value):
if not isinstance(value, list):
raise ValidationError('ports must be a list')
if len(value) != 1:
raise ValidationError('ports must be a list with exactly one'
'element')
return str(value.pop())
@post_load
def post_load_parent(self, data):
"""Gets the host_id from parent attribute. Pops it and tries to
get a Host with that id in the corresponding workspace.
"""
host_id = data.pop('host_id', None)
if self.context['updating']:
if host_id is None:
# Partial update?
return
if host_id != self.context['object'].parent.id:
raise ValidationError('Can\'t change service parent.')
else:
if not host_id:
raise ValidationError(
'Parent id is required when creating a service.')
try:
data['host'] = Host.query.join(Workspace).filter(
Workspace.name == self.context['workspace_name'],
Host.id == host_id).one()
except NoResultFound:
raise ValidationError(
'Host with id {} not found'.format(host_id))
class Meta:
model = Service
fields = ('id', '_id', 'status', 'parent', 'type', 'protocol',
'description', '_rev', 'owned', 'owner', 'credentials',
'vulns', 'name', 'version', '_id', 'port', 'ports',
'metadata', 'summary', 'host_id')
class VulnerabilitySchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
_rev = fields.String(dump_only=True, default='')
_attachments = fields.Method(serialize='get_attachments', deserialize='load_attachments', default=[])
owned = fields.Boolean(dump_only=True, default=False)
owner = PrimaryKeyRelatedField('username', dump_only=True, attribute='creator')
impact = SelfNestedField(ImpactSchema())
desc = fields.String(attribute='description')
description = fields.String(dump_only=True)
policyviolations = fields.List(fields.String,
attribute='policy_violations')
refs = fields.List(fields.String(), attribute='references')
issuetracker = fields.Method(serialize='get_issuetracker')
parent = fields.Method(serialize='get_parent', deserialize='load_parent', required=True)
parent_type = MutableField(fields.Method('get_parent_type'),
fields.String(),
required=True)
tags = PrimaryKeyRelatedField('name', dump_only=True, many=True)
easeofresolution = fields.String(
attribute='ease_of_resolution',
validate=OneOf(Vulnerability.EASE_OF_RESOLUTIONS),
allow_none=True)
hostnames = PrimaryKeyRelatedField('name', many=True, dump_only=True)
service = fields.Nested(ServiceSchema(only=[
'_id', 'ports', 'status', 'protocol', 'name', 'version', 'summary'
]), dump_only=True)
host = fields.Integer(dump_only=True, attribute='host_id')
severity = SeverityField(required=True)
status = fields.Method(
serialize='get_status',
validate=OneOf(Vulnerability.STATUSES + ['opened']),
deserialize='load_status')
type = fields.Method(serialize='get_type',
deserialize='load_type',
required=True)
obj_id = fields.String(dump_only=True, attribute='id')
target = fields.String(dump_only=True, attribute='target_host_ip')
host_os = fields.String(dump_only=True, attribute='target_host_os')
metadata = SelfNestedField(CustomMetadataSchema())
date = fields.DateTime(attribute='create_date',
dump_only=True) # This is only used for sorting
class Meta:
model = Vulnerability
fields = (
'_id', 'status',
'issuetracker', 'description', 'parent', 'parent_type',
'tags', 'severity', '_rev', 'easeofresolution', 'owned',
'hostnames', 'owner',
'date', 'data', 'refs',
'desc', 'impact', 'confirmed', 'name',
'service', 'obj_id', 'type', 'policyviolations',
'_attachments',
'target', 'host_os', 'resolution', 'metadata')
def get_type(self, obj):
return obj.__class__.__name__
def get_attachments(self, obj):
res = {}
for file_obj in obj.evidence:
ret, errors = EvidenceSchema().dump(file_obj)
if errors:
raise ValidationError(errors, data=ret)
res[file_obj.filename] = ret
return res
def load_attachments(self, value):
return value
def get_parent(self, obj):
return obj.service_id or obj.host_id
def get_parent_type(self, obj):
assert obj.service_id is not None or obj.host_id is not None
return 'Service' if obj.service_id is not None else 'Host'
def get_status(self, obj):
if obj.status == 'open':
return 'opened'
return obj.status
def get_issuetracker(self, obj):
return {}
def load_status(self, value):
if value == 'opened':
return 'open'
return value
def load_type(self, value):
if value == 'Vulnerability':
return 'vulnerability'
if value == 'VulnerabilityWeb':
return 'vulnerability_web'
def load_parent(self, value):
try:
# sometimes api requests send str or unicode.
value = int(value)
except ValueError:
raise ValidationError("Invalid parent type")
return value
@post_load
def post_load_impact(self, data):
# Unflatten impact (move data[impact][*] to data[*])
impact = data.pop('impact', None)
if impact:
data.update(impact)
return data
@post_load
def post_load_parent(self, data):
# schema guarantees that parent_type exists.
parent_class = None
parent_type = data.pop('parent_type', None)
parent_id = data.pop('parent', None)
if not (parent_type and parent_id):
# Probably a partial load, since they are required
return
if parent_type == 'Host':
parent_class = Host
parent_field = 'host_id'
if parent_type == 'Service':
parent_class = Service
parent_field = 'service_id'
if not parent_class:
raise ValidationError('Unknown 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 ValidationError('Parent id not found: {}'.format(parent_id))
data[parent_field] = parent.id
# TODO migration: check what happens when updating the parent from
# service to host or viceverse
return data
class ProfileSchema(Schema):
user = PrimaryKeyRelatedField('username')
first_name = fields.String()
class UserSchema(Schema):
username = fields.String()
blogposts = PrimaryKeyRelatedField(many=True)
class UserSchemaWithTitle(UserSchema):
blogposts = PrimaryKeyRelatedField('title', many=True)