class ProjectDatabase(Base):
__tablename__ = 'projects'
project_uuid = Column(Integer, primary_key=True, autoincrement=True)
project_name = Column(String)
comment = Column(String, default='')
# Marks whether any user has locked ips (hosts), disallowing adding new ones
ips_locked = Column(Boolean, default=False)
hosts_locked = Column(Boolean, default=False)
date_added = Column(DateTime, default=datetime.datetime.utcnow)
ips_relationship = relationship('IPDatabase', cascade="all, delete-orphan")
hosts_relationship = relationship('HostDatabase',
cascade="all, delete-orphan")
tasks_relationship = relationship('TaskDatabase',
cascade="all, delete-orphan")
scans_relationship = relationship('ScanDatabase',
cascade="all, delete-orphan")
session_spawner = Sessions()
def dict(self):
return {
"project_uuid": self.project_uuid,
"project_name": self.project_name,
"comment": self.comment,
"ips_locked": self.ips_locked,
"hosts_locked": self.hosts_locked,
"date_added": str(self.date_added)
}
@classmethod
async def create(cls, project_name):
find_result = await cls.find(project_name=project_name)
if find_result["status"] == "success":
if not find_result["projects"]:
project = cls(project_name=project_name)
try:
with cls.session_spawner.get_session() as session:
session.add(project)
return {"status": "success", "project": project}
except Exception as exc:
return {"status": "error", "text": str(exc)}
return {
"status": "error",
"text": "A project with that name exists"
}
return find_result
@classmethod
def _find(cls, project_name=None, project_uuid=None):
try:
with cls.session_spawner.get_session() as session:
project = session.query(cls)
if project_name is not None:
project = project.filter(cls.project_name == project_name)
if project_uuid is not None:
project = project.filter(cls.project_uuid == project_uuid)
projects = project.all()
return {"status": "success", "projects": projects}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
async def find(cls, *args, **kwargs):
return await asyncify(cls._find)(*args, **kwargs)
@classmethod
@asyncify
def delete(cls, project_uuid=None):
find_result = cls._find(project_uuid=project_uuid)
if find_result["status"] == "success":
if find_result["projects"]:
project = find_result["projects"][0]
project_uuid = project.project_uuid
project_name = project.project_name
try:
with cls.session_spawner.get_session() as session:
session.delete(project)
return {
"status": "success",
"project_uuid": project_uuid,
"project_name": project_name
}
except Exception as exc:
return {"status": "error", "text": str(exc)}
return {
"status": "error",
"text": "This project is already deleted"
}
# Resend the error which occured during find
return find_result
@classmethod
@asyncify
def update(cls,
project_uuid,
new_name=None,
new_comment=None,
ips_locked=None,
hosts_locked=None):
find_result = cls._find(project_uuid=project_uuid)
if find_result["status"] == "success":
if find_result["projects"]:
project = find_result["projects"][0]
try:
if new_name is not None:
project.project_name = new_name
if new_comment is not None:
project.comment = new_comment
if ips_locked is not None:
project.ips_locked = ips_locked
if hosts_locked is not None:
project.hosts_locked = hosts_locked
with cls.session_spawner.get_session() as session:
session.add(project)
return {"status": "success", "project": project}
except Exception as exc:
return {"status": "error", "text": str(exc)}
return {"status": "error", "text": "This project no longer exists"}
# Resend the error which occured during find
return find_result
def __repr__(self):
return "<Project(project_name='%s')>" % (self.project_name)
class CredDatabase(Base):
""" Keeps data on the found credentials """
__tablename__ = "creds"
__table_args__ = (
UniqueConstraint(
'target', 'port_number', 'code', 'candidate', 'mesg', 'project_uuid'
),
)
# Primary key
id = Column(Integer, primary_key=True, autoincrement=True)
# Status code of the request
code = Column(String)
# Response size
size = Column(Integer)
# Time taken to complete the request
time = Column(String)
# Login and password divided by ':'
candidate = Column(String)
# Number of the current request in the whole task.
# Not sure we need it, but patator keeps it
num = Column(Integer)
# Response message for this request
mesg = Column(String)
# Other info from the requester which did not fit to the above fields
other = Column(String)
# Name of the module used to bruteforce the target
# Can be used to extract the name of the service
service = Column(String)
# Target (this is either IP or hostname)
target = Column(String, index=True)
# Port
port_number = Column(Integer)
# ID of the related task (the task, which resulted in the current data)
task_id = Column(String, ForeignKey('tasks.task_id'))
# The name of the related project
project_uuid = Column(
Integer, ForeignKey('projects.project_uuid', ondelete='CASCADE'), index=True
)
# Date of added
date_added = Column(DateTime, default=datetime.datetime.utcnow)
session_spawner = Sessions()
def dict(self):
return {
"id": self.id,
"code": self.code,
"size": self.size,
"time": self.time,
"candidate": self.candidate,
"num": self.num,
"mesg": self.mesg,
"other": self.other,
"service": self.service,
"target": self.target,
"port_number": self.port_number,
"task_id": self.task_id,
"project_uuid": self.project_uuid,
"date_added": str(self.date_added)
}
@classmethod
def create(cls, code=None, size=None, time=None, candidate=None,
num=None, mesg=None, service=service, target=None,
port_number=None, task_id=None, project_uuid=None,
**args
):
db_object = cls(
code=code,
size=size,
time=time,
candidate=candidate,
num=num,
mesg=mesg,
service=service,
target=target,
port_number=port_number,
task_id=task_id,
project_uuid=project_uuid,
other=str(args)
)
try:
with cls.session_spawner.get_session() as session:
session.add(db_object)
except IntegrityError:
return {"status": "success", "target": target}
except Exception as exc:
print(str(exc))
return {"status": "error", "text": str(exc), "target": target}
else:
return {"status": "success", "target": target}
@classmethod
def find(cls, project_uuid, targets=None, port_number=None):
try:
with cls.session_spawner.get_session() as session:
creds_request = session.query(cls).filter(
cls.project_uuid == project_uuid
)
if targets:
creds_request = creds_request.filter(cls.target.in_(targets))
if port_number:
creds_request = creds_request.filter(cls.port_number == port_number)
creds = creds_request.all()
return {"status": "success", "creds": creds}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def count(cls, project_uuid):
try:
with cls.session_spawner.get_session() as session:
amount = session.query(cls).filter(
cls.project_uuid == project_uuid
).count()
return {"status": "success", "amount": amount}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def delete(cls, project_uuid, targets, port_number):
try:
with cls.session_spawner.get_session() as session:
to_delete_creds = session.query(cls).filter(
cls.project_uuid == project_uuid
)
if targets:
to_delete_creds = to_delete_creds.filter(
cls.target.in_(targets)
)
if port_number:
to_delete_creds = to_delete_creds.filter(
cls.port_number == port_number
)
to_delete_creds.delete('fetch')
return {"status": "success"}
except Exception as exc:
return {"status": "error", "text": str(exc)}
class TaskDatabase(Base):
""" Keeps the data of all the tasks that ever existed
in the system (though not the deleted ones) """
__tablename__ = 'tasks'
# Primary key (uuid4)
task_id = Column(String, primary_key=True)
# Literal name of the task: dnsscan, nmap etc.
task_type = Column(String)
# Target can be: hostname, ip, URL
# (depending on the task_type)
target = Column(String)
# Params of the lanuched task (flags, dictionaries etc.)
params = Column(String)
# {New, Working, Finished, Aborted, ...}
status = Column(String)
# Progress in percents
progress = Column(Integer)
# Special note. E.x. error
text = Column(String)
# Stdout
stdout = Column(String)
# Stderr
stderr = Column(String)
# Time of adding
date_added = Column(DateTime, default=datetime.datetime.utcnow)
# Time of adding
date_finished = Column(DateTime)
# The name of the related project
project_uuid = Column(Integer,
ForeignKey('projects.project_uuid',
ondelete='CASCADE'),
index=True)
session_spawner = Sessions()
def dict(self):
return {
"task_id": self.task_id,
"task_type": self.task_type,
"target": self.target,
"params": self.params,
"status": self.status,
"progress": self.progress,
"text": self.text,
"stdout": self.stdout,
"stderr": self.stderr,
"project_uuid": self.project_uuid,
"date_added": self.date_added,
"date_finished": self.date_finished
}
@classmethod
@asyncify
def get_tasks(cls, project_uuid, ips=None, hosts=None):
try:
with cls.session_spawner.get_session() as session:
tasks = session.query(cls)
if ips is not None:
like_filters = list(
map(lambda ip: cls.target.like("%,{},%".format(ip)),
ips)
) + list(
map(lambda ip: cls.target.like("%,{}".format(ip)), ips)
) + list(
map(lambda ip: cls.target.like("{},%".format(ip)),
ips))
tasks = tasks.filter(
or_(cls.target.in_(ips), *like_filters))
if hosts is not None:
like_filters = list(
map(lambda host: cls.target.like("{}:%".format(host)),
hosts))
tasks = tasks.filter(
or_(
# cls.target.in_(hosts),
*like_filters))
if project_uuid is not None:
tasks = tasks.filter(cls.project_uuid == project_uuid)
finished = tasks.filter(
or_(cls.status == 'Finished',
cls.status == 'Aborted')).all()
active = tasks.filter(cls.status != 'Finished',
cls.status != 'Aborted').all()
return {
"status": "success",
"finished": finished,
"active": active
}
except Exception as exc:
return {"status": "error", "text": str(exc)}
class FileDatabase(Base):
""" Keeps data on the found file """
__tablename__ = "files"
__table_args__ = (
UniqueConstraint('file_path', 'status_code', 'content_length', 'project_uuid'),
)
# Primary key
file_id = Column(String, primary_key=True)
# Name of the file
file_name = Column(String)
host_id = Column(Integer, ForeignKey('hosts.id'), index=True)
ip_id = Column(Integer, ForeignKey('ips.id'), index=True)
# Port
port_number = Column(Integer)
# File path
file_path = Column(String)
# Status code of that file
status_code = Column(Integer)
# Content length of that response
content_length = Column(String)
# Special note (for now, we will keep only redirects)
special_note = Column(String)
# ID of the related task (the task, which resulted in the current data)
task_id = Column(String, ForeignKey('tasks.task_id'))
# The name of the related project
project_uuid = Column(
Integer, ForeignKey('projects.project_uuid', ondelete='CASCADE'), index=True
)
# Date of added
date_added = Column(DateTime, default=datetime.datetime.utcnow)
session_spawner = Sessions()
@classmethod
def count(cls, project_uuid):
with cls.session_spawner.get_session() as session:
try:
amount = cls._query_by_project_uuid(session, project_uuid).count()
except Exception as exc:
return {
"status": "error",
"text": str(exc)
}
return {
"status": "success",
"amount": amount
}
@classmethod
def _query_by_project_uuid(cls, session, project_uuid):
return (
session.query(
cls
)
.filter(
cls.project_uuid == project_uuid if project_uuid is not None else True
)
)
@classmethod
def get_stats_for_ips(cls, project_uuid, ip_ids, filters):
stats = {}
try:
with cls.session_spawner.get_session() as session:
status_code_filters = []
if filters and filters[0] != '%':
status_code_filters.append(FileDatabase.status_code.in_(filters))
for ip_id in ip_ids:
prepared_filters = [FileDatabase.ip_id == ip_id] + status_code_filters
files_stats = (
session.query(
FileDatabase.status_code,
FileDatabase.port_number,
func.count(FileDatabase.status_code)
)
.filter(*prepared_filters)
.group_by(
FileDatabase.status_code,
FileDatabase.port_number
)
.all()
)
stats[ip_id] = {}
for status_code, port_number, res in files_stats:
if port_number not in stats[ip_id]:
stats[ip_id][port_number] = {}
stats[ip_id][port_number][status_code] = res
for port_number, stats_for_port in stats[ip_id].items():
stats[ip_id][port_number]['total'] = reduce(
lambda x, y: x + y,
map(lambda stat: stat[1], stats_for_port.items())
)
return {"status": "success", "stats": stats}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def get_stats_for_hosts(cls, project_uuid, host_ids, filters):
stats = {}
try:
with cls.session_spawner.get_session() as session:
status_code_filters = []
if filters and filters[0] != '%':
status_code_filters.append(FileDatabase.status_code.in_(filters))
for host_id in host_ids:
prepared_filters = [FileDatabase.host_id == host_id] + status_code_filters
files_stats = (
session.query(
FileDatabase.status_code,
FileDatabase.port_number,
func.count(FileDatabase.status_code)
)
.filter(*prepared_filters)
.group_by(
FileDatabase.status_code,
FileDatabase.port_number
)
.all()
)
stats[host_id] = {}
for status_code, port_number, res in files_stats:
if port_number not in stats[host_id]:
stats[host_id][port_number] = {}
stats[host_id][port_number][status_code] = res
for port_number, stats_for_port in stats[host_id].items():
stats[host_id][port_number]['total'] = reduce(
lambda x, y: x + y,
map(lambda stat: stat[1], stats_for_port.items())
)
return {"status": "success", "stats": stats}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def get_files_ip(cls, ip, port_number, limit, offset, filters):
try:
files = []
with cls.session_spawner.get_session() as session:
prepared_filters = [
FileDatabase.ip_id == ip,
FileDatabase.port_number == port_number
]
if filters and filters[0] != '%':
prepared_filters.append(FileDatabase.status_code.in_(filters))
files = (
session.query(
FileDatabase,
)
.filter(
*prepared_filters
)
.order_by(FileDatabase.status_code, FileDatabase.file_name)
.limit(limit)
.offset(offset)
.all()
)
return {"status": "success", "files": files}
except Exception as exc:
print(exc)
return {"status": "error", "text": str(exc)}
@classmethod
def get_files_host(cls, host_id, port_number, limit, offset, filters):
try:
files = []
with cls.session_spawner.get_session() as session:
prepared_filters = [
FileDatabase.host_id == host_id,
FileDatabase.port_number == port_number
]
if filters and filters[0] != '%':
prepared_filters.append(FileDatabase.status_code.in_(filters))
files = (
session.query(
FileDatabase,
)
.filter(
*prepared_filters
)
.order_by(FileDatabase.status_code, FileDatabase.file_name)
.limit(limit)
.offset(offset)
.all()
)
return {"status": "success", "files": files}
except Exception as exc:
return {"status": "error", "text": str(exc)}
def dict(self):
return {
"file_id": self.file_id,
"host_id": self.host_id,
"ip_id": self.ip_id,
"file_name": self.file_name,
"port_number": self.port_number,
"file_path": self.file_path,
"status_code": self.status_code,
"content_length": self.content_length,
"special_note": self.special_note,
"task_id": self.task_id,
"project_uuid": self.project_uuid
}
def __repr__(self):
return """
<FileDatabase(file_id='%s', project_uuid='%s', file_name='%s')>""" % (
self.file_id, self.project_uuid,
self.file_name
)
class DictDatabase(Base):
""" Keeps data on the used dictionaries """
__tablename__ = "dicts"
# Primary key
id = Column(Integer, primary_key=True, autoincrement=True)
# Name of the dictionary, used to store the dict on
# the disk of container
name = Column(String)
# Type of the dictionary. Specifically,
# the name of the task which will use the dict
dict_type = Column(String)
# The content of the dictionary
content = Column(Text)
# Amount of lines in the file
lines_count = Column(Integer)
# md5 hashsum of the dictionary to make sure we
# always use the latest version
hashsum = Column(String)
# The name of the related project
project_uuid = Column(Integer,
ForeignKey('projects.project_uuid',
ondelete='CASCADE'),
index=True)
# Date of added
date_added = Column(DateTime, default=datetime.datetime.utcnow)
session_spawner = Sessions()
def dict(self):
return {
"id": self.id,
"name": self.name,
"content": self.content,
"lines_count": self.lines_count,
"hashsum": self.hashsum,
"project_uuid": self.project_uuid,
"date_added": str(self.date_added)
}
@classmethod
def create(cls,
name,
dict_type,
content,
project_uuid,
lines_count=None,
hashsum=None):
db_object = cls(name=name,
dict_type=dict_type,
content=content,
project_uuid=project_uuid,
lines_count=content.count('\n'),
hashsum=hashlib.md5(
content.encode('utf-8')).hexdigest())
try:
with cls.session_spawner.get_session() as session:
session.add(db_object)
except IntegrityError:
return {"status": "success", "name": name, "dictionary": db_object}
except Exception as exc:
return {"status": "error", "text": str(exc), "name": name}
else:
return {"status": "success", "name": name, "dictionary": db_object}
@classmethod
def get(cls, project_uuid=None, dict_id=None, name=None):
try:
with cls.session_spawner.get_session() as session:
dict_request = session.query(
# cls.id,
# cls.name,
# cls.dict_type,
# # cls.content, ## This is too big to keep in memory
# cls.lines_count,
# cls.hashsum,
# cls.project_uuid
cls)
if project_uuid:
dict_request = dict_request.filter(
cls.project_uuid == project_uuid)
if dict_id:
dict_request = dict_request.filter(cls.id == dict_id)
if name:
dict_request = dict_request.filter(cls.name == name)
dicts = dict_request.all()
return {"status": "success", "dicts": dicts}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def count(cls, project_uuid):
try:
with cls.session_spawner.get_session() as session:
amount = session.query(cls).filter(
cls.project_uuid == project_uuid).count()
return {"status": "success", "amount": amount}
except Exception as exc:
return {"status": "error", "text": str(exc)}
@classmethod
def delete(cls, project_uuid, dict_id=None, name=None):
try:
with cls.session_spawner.get_session() as session:
to_delete_dict = session.query(cls).filter(
cls.project_uuid == project_uuid)
if dict_id:
to_delete_dict = to_delete_dict.filter(cls.id == id)
if name:
to_delete_dict = to_delete_dict.filter(cls.name == name)
to_delete_dict.delete('fetch')
return {"status": "success", "dict_id": to_delete_dict.dict_id}
except Exception as exc:
return {"status": "error", "text": str(exc)}
class HostDatabase(Base):
""" Keeps hosts that point to relative IPs """
__tablename__ = 'hosts'
__table_args__ = (UniqueConstraint('target', 'project_uuid'), )
# Primary key (probably uuid4)
id = Column(Integer, primary_key=True, autoincrement=True)
# IP address is a string (probably None, but not sure if
# is needed)
target = Column(String)
# Comment field, as requested by VI
comment = Column(String, default="")
# A list of files which is associated with the current scope
files = relationship('FileDatabase',
cascade="all, delete-orphan",
lazy='select')
# The name of the related project
project_uuid = Column(Integer,
ForeignKey('projects.project_uuid',
ondelete="CASCADE"),
index=True)
# References the task that got this record
# Default is None, as scope can be given by the user manually.
task_id = Column(String,
ForeignKey('tasks.task_id', ondelete='SET NULL'),
default=None)
# Date of adding
date_added = Column(DateTime, default=datetime.datetime.utcnow)
# IP addresses of that host
ip_addresses = relationship("IPDatabase",
secondary=association_table,
back_populates="hostnames",
lazy='noload')
__mapper_args__ = {'concrete': True}
session_spawner = Sessions()
@classmethod
def _find(cls, target, project_uuid):
""" Finds scope (host or ip) in the database """
with cls.session_spawner.get_session() as session:
scope_from_db = session.query(cls).filter(
cls.project_uuid == project_uuid,
cls.target == target).one_or_none()
return scope_from_db
@classmethod
async def find(cls, *args, **kwargs):
return await asyncio.get_event_loop().run_in_executor(
None, lambda: cls._find(*args, **kwargs))
@classmethod
@asyncify
def create(cls, target, project_uuid, task_id=None):
""" Creates a new scope if it is not in the db yet """
if cls._find(target, project_uuid) is None:
try:
new_scope = cls(target=target,
project_uuid=project_uuid,
task_id=task_id)
with cls.session_spawner.get_session() as session:
session.add(new_scope)
except Exception as exc:
return {"status": "error", "text": str(exc)}
else:
return {"status": "success", "new_scope": new_scope}
return {"status": "duplicate", "text": "duplicate"}
@classmethod
@asyncify
def get_or_create(cls, target, project_uuid, task_id=None):
""" Returns a tuple (new_scope, created).
Second value:
False if existed
True if created
"""
found = cls._find(target, project_uuid)
if found is None:
try:
new_scope = cls(target=target,
project_uuid=project_uuid,
task_id=task_id)
with cls.session_spawner.get_session() as session:
session.add(new_scope)
except Exception as exc:
return (None, None)
else:
return (new_scope, True)
return (found, False)
@classmethod
@asyncify
def update(cls, scope_id, comment):
try:
with cls.session_spawner.get_session() as session:
db_object = session.query(cls).filter(cls.id == scope_id).one()
target = db_object.target
db_object.comment = comment
session.add(db_object)
except Exception as exc:
return {"status": "error", "text": str(exc)}
else:
return {"status": "success", "target": target}
@classmethod
def count(cls, project_uuid):
with cls.session_spawner.get_session() as session:
return session.query(cls).filter(
cls.project_uuid == project_uuid).count()
@classmethod
@asyncify
def delete_scope(cls, scope_id):
""" Deletes scope by its id """
try:
with cls.session_spawner.get_session() as session:
db_object = (session.query(HostDatabase).filter(
HostDatabase.id == scope_id).options(
joinedload(HostDatabase.ip_addresses)).one())
target = db_object.target
session.delete(db_object)
except Exception as exc:
print(str(exc))
return {"status": "error", "text": str(exc), "target": scope_id}
else:
return {"status": "success", "target": target}
def dict(self,
include_ports=False,
include_ips=False,
include_files=False,
files_statsified=False):
return {
"host_id":
self.id,
"hostname":
self.target,
"comment":
self.comment,
"project_uuid":
self.project_uuid,
"task_id":
self.task_id,
"ip_addresses":
list(
map(
lambda hostname: hostname.dict(include_ports=include_ports
), self.ip_addresses))
if include_ips else [],
"files": []
}
def __repr__(self):
return """<HostDatabase(host_id='%s', hostname='%s',
ip_addresses='%s', project_uuid='%s')>""" % (
self.id, self.target, self.ip_addresses, self.project_uuid)