mirror/data-storage
1
0
Fork 0
mirror of https://github.com/fuckpiracyshield/data-storage.git synced 2024-11-23 05:39:47 +01:00
Code Issues Projects Releases Packages Wiki Activity

Various fixes.

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

27
src/piracyshield_data_storage/account/storage.py
View file

@ -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.

15
src/piracyshield_data_storage/blob/drivers/azure.py
View file

@ -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
src/piracyshield_data_storage/blob/storage.py
View file

@ -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

4
src/piracyshield_data_storage/database/arangodb/document.py
View file

@ -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()

69
src/piracyshield_data_storage/dda/storage.py
View file

@ -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.

4
src/piracyshield_data_storage/forensic/storage.py
View file

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

59
src/piracyshield_data_storage/ticket/item/storage.py
View file

@ -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,

31
src/piracyshield_data_storage/ticket/storage.py
View file

@ -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
)

30
src/piracyshield_data_storage/whitelist/storage.py
View file

@ -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
"""