diff --git a/src/piracyshield_data_storage/account/storage.py b/src/piracyshield_data_storage/account/storage.py index b2ab36a..f9f86b5 100644 --- a/src/piracyshield_data_storage/account/storage.py +++ b/src/piracyshield_data_storage/account/storage.py @@ -86,6 +86,33 @@ class AccountStorage(DatabaseArangodbDocument): except: 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: """ Checks if an account with this identifier is in the collection. diff --git a/src/piracyshield_data_storage/blob/drivers/azure.py b/src/piracyshield_data_storage/blob/drivers/azure.py index d2623ac..cd56e14 100644 --- a/src/piracyshield_data_storage/blob/drivers/azure.py +++ b/src/piracyshield_data_storage/blob/drivers/azure.py @@ -31,7 +31,7 @@ class AzureBlobStorage(BlobStorageDriver): def create_container(self): """ Creates the container. - Currently used to mock a blob storage. + Currently used to create a mock blob storage. """ return self.container_client.create_container() @@ -79,7 +79,10 @@ class AzureBlobStorage(BlobStorageDriver): return True - except (ResourceNotFoundError, ResourceExistsError, AzureError): + except ResourceExistsError: + raise AzureBlobStorageAlreadyExistsException() + + except (ResourceNotFoundError, AzureError): raise AzureBlobStorageUploadException() def get_list(self): @@ -107,6 +110,14 @@ class AzureBlobStorage(BlobStorageDriver): except (ResourceNotFoundError, AzureError): raise AzureBlobStorageRemoveException() +class AzureBlobStorageAlreadyExistsException(Exception): + + """ + The blob already exists. + """ + + pass + class AzureBlobStorageUploadException(Exception): """ diff --git a/src/piracyshield_data_storage/blob/storage.py b/src/piracyshield_data_storage/blob/storage.py index 6a30b33..d45ba91 100644 --- a/src/piracyshield_data_storage/blob/storage.py +++ b/src/piracyshield_data_storage/blob/storage.py @@ -12,19 +12,7 @@ class BlobStorage: self.driver = driver def upload(self, blob_name: str, file_path: str): - try: - return self.driver.upload( - blob_name = blob_name, - file_path = file_path - ) - - except: - raise BlobStorageUploadException() - -class BlobStorageUploadException(Exception): - - """ - Cannot upload the file. - """ - - pass + return self.driver.upload( + blob_name = blob_name, + file_path = file_path + ) diff --git a/src/piracyshield_data_storage/database/arangodb/document.py b/src/piracyshield_data_storage/database/arangodb/document.py index 602b4ef..37c33fa 100644 --- a/src/piracyshield_data_storage/database/arangodb/document.py +++ b/src/piracyshield_data_storage/database/arangodb/document.py @@ -4,6 +4,8 @@ from arango.exceptions import AQLQueryExecuteError class DatabaseArangodbDocument(DatabaseArangodbConnection): + max_batch_size = 50000 + def collection(self, collection): try: return self.instance.collection(collection) @@ -13,7 +15,7 @@ class DatabaseArangodbDocument(DatabaseArangodbConnection): def query(self, aql, **kwargs): try: - return self.instance.aql.execute(aql, **kwargs) + return self.instance.aql.execute(aql, batch_size = self.max_batch_size, **kwargs) except AQLQueryExecuteError: raise DatabaseArangodbQueryException() diff --git a/src/piracyshield_data_storage/dda/storage.py b/src/piracyshield_data_storage/dda/storage.py index 8b28001..f4f653e 100644 --- a/src/piracyshield_data_storage/dda/storage.py +++ b/src/piracyshield_data_storage/dda/storage.py @@ -27,6 +27,75 @@ class DDAStorage(DatabaseArangodbDocument): except: 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: """ Fetches all the DDA instances. diff --git a/src/piracyshield_data_storage/forensic/storage.py b/src/piracyshield_data_storage/forensic/storage.py index 19475d2..81d2936 100644 --- a/src/piracyshield_data_storage/forensic/storage.py +++ b/src/piracyshield_data_storage/forensic/storage.py @@ -52,9 +52,7 @@ class ForensicStorage(DatabaseArangodbDocument): FILTER document.hash_string == @hash_string - COLLECT WITH COUNT INTO length - - RETURN length + RETURN document """ try: diff --git a/src/piracyshield_data_storage/ticket/item/storage.py b/src/piracyshield_data_storage/ticket/item/storage.py index e4ecd5c..8f9079e 100644 --- a/src/piracyshield_data_storage/ticket/item/storage.py +++ b/src/piracyshield_data_storage/ticket/item/storage.py @@ -29,6 +29,20 @@ class TicketItemStorage(DatabaseArangodbDocument): except: 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: """ Gets all ticket items. Values only. @@ -76,7 +90,7 @@ class TicketItemStorage(DatabaseArangodbDocument): FOR parent_ticket IN {self.ticket_collection_name} FILTER 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 """ @@ -179,7 +193,7 @@ class TicketItemStorage(DatabaseArangodbDocument): FOR parent_ticket IN {self.ticket_collection_name} FILTER 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 """ @@ -404,7 +418,13 @@ class TicketItemStorage(DatabaseArangodbDocument): document.is_whitelisted == false AND 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, 'status': document.status, 'timestamp': document.timestamp, @@ -510,6 +530,37 @@ class TicketItemStorage(DatabaseArangodbDocument): except: 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: """ Searches for a duplicate. @@ -582,7 +633,7 @@ class TicketItemStorage(DatabaseArangodbDocument): FOR parent_ticket IN {self.ticket_collection_name} FILTER 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 {{ status: @status, diff --git a/src/piracyshield_data_storage/ticket/storage.py b/src/piracyshield_data_storage/ticket/storage.py index c5998f5..3e6a5c1 100644 --- a/src/piracyshield_data_storage/ticket/storage.py +++ b/src/piracyshield_data_storage/ticket/storage.py @@ -63,6 +63,7 @@ class TicketStorage(DatabaseArangodbDocument): RETURN {{ 'ticket_id': document.ticket_id, + 'dda_id': document.dda_id, 'description': document.description, 'fqdn': document.fqdn, 'ipv4': document.ipv4, @@ -131,6 +132,7 @@ class TicketStorage(DatabaseArangodbDocument): RETURN {{ 'ticket_id': document.ticket_id, + 'dda_id': document.dda_id, 'description': document.description, 'fqdn': document.fqdn, 'ipv4': document.ipv4, @@ -141,7 +143,8 @@ class TicketStorage(DatabaseArangodbDocument): }}, 'settings': {{ '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 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 LET ticket_items = ( @@ -197,7 +200,6 @@ class TicketStorage(DatabaseArangodbDocument): RETURN {{ 'ticket_id': document.ticket_id, - 'description': document.description, 'fqdn': fqdn_items[0] or [], 'ipv4': ipv4_items[0] or [], 'ipv6': ipv6_items[0] or [], @@ -234,6 +236,8 @@ class TicketStorage(DatabaseArangodbDocument): RETURN a['name'] )[0] + SORT document.metadata.created_at DESC + RETURN {{ 'ticket_id': document.ticket_id, 'description': document.description, @@ -272,6 +276,8 @@ class TicketStorage(DatabaseArangodbDocument): FILTER document.metadata.created_by == @account_id + SORT document.metadata.created_at DESC + RETURN {{ 'ticket_id': document.ticket_id, 'description': document.description, @@ -307,7 +313,7 @@ class TicketStorage(DatabaseArangodbDocument): FOR document IN {self.collection_name} FILTER - document.status != 'created' AND + (document.status == 'open' OR document.status == 'closed') AND POSITION(document.assigned_to, @account_id) == true 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 ipv6_items = (FOR item IN ticket_items FILTER item.genre == 'ipv6' RETURN item.items) + SORT document.metadata.created_at DESC + RETURN {{ 'ticket_id': document.ticket_id, - 'description': document.description, 'fqdn': fqdn_items[0] or [], 'ipv4': ipv4_items[0] or [], 'ipv6': ipv6_items[0] or [], @@ -394,12 +401,12 @@ class TicketStorage(DatabaseArangodbDocument): except: 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. :param ticket_id: ticket identifier. - :param task_id: task identifier. + :param task_ids: a list of tasks identifiers. :return: list of updated rows. """ @@ -407,7 +414,12 @@ class TicketStorage(DatabaseArangodbDocument): FOR document IN {self.collection_name} 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 """ @@ -417,7 +429,8 @@ class TicketStorage(DatabaseArangodbDocument): aql, bind_vars = { 'ticket_id': ticket_id, - 'task_id': task_id + 'task_ids': task_ids, + 'updated_at': updated_at }, count = True ) diff --git a/src/piracyshield_data_storage/whitelist/storage.py b/src/piracyshield_data_storage/whitelist/storage.py index 19fe7e5..50bd652 100644 --- a/src/piracyshield_data_storage/whitelist/storage.py +++ b/src/piracyshield_data_storage/whitelist/storage.py @@ -98,6 +98,33 @@ class WhitelistStorage(DatabaseArangodbDocument): except: 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: """ Searches for an item. @@ -109,7 +136,8 @@ class WhitelistStorage(DatabaseArangodbDocument): aql = f""" FOR document IN {self.collection_name} - FILTER document.value == @value + FILTER + document.value == @value RETURN document """