class HostSchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
id = fields.Integer()
_rev = fields.String(default='', dump_only=True)
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)
versions = fields.Method('get_service_version', dump_only=True)
important = fields.Boolean(default=False)
severity_counts = SelfNestedField(HostCountSchema(), 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', 'versions', 'important',
'severity_counts')
def get_service_summaries(self, obj):
return [
service.summary for service in obj.services
if service.status == 'open'
]
def get_service_version(self, obj):
return [
service.version 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 = MutableField(fields.Method(serialize='get_duration'),
fields.Integer(),
allow_none=True)
workspace = PrimaryKeyRelatedField('name', dump_only=True)
creator = PrimaryKeyRelatedField('username', dump_only=True)
metadata = SelfNestedField(MetadataSchema())
def load_itime(self, value):
try:
return datetime.datetime.utcfromtimestamp(value)
except ValueError:
raise ValidationError('Invalid Itime Value')
def get_itime(self, obj):
return obj.start_date.replace(tzinfo=pytz.utc).timestamp() * 1000
def get_duration(self, obj):
# obj.start_date can't be None
if obj.end_date:
return (obj.end_date - obj.start_date).seconds + (
(obj.end_date - obj.start_date).microseconds / 1000000.0)
else:
if (datetime.datetime.utcnow() - obj.start_date
).total_seconds() > 86400: # 86400 is 1d TODO BY CONFIG
return 'Timeout'
return 'In progress'
@post_load
def post_load_set_end_date_with_duration(self, data, **kwargs):
# 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)
return data
class Meta:
model = Command
fields = ('_id', 'command', 'duration', 'itime', 'ip', 'hostname',
'params', 'user', 'creator', 'workspace', 'tool',
'import_source', 'metadata')
class CredentialSchema(AutoSchema):
_id = fields.Integer(dump_only=True, attribute='id')
_rev = fields.String(default='', dump_only=True)
owned = fields.Boolean(default=False, dump_only=True)
owner = fields.String(dump_only=True, attribute='creator.username',
default='')
username = fields.String(default='', required=True,
validate=validate.Length(min=1, error="Username must be defined"))
password = fields.String(default='')
description = fields.String(default='')
couchdbid = fields.String(default='') # backwards compatibility
parent_type = MutableField(fields.Method('get_parent_type'),
fields.String(),
required=True)
parent = MutableField(fields.Method('get_parent'),
fields.Integer(),
required=True)
host_ip = fields.String(dump_only=True, attribute="host.ip",
default=None)
service_name = fields.String(dump_only=True, attribute="service.name",
default=None)
target = fields.String(dump_only=True, attribute="target_ip")
# for filtering
host_id = fields.Integer(load_only=True)
service_id = fields.Integer(load_only=True)
metadata = SelfNestedField(MetadataSchema())
def get_parent(self, obj):
return obj.host_id or obj.service_id
def get_parent_type(self, obj):
assert obj.host_id is not None or obj.service_id is not None
return 'Service' if obj.service_id is not None else 'Host'
def get_target(self, obj):
if obj.host is not None:
return obj.host.ip
else:
return obj.service.host.ip + '/' + obj.service.name
class Meta:
model = Credential
fields = ('id', '_id', "_rev", 'parent', 'username', 'description',
'name', 'password', 'owner', 'owned', 'couchdbid', 'parent',
'parent_type', 'metadata', 'host_ip', 'service_name',
'target',
)
@post_load
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
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, required=True,
validate=[Range(min=0, error="The value must be greater than or equal to 0")]) # Port is loaded via ports
ports = MutableField(fields.Integer(required=True,
validate=[Range(min=0, error="The value must be greater than or equal to 0")]),
fields.Method(deserialize='load_ports'),
required=True,
attribute='port')
status = fields.String(missing='open', validate=OneOf(Service.STATUSES),
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)
summary = fields.String(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')
port = value.pop()
if isinstance(port, str):
try:
port = int(port)
except ValueError:
raise ValidationError('The value must be a number')
if port > 65535 or port < 1:
raise ValidationError('The value must be in the range [1-65535]')
return str(port)
@post_load
def post_load_parent(self, data, **kwargs):
"""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 data
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(f'Host with id {host_id} not found')
return data
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')