Various fixes.

This commit is contained in:
Daniele Maglie 2024-02-07 14:54:02 +01:00
parent 5397063bbd
commit c3c7933d2a
9 changed files with 223 additions and 36 deletions

View file

@ -86,6 +86,33 @@ class AccountStorage(DatabaseArangodbDocument):
except: except:
raise AccountStorageGetException() raise AccountStorageGetException()
def get_active(self) -> Cursor | Exception:
aql = f"""
FOR document IN {self.collection_name}
FILTER document.is_active == true
RETURN {{
'account_id': document.account_id,
'name': document.name,
'email': document.email,
'role': document.role
}}
"""
try:
return self.query(aql)
except:
raise AccountStorageGetException()
def get_total(self) -> int | Exception:
try:
return self.collection_instance.count()
except:
raise AccountStorageGetException()
def exists_by_identifier(self, identifier: str) -> Cursor | Exception: def exists_by_identifier(self, identifier: str) -> Cursor | Exception:
""" """
Checks if an account with this identifier is in the collection. Checks if an account with this identifier is in the collection.

View file

@ -31,7 +31,7 @@ class AzureBlobStorage(BlobStorageDriver):
def create_container(self): def create_container(self):
""" """
Creates the container. Creates the container.
Currently used to mock a blob storage. Currently used to create a mock blob storage.
""" """
return self.container_client.create_container() return self.container_client.create_container()
@ -79,7 +79,10 @@ class AzureBlobStorage(BlobStorageDriver):
return True return True
except (ResourceNotFoundError, ResourceExistsError, AzureError): except ResourceExistsError:
raise AzureBlobStorageAlreadyExistsException()
except (ResourceNotFoundError, AzureError):
raise AzureBlobStorageUploadException() raise AzureBlobStorageUploadException()
def get_list(self): def get_list(self):
@ -107,6 +110,14 @@ class AzureBlobStorage(BlobStorageDriver):
except (ResourceNotFoundError, AzureError): except (ResourceNotFoundError, AzureError):
raise AzureBlobStorageRemoveException() raise AzureBlobStorageRemoveException()
class AzureBlobStorageAlreadyExistsException(Exception):
"""
The blob already exists.
"""
pass
class AzureBlobStorageUploadException(Exception): class AzureBlobStorageUploadException(Exception):
""" """

View file

@ -12,19 +12,7 @@ class BlobStorage:
self.driver = driver self.driver = driver
def upload(self, blob_name: str, file_path: str): def upload(self, blob_name: str, file_path: str):
try:
return self.driver.upload( return self.driver.upload(
blob_name = blob_name, blob_name = blob_name,
file_path = file_path file_path = file_path
) )
except:
raise BlobStorageUploadException()
class BlobStorageUploadException(Exception):
"""
Cannot upload the file.
"""
pass

View file

@ -4,6 +4,8 @@ from arango.exceptions import AQLQueryExecuteError
class DatabaseArangodbDocument(DatabaseArangodbConnection): class DatabaseArangodbDocument(DatabaseArangodbConnection):
max_batch_size = 50000
def collection(self, collection): def collection(self, collection):
try: try:
return self.instance.collection(collection) return self.instance.collection(collection)
@ -13,7 +15,7 @@ class DatabaseArangodbDocument(DatabaseArangodbConnection):
def query(self, aql, **kwargs): def query(self, aql, **kwargs):
try: try:
return self.instance.aql.execute(aql, **kwargs) return self.instance.aql.execute(aql, batch_size = self.max_batch_size, **kwargs)
except AQLQueryExecuteError: except AQLQueryExecuteError:
raise DatabaseArangodbQueryException() raise DatabaseArangodbQueryException()

View file

