mirror of
https://github.com/fuckpiracyshield/data-storage.git
synced 2024-11-23 05:39:47 +01:00
Various fixes.
This commit is contained in:
parent
5397063bbd
commit
c3c7933d2a
9 changed files with 223 additions and 36 deletions
|
@ -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.
|
||||
|
|
|
@ -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):
|
||||
|
||||
"""
|
||||
|
|
|
@ -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
|
||||
)
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -52,9 +52,7 @@ class ForensicStorage(DatabaseArangodbDocument):
|
|||
|
||||
FILTER document.hash_string == @hash_string
|
||||
|
||||
COLLECT WITH COUNT INTO length
|
||||
|
||||
RETURN length
|
||||
RETURN document
|
||||
"""
|
||||
|
||||
try:
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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
|
||||
)
|
||||
|
|
|
@ -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
|
||||
"""
|
||||
|
|
Loading…
Reference in a new issue