@ -27,6 +27,75 @@ class DDAStorage(DatabaseArangodbDocument):
except: except:
raise DDAStorageCreateException() raise DDAStorageCreateException()
def get_by_identifier(self, dda_id: str) -> Cursor | Exception:
"""
Fetches a single DDA instance by its identifier.
:param dda_id: a valid DDA identifier.
:return: cursor with the requested data.
"""
aql = f"""
FOR document IN {self.collection_name}
FILTER
document.dda_id == @dda_id
RETURN {{
'dda_id': document.dda_id,
'description': document.description,
'instance': document.instance,
'is_active': document.is_active,
'metadata': {{
'created_at': document.metadata.created_at
}}
}}
"""
try:
return self.query(aql, bind_vars = {
'dda_id': dda_id
})
except:
raise DDAStorageGetException()
def get_by_identifier_for_reporter(self, dda_id: str, reporter_id: str) -> Cursor | Exception:
"""
Fetches a single DDA instance by its identifier for reporter.
:param dda_id: a valid DDA identifier.
:param account_id: a valid reporter identifier.
:return: cursor with the requested data.
"""
aql = f"""
FOR document IN {self.collection_name}
FILTER
document.dda_id == @dda_id AND
document.account_id == @reporter_id
RETURN {{
'dda_id': document.dda_id,
'description': document.description,
'instance': document.instance,
'is_active': document.is_active,
'metadata': {{
'created_at': document.metadata.created_at
}}
}}
"""
try:
return self.query(aql, bind_vars = {
'dda_id': dda_id,
'reporter_id': reporter_id
})
except:
raise DDAStorageGetException()
def get_global(self) -> Cursor | Exception: def get_global(self) -> Cursor | Exception:
""" """
Fetches all the DDA instances. Fetches all the DDA instances.

View file

@ -52,9 +52,7 @@ class ForensicStorage(DatabaseArangodbDocument):
FILTER document.hash_string == @hash_string FILTER document.hash_string == @hash_string
COLLECT WITH COUNT INTO length RETURN document
RETURN length
""" """
try: try:

View file

@ -29,6 +29,20 @@ class TicketItemStorage(DatabaseArangodbDocument):
except: except:
raise TicketItemStorageCreateException() raise TicketItemStorageCreateException()
def insert_many(self, documents: list) -> dict | Exception:
"""
Adds a batch of new ticket items.
:param documents: a list of dictionary ticket items.
:return: cursor with the inserted data.
"""
try:
return self.collection_instance.insert_many(documents)
except:
raise TicketItemStorageCreateException()
def get_all_items_with_genre(self, genre: str) -> Cursor: def get_all_items_with_genre(self, genre: str) -> Cursor:
""" """
Gets all ticket items. Values only. Gets all ticket items. Values only.
@ -76,7 +90,7 @@ class TicketItemStorage(DatabaseArangodbDocument):
FOR parent_ticket IN {self.ticket_collection_name} FOR parent_ticket IN {self.ticket_collection_name}
FILTER FILTER
parent_ticket.ticket_id == document.ticket_id AND parent_ticket.ticket_id == document.ticket_id AND
parent_ticket.status.ticket != 'created' (parent_ticket.status == 'open' OR parent_ticket.status == 'closed')
RETURN DISTINCT document.value RETURN DISTINCT document.value
""" """
@ -179,7 +193,7 @@ class TicketItemStorage(DatabaseArangodbDocument):
FOR parent_ticket IN {self.ticket_collection_name} FOR parent_ticket IN {self.ticket_collection_name}
FILTER FILTER
parent_ticket.ticket_id == document.ticket_id AND parent_ticket.ticket_id == document.ticket_id AND
parent_ticket.status.ticket != 'created' (parent_ticket.status == 'open' OR parent_ticket.status == 'closed')
RETURN DISTINCT document.value RETURN DISTINCT document.value
""" """
@ -404,7 +418,13 @@ class TicketItemStorage(DatabaseArangodbDocument):
document.is_whitelisted == false AND document.is_whitelisted == false AND
document.is_error == false document.is_error == false
RETURN {{ // ensure only available tickets are considered
FOR parent_ticket IN {self.ticket_collection_name}
FILTER
parent_ticket.ticket_id == document.ticket_id AND
(parent_ticket.status == 'open' OR parent_ticket.status == 'closed')
RETURN DISTINCT {{
'value': document.value, 'value': document.value,
'status': document.status, 'status': document.status,
'timestamp': document.timestamp, 'timestamp': document.timestamp,
@ -510,6 +530,37 @@ class TicketItemStorage(DatabaseArangodbDocument):
except: except:
raise TicketItemStorageGetException() raise TicketItemStorageGetException()
def get_active(self) -> Cursor | Exception:
"""
Get all the ticket items without any unblocked items, grouped by genre.
:return: cursor with the requested data.
"""
aql = f"""
LET result = (
FOR document IN {self.collection_name}
// filter out ticket items removed by reported errors
FILTER document.is_error == false
COLLECT genre = document.genre INTO groupedDocuments
RETURN {{
'genre': genre,
'values': UNIQUE(groupedDocuments[*].document.value)
}}
)
RETURN MERGE(FOR r IN result RETURN {{[r.genre]: r.values}})
"""
try:
return self.query(aql)
except:
raise TicketItemStorageGetException()
def exists_by_value(self, genre: str, value: str) -> Cursor | Exception: def exists_by_value(self, genre: str, value: str) -> Cursor | Exception:
""" """
Searches for a duplicate. Searches for a duplicate.
@ -582,7 +633,7 @@ class TicketItemStorage(DatabaseArangodbDocument):
FOR parent_ticket IN {self.ticket_collection_name} FOR parent_ticket IN {self.ticket_collection_name}
FILTER FILTER
parent_ticket.ticket_id == document.ticket_id AND parent_ticket.ticket_id == document.ticket_id AND
parent_ticket.status.ticket != 'created' (parent_ticket.status == 'open' OR parent_ticket.status == 'closed')
UPDATE document WITH {{ UPDATE document WITH {{
status: @status, status: @status,

View file

@ -63,6 +63,7 @@ class TicketStorage(DatabaseArangodbDocument):
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'dda_id': document.dda_id,
'description': document.description, 'description': document.description,
'fqdn': document.fqdn, 'fqdn': document.fqdn,
'ipv4': document.ipv4, 'ipv4': document.ipv4,
@ -131,6 +132,7 @@ class TicketStorage(DatabaseArangodbDocument):
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'dda_id': document.dda_id,
'description': document.description, 'description': document.description,
'fqdn': document.fqdn, 'fqdn': document.fqdn,
'ipv4': document.ipv4, 'ipv4': document.ipv4,
@ -141,7 +143,8 @@ class TicketStorage(DatabaseArangodbDocument):
}}, }},
'settings': {{ 'settings': {{
'revoke_time': document.settings.revoke_time, 'revoke_time': document.settings.revoke_time,
'autoclose_time': document.settings.autoclose_time 'autoclose_time': document.settings.autoclose_time,
'report_error_time': document.settings.report_error_time
}} }}
}} }}
""" """
@ -168,7 +171,7 @@ class TicketStorage(DatabaseArangodbDocument):
FILTER FILTER
document.ticket_id == @ticket_id AND document.ticket_id == @ticket_id AND
document.status != 'created' AND (document.status == 'open' OR document.status == 'closed') AND
POSITION(document.assigned_to, @account_id) == true POSITION(document.assigned_to, @account_id) == true
LET ticket_items = ( LET ticket_items = (
@ -197,7 +200,6 @@ class TicketStorage(DatabaseArangodbDocument):
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'description': document.description,
'fqdn': fqdn_items[0] or [], 'fqdn': fqdn_items[0] or [],
'ipv4': ipv4_items[0] or [], 'ipv4': ipv4_items[0] or [],
'ipv6': ipv6_items[0] or [], 'ipv6': ipv6_items[0] or [],
@ -234,6 +236,8 @@ class TicketStorage(DatabaseArangodbDocument):
RETURN a['name'] RETURN a['name']
)[0] )[0]
SORT document.metadata.created_at DESC
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'description': document.description, 'description': document.description,
@ -272,6 +276,8 @@ class TicketStorage(DatabaseArangodbDocument):
FILTER document.metadata.created_by == @account_id FILTER document.metadata.created_by == @account_id
SORT document.metadata.created_at DESC
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'description': document.description, 'description': document.description,
@ -307,7 +313,7 @@ class TicketStorage(DatabaseArangodbDocument):
FOR document IN {self.collection_name} FOR document IN {self.collection_name}
FILTER FILTER
document.status != 'created' AND (document.status == 'open' OR document.status == 'closed') AND
POSITION(document.assigned_to, @account_id) == true POSITION(document.assigned_to, @account_id) == true
LET ticket_items = ( LET ticket_items = (
@ -334,9 +340,10 @@ class TicketStorage(DatabaseArangodbDocument):
LET ipv4_items = (FOR item IN ticket_items FILTER item.genre == 'ipv4' RETURN item.items) LET ipv4_items = (FOR item IN ticket_items FILTER item.genre == 'ipv4' RETURN item.items)
LET ipv6_items = (FOR item IN ticket_items FILTER item.genre == 'ipv6' RETURN item.items) LET ipv6_items = (FOR item IN ticket_items FILTER item.genre == 'ipv6' RETURN item.items)
SORT document.metadata.created_at DESC
RETURN {{ RETURN {{
'ticket_id': document.ticket_id, 'ticket_id': document.ticket_id,
'description': document.description,
'fqdn': fqdn_items[0] or [], 'fqdn': fqdn_items[0] or [],
'ipv4': ipv4_items[0] or [], 'ipv4': ipv4_items[0] or [],
'ipv6': ipv6_items[0] or [], 'ipv6': ipv6_items[0] or [],
@ -394,12 +401,12 @@ class TicketStorage(DatabaseArangodbDocument):
except: except:
raise TicketStorageGetException() raise TicketStorageGetException()
def update_task_list(self, ticket_id: str, task_id: str) -> Cursor | Exception: def update_task_list(self, ticket_id: str, task_ids: list, updated_at: str) -> Cursor | Exception:
""" """
Appends a new task. Appends a new task.
:param ticket_id: ticket identifier. :param ticket_id: ticket identifier.
:param task_id: task identifier. :param task_ids: a list of tasks identifiers.
:return: list of updated rows. :return: list of updated rows.
""" """
@ -407,7 +414,12 @@ class TicketStorage(DatabaseArangodbDocument):
FOR document IN {self.collection_name} FOR document IN {self.collection_name}
FILTER document.ticket_id == @ticket_id FILTER document.ticket_id == @ticket_id
UPDATE document WITH {{ "tasks": PUSH(document.tasks, @task_id) }} IN {self.collection_name} UPDATE document WITH {{
"tasks": APPEND(document.tasks, @task_ids, true),
"metadata": {{
"updated_at": @updated_at
}}
}} IN {self.collection_name}
RETURN NEW RETURN NEW
""" """
@ -417,7 +429,8 @@ class TicketStorage(DatabaseArangodbDocument):
aql, aql,
bind_vars = { bind_vars = {
'ticket_id': ticket_id, 'ticket_id': ticket_id,
'task_id': task_id 'task_ids': task_ids,
'updated_at': updated_at
}, },
count = True count = True
) )

View file

@ -98,6 +98,33 @@ class WhitelistStorage(DatabaseArangodbDocument):
except: except:
raise WhitelistStorageGetException() raise WhitelistStorageGetException()
def get_active(self) -> Cursor | Exception:
"""
Fetches the active whitelist items.
:return: cursor with the requested data.
"""
aql = f"""
LET result = (
FOR document IN {self.collection_name}
FILTER document.is_active == true
COLLECT genre = document.genre INTO groupedDocuments
RETURN {{
'genre': genre,
'values': groupedDocuments[*].document.value
}}
)
RETURN MERGE(FOR r IN result RETURN {{[r.genre]: r.values}})
"""
try:
return self.query(aql)
except:
raise WhitelistStorageGetException()
def exists_by_value(self, value: str) -> Cursor | Exception: def exists_by_value(self, value: str) -> Cursor | Exception:
""" """
Searches for an item. Searches for an item.
@ -109,7 +136,8 @@ class WhitelistStorage(DatabaseArangodbDocument):
aql = f""" aql = f"""
FOR document IN {self.collection_name} FOR document IN {self.collection_name}
FILTER document.value == @value FILTER
document.value == @value
RETURN document RETURN document
""" """