From 4acfb7e91c0bbc196c72928f2f303990c6f0cf3e Mon Sep 17 00:00:00 2001 From: Daniele Maglie Date: Fri, 19 Jan 2024 16:01:41 +0100 Subject: [PATCH] Initial commit. --- .gitignore | 6 + README.md | 3 + pyproject.toml | 5 + setup.cfg | 16 + src/piracyshield_data_storage/__init__.py | 1 + .../account/__init__.py | 1 + .../account/general/__init__.py | 1 + .../account/general/storage.py | 86 ++ .../account/storage.py | 284 +++++++ .../authentication/__init__.py | 1 + .../authentication/storage.py | 42 + src/piracyshield_data_storage/base.py | 16 + .../blob/__init__.py | 1 + src/piracyshield_data_storage/blob/driver.py | 7 + .../blob/drivers/__init__.py | 1 + .../blob/drivers/azure.py | 124 +++ src/piracyshield_data_storage/blob/storage.py | 30 + .../cache/__init__.py | 1 + .../cache/storage.py | 77 ++ .../database/__init__.py | 1 + .../database/arangodb/__init__.py | 1 + .../database/arangodb/connection.py | 62 ++ .../database/arangodb/document.py | 27 + .../database/redis/__init__.py | 1 + .../database/redis/connection.py | 31 + .../database/redis/document.py | 34 + src/piracyshield_data_storage/dda/__init__.py | 1 + src/piracyshield_data_storage/dda/storage.py | 253 ++++++ .../forensic/__init__.py | 1 + .../forensic/storage.py | 268 ++++++ .../guest/__init__.py | 1 + .../guest/storage.py | 10 + .../internal/__init__.py | 1 + .../internal/storage.py | 10 + src/piracyshield_data_storage/log/__init__.py | 1 + src/piracyshield_data_storage/log/storage.py | 114 +++ .../log/ticket/__init__.py | 1 + .../log/ticket/item/__init__.py | 1 + .../log/ticket/item/storage.py | 114 +++ .../log/ticket/storage.py | 114 +++ .../provider/__init__.py | 1 + .../provider/storage.py | 10 + .../reporter/__init__.py | 1 + .../reporter/storage.py | 10 + .../security/__init__.py | 1 + .../security/anti_brute_force/__init__.py | 1 + .../security/anti_brute_force/memory.py | 98 +++ .../security/blacklist/__init__.py | 1 + .../security/blacklist/memory.py | 82 ++ .../ticket/__init__.py | 1 + .../ticket/error/__init__.py | 1 + .../ticket/error/storage.py | 154 ++++ .../ticket/item/__init__.py | 1 + .../ticket/item/storage.py | 795 ++++++++++++++++++ .../ticket/storage.py | 529 ++++++++++++ .../whitelist/__init__.py | 1 + .../whitelist/storage.py | 227 +++++ 57 files changed, 3664 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 setup.cfg create mode 100644 src/piracyshield_data_storage/__init__.py create mode 100644 src/piracyshield_data_storage/account/__init__.py create mode 100644 src/piracyshield_data_storage/account/general/__init__.py create mode 100644 src/piracyshield_data_storage/account/general/storage.py create mode 100644 src/piracyshield_data_storage/account/storage.py create mode 100644 src/piracyshield_data_storage/authentication/__init__.py create mode 100644 src/piracyshield_data_storage/authentication/storage.py create mode 100644 src/piracyshield_data_storage/base.py create mode 100644 src/piracyshield_data_storage/blob/__init__.py create mode 100644 src/piracyshield_data_storage/blob/driver.py create mode 100644 src/piracyshield_data_storage/blob/drivers/__init__.py create mode 100644 src/piracyshield_data_storage/blob/drivers/azure.py create mode 100644 src/piracyshield_data_storage/blob/storage.py create mode 100644 src/piracyshield_data_storage/cache/__init__.py create mode 100644 src/piracyshield_data_storage/cache/storage.py create mode 100644 src/piracyshield_data_storage/database/__init__.py create mode 100644 src/piracyshield_data_storage/database/arangodb/__init__.py create mode 100644 src/piracyshield_data_storage/database/arangodb/connection.py create mode 100644 src/piracyshield_data_storage/database/arangodb/document.py create mode 100644 src/piracyshield_data_storage/database/redis/__init__.py create mode 100644 src/piracyshield_data_storage/database/redis/connection.py create mode 100644 src/piracyshield_data_storage/database/redis/document.py create mode 100644 src/piracyshield_data_storage/dda/__init__.py create mode 100644 src/piracyshield_data_storage/dda/storage.py create mode 100644 src/piracyshield_data_storage/forensic/__init__.py create mode 100644 src/piracyshield_data_storage/forensic/storage.py create mode 100644 src/piracyshield_data_storage/guest/__init__.py create mode 100644 src/piracyshield_data_storage/guest/storage.py create mode 100644 src/piracyshield_data_storage/internal/__init__.py create mode 100644 src/piracyshield_data_storage/internal/storage.py create mode 100644 src/piracyshield_data_storage/log/__init__.py create mode 100644 src/piracyshield_data_storage/log/storage.py create mode 100644 src/piracyshield_data_storage/log/ticket/__init__.py create mode 100644 src/piracyshield_data_storage/log/ticket/item/__init__.py create mode 100644 src/piracyshield_data_storage/log/ticket/item/storage.py create mode 100644 src/piracyshield_data_storage/log/ticket/storage.py create mode 100644 src/piracyshield_data_storage/provider/__init__.py create mode 100644 src/piracyshield_data_storage/provider/storage.py create mode 100644 src/piracyshield_data_storage/reporter/__init__.py create mode 100644 src/piracyshield_data_storage/reporter/storage.py create mode 100644 src/piracyshield_data_storage/security/__init__.py create mode 100644 src/piracyshield_data_storage/security/anti_brute_force/__init__.py create mode 100644 src/piracyshield_data_storage/security/anti_brute_force/memory.py create mode 100644 src/piracyshield_data_storage/security/blacklist/__init__.py create mode 100644 src/piracyshield_data_storage/security/blacklist/memory.py create mode 100644 src/piracyshield_data_storage/ticket/__init__.py create mode 100644 src/piracyshield_data_storage/ticket/error/__init__.py create mode 100644 src/piracyshield_data_storage/ticket/error/storage.py create mode 100644 src/piracyshield_data_storage/ticket/item/__init__.py create mode 100644 src/piracyshield_data_storage/ticket/item/storage.py create mode 100644 src/piracyshield_data_storage/ticket/storage.py create mode 100644 src/piracyshield_data_storage/whitelist/__init__.py create mode 100644 src/piracyshield_data_storage/whitelist/storage.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3041614 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +__pycache__/ +build/ +eggs/ +.eggs/ +*.egg +*.egg-info/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..cf56a5c --- /dev/null +++ b/README.md @@ -0,0 +1,3 @@ +### Data Storage + +Storage and filesystem management. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..1a1ed96 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,5 @@ +[build-system] +requires = [ + "setuptools>=54", +] +build-backend = "setuptools.build_meta" diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..6a61c8a --- /dev/null +++ b/setup.cfg @@ -0,0 +1,16 @@ +[metadata] +name = piracyshield_data_storage +version = 1.0.0 +description = Data Storage + +[options] +package_dir= + =src +packages = find: +python_requires = >= 3.10 +install_requires = + python-arango + redis + +[options.packages.find] +where = src diff --git a/src/piracyshield_data_storage/__init__.py b/src/piracyshield_data_storage/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/account/__init__.py b/src/piracyshield_data_storage/account/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/account/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/account/general/__init__.py b/src/piracyshield_data_storage/account/general/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/account/general/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/account/general/storage.py b/src/piracyshield_data_storage/account/general/storage.py new file mode 100644 index 0000000..62d981b --- /dev/null +++ b/src/piracyshield_data_storage/account/general/storage.py @@ -0,0 +1,86 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class GeneralAccountStorage(DatabaseArangodbDocument): + + view_name = 'accounts_view' + + def __init__(self): + super().__init__() + + def get(self, account_id: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.view_name} + + FILTER document.account_id == @account_id + + // resolve account identifier + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a['name'] + )[0] + + RETURN {{ + 'account_id': document.account_id, + 'name': document.name, + 'email': document.email, + 'role': document.role, + 'is_active': document.is_active, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise GeneralAccountStorageGetException() + + def get_all(self) -> Cursor | Exception: + aql = f""" + FOR document IN {self.view_name} + + // resolve account identifier + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a['name'] + )[0] + + RETURN {{ + 'account_id': document.account_id, + 'name': document.name, + 'email': document.email, + 'role': document.role, + 'is_active': document.is_active, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }} + }} + """ + + try: + return self.query(aql) + + except: + raise GeneralAccountStorageGetException() + +class GeneralAccountStorageGetException(Exception): + + """ + Cannot get the account. + """ + + pass diff --git a/src/piracyshield_data_storage/account/storage.py b/src/piracyshield_data_storage/account/storage.py new file mode 100644 index 0000000..b2ab36a --- /dev/null +++ b/src/piracyshield_data_storage/account/storage.py @@ -0,0 +1,284 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class AccountStorage(DatabaseArangodbDocument): + + collection_name = None + + collection_instance = None + + def __init__(self, collection_name: str): + super().__init__() + + self.collection_name = collection_name + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + try: + return self.collection_instance.insert(document) + + except: + raise AccountStorageCreateException() + + def get(self, account_id: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + FILTER document.account_id == @account_id + RETURN {{ + 'account_id': document.account_id, + 'name': document.name, + 'email': document.email, + 'role': document.role, + 'is_active': document.is_active, + 'created_by': document.created_by + }} + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise AccountStorageGetException() + + def get_complete(self, account_id: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise AccountStorageGetException() + + def get_all(self) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + RETURN {{ + 'account_id': document.account_id, + 'name': document.name, + 'email': document.email, + 'is_active': document.is_active, + '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. + + :param value: a valid account identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @identifier + + LIMIT 1 + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'identifier': identifier + }) + + except: + raise AccountStorageGetException() + + def set_flag(self, account_id: str, flag: str, value: any) -> Cursor | Exception: + """ + Sets a flag with a status. + + :param account_id: account identifier. + :param flag: the flag to update. + :param value: any value. + :return: the number of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + + UPDATE {{ + '_key': document._key, + 'flags': {{ + @flag: @value + }} + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'account_id': account_id, + 'flag': flag, + 'value': value + }, + count = True + ) + + return affected_rows + + except: + raise AccountStorageUpdateException() + + def change_password(self, account_id: str, password: str) -> Cursor | Exception: + """ + Changes the accounts' password. + + :param account_id: account identifier. + :param password: the encoded password. + :return: the number of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + + UPDATE {{ + '_key': document._key, + 'password': @password + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'account_id': account_id, + 'password': password + }, + count = True + ) + + return affected_rows + + except: + raise AccountStorageUpdateException() + + def update_status(self, account_id: str, value: bool) -> Cursor | Exception: + """ + Updates the account status. + + :param account_id: account identifier. + :param is_active: account status. + :return: the number of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + + UPDATE {{ + '_key': document._key, + 'is_active': @value + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'account_id': account_id, + 'value': value + }, + count = True + ) + + return affected_rows + + except: + raise AccountStorageUpdateException() + + def remove(self, account_id: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'account_id': account_id + }, + count = True + ) + + return affected_rows + + except: + raise AccountStorageRemoveException() + +class AccountStorageCreateException(Exception): + + """ + Cannot create the account. + """ + + pass + +class AccountStorageGetException(Exception): + + """ + Cannot get the account. + """ + + pass + +class AccountStorageUpdateException(Exception): + + """ + Cannot update the account. + """ + + pass + +class AccountStorageRemoveException(Exception): + + """ + Cannot remove the account. + """ + + pass diff --git a/src/piracyshield_data_storage/authentication/__init__.py b/src/piracyshield_data_storage/authentication/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/authentication/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/authentication/storage.py b/src/piracyshield_data_storage/authentication/storage.py new file mode 100644 index 0000000..ef0de28 --- /dev/null +++ b/src/piracyshield_data_storage/authentication/storage.py @@ -0,0 +1,42 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class AuthenticationStorage(DatabaseArangodbDocument): + + view_name = 'accounts_view' + + def __init__(self): + super().__init__() + + def get(self, email: str) -> Cursor | Exception: + """ + Returns the account via email. + + :param email: a valid e-mail address. + :return: cursor with the requested data. + """ + + try: + aql = f""" + FOR document IN {self.view_name} + + FILTER document.email == @email + + RETURN document + """ + + return self.query(aql, bind_vars = { + 'email': email + }) + + except: + raise AuthenticationStorageGetException() + +class AuthenticationStorageGetException(Exception): + + """ + Cannot get the account email. + """ + + pass diff --git a/src/piracyshield_data_storage/base.py b/src/piracyshield_data_storage/base.py new file mode 100644 index 0000000..85d33e1 --- /dev/null +++ b/src/piracyshield_data_storage/base.py @@ -0,0 +1,16 @@ +import time + +class BaseStorage(DatabaseArangodbDocument): + + start_counter = 0 + + stop_counter = 0 + + def __init__(self): + super().__init__() + + def _start_counter(self): + self.start_counter = time.time() + + def _stop_counter(self): + self.stop_counter = time.time() diff --git a/src/piracyshield_data_storage/blob/__init__.py b/src/piracyshield_data_storage/blob/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/blob/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/blob/driver.py b/src/piracyshield_data_storage/blob/driver.py new file mode 100644 index 0000000..7c36e1f --- /dev/null +++ b/src/piracyshield_data_storage/blob/driver.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + +class BlobStorageDriver(ABC): + + @abstractmethod + def upload(self, file_path: str): + pass diff --git a/src/piracyshield_data_storage/blob/drivers/__init__.py b/src/piracyshield_data_storage/blob/drivers/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/blob/drivers/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/blob/drivers/azure.py b/src/piracyshield_data_storage/blob/drivers/azure.py new file mode 100644 index 0000000..d2623ac --- /dev/null +++ b/src/piracyshield_data_storage/blob/drivers/azure.py @@ -0,0 +1,124 @@ +from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient, BlobType + +from azure.core.exceptions import ResourceNotFoundError, ResourceExistsError, AzureError + +from piracyshield_data_storage.blob.driver import BlobStorageDriver + +class AzureBlobStorage(BlobStorageDriver): + + """ + Azure blob storage manage. + """ + + def __init__(self, connection_string: str, container_name: str): + """ + Initializes the client expecting a connection string. + The connection string exploded looks something like: + DefaultEndpointsProtocol=http; + AccountName=ACCOUNT_NAME; + AccountKey=ACCOUNT_KEY; + BlobEndpoint=http://ACCOUNT_NAME.blob.localhost:10000; + QueueEndpoint=http://ACCOUNT_NAME.queue.localhost:10001; + TableEndpoint=http://ACCOUNT_NAME.table.localhost:10002; + + :param connection_string: Azure connection string (https://learn.microsoft.com/en-us/azure/storage/common/storage-configure-connection-string). + :param container_name: container name in use. + """ + self.blob_service_client = BlobServiceClient.from_connection_string(connection_string) + + self.container_client = self.blob_service_client.get_container_client(container_name) + + def create_container(self): + """ + Creates the container. + Currently used to mock a blob storage. + """ + + return self.container_client.create_container() + + def upload(self, blob_name: str, file_path: str) -> bool | Exception: + """ + Uploads a new file in chunks to handle big files safely. + + :param blob_name: name of the file to store. + :param data: file content. + :return bool if correct. + """ + + try: + blob_client = self.container_client.get_blob_client(blob_name) + + blob_client.upload_blob("", blob_type = BlobType.BlockBlob) + + # upload in small chunks + with open(file_path, "rb") as file: + block_size = 4 * 1024 * 1024 # 4 MB chunks + + block_ids = [] + + block_id_prefix = "block" + + block_number = 0 + + while True: + data = file.read(block_size) + + if not data: + break + + block_id = block_id_prefix + str(block_number).zfill(6) + + blob_client.stage_block(block_id, data) + + block_ids.append(block_id) + + block_number += 1 + + # commit the blocks and create the blob + blob_client.commit_block_list(block_ids) + + return True + + except (ResourceNotFoundError, ResourceExistsError, AzureError): + raise AzureBlobStorageUploadException() + + def get_list(self): + """ + Returns a list of blobs in the container. + + :return a paginated list of the files. + """ + + return self.container_client.list_blobs() + + def remove(self, blob_name: str) -> bool | Exception: + """ + Removes the blob from the container. + + :param blob_name: the file name to be removed. + :return bool if correct. + """ + + try: + blob_client = self.container_client.get_blob_client(blob_name) + + blob_client.delete_blob() + + except (ResourceNotFoundError, AzureError): + raise AzureBlobStorageRemoveException() + +class AzureBlobStorageUploadException(Exception): + + """ + The blob cannot be uploaded. + """ + + pass + +class AzureBlobStorageRemoveException(Exception): + + """ + The blob cannot be removed. + """ + + pass diff --git a/src/piracyshield_data_storage/blob/storage.py b/src/piracyshield_data_storage/blob/storage.py new file mode 100644 index 0000000..6a30b33 --- /dev/null +++ b/src/piracyshield_data_storage/blob/storage.py @@ -0,0 +1,30 @@ +from piracyshield_data_storage.blob.driver import BlobStorageDriver + +class BlobStorage: + + """ + Blob storage manager. + """ + + driver = None + + def __init__(self, driver: BlobStorageDriver): + 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 diff --git a/src/piracyshield_data_storage/cache/__init__.py b/src/piracyshield_data_storage/cache/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/cache/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/cache/storage.py b/src/piracyshield_data_storage/cache/storage.py new file mode 100644 index 0000000..20226d9 --- /dev/null +++ b/src/piracyshield_data_storage/cache/storage.py @@ -0,0 +1,77 @@ +from piracyshield_component.environment import Environment +from piracyshield_component.log.logger import Logger + +import os + +# TODO: use custom exceptions. + +class CacheStorage: + + logger = None + + def __init__(self): + self.logger = Logger('storage') + + if not os.path.exists(Environment.CACHE_PATH): + raise FileNotFoundError(f'The specified folder `{Environment.CACHE_PATH}` does not exist') + + if not os.path.isdir(Environment.CACHE_PATH): + raise NotADirectoryError(f'The specified path `{Environment.CACHE_PATH}` is not a directory') + + def write(self, filename: str, content: bytes) -> str | Exception: + path = self._get_absolute_path(filename) + + try: + with open(path, 'wb') as handle: + handle.write(content) + + except IOError: + raise IOError(f'Failed to write content to file `{filename}`') + + # return the absolute path + return path + + def get(self, filename: str) -> str | Exception: + return self._get_absolute_path(filename) + + def read(self, filename: str) -> bytes | Exception: + path = self._get_absolute_path(filename) + + try: + with open(path, 'r') as handle: + return handle.read() + + except FileNotFoundError: + raise FileNotFoundError(f'The specified file `{filename}` does not exist') + + except IOError: + raise IOError(f'Failed to read content from file `{filename}`') + + def exists(self, filename: str) -> bool: + path = self._get_absolute_path(filename) + + return os.path.exists(path) + + def get_all(self) -> list | Exception: + try: + return os.listdir(Environment.CACHE_PATH) + + except IOError: + raise IOError(f'Failed to get files from folder `{Environment.CACHE_PATH}`') + + def remove(self, filename: str) -> bool | Exception: + path = self._get_absolute_path(filename) + + try: + os.remove(path) + + except FileNotFoundError: + raise FileNotFoundError(f'The specified file `{filename}` does not exist') + + except IOError: + raise IOError(f'Failed to remove file `{filename}`.') + + return True + + def _get_absolute_path(self, filename: str) -> str: + return os.path.join(Environment.CACHE_PATH, filename) diff --git a/src/piracyshield_data_storage/database/__init__.py b/src/piracyshield_data_storage/database/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/database/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/database/arangodb/__init__.py b/src/piracyshield_data_storage/database/arangodb/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/database/arangodb/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/database/arangodb/connection.py b/src/piracyshield_data_storage/database/arangodb/connection.py new file mode 100644 index 0000000..7eb7ce8 --- /dev/null +++ b/src/piracyshield_data_storage/database/arangodb/connection.py @@ -0,0 +1,62 @@ +from piracyshield_component.config import Config + +from arango import ArangoClient + +class DatabaseArangodbConnection: + + instance = None + + database_config = None + + def __init__(self, as_root = False): + self._prepare_configs() + + self._prepare_settings() + + self._prepare_credentials(as_root) + + self.establish() + + def establish(self): + client = ArangoClient(hosts = f'{self.protocol}://{self.host}:{self.port}') + + self.instance = client.db(self.database, self.username, self.password, self.verify) + + def _prepare_settings(self): + connection = self.database_config.get('connection') + + try: + self.protocol = connection['protocol'] + + self.host = connection['host'] + + self.port = connection['port'] + + self.verify = connection['verify'] + + except KeyError: + DatabaseArangodbConnectionException('Cannot find the database settings') + + try: + self.database = connection['database'] + + except KeyError: + DatabaseArangodbConnectionException('Cannot find the database name') + + def _prepare_credentials(self, as_root): + credentials = self.database_config.get('root_credentials') if as_root else self.database_config.get('user_credentials') + + try: + self.username = credentials['username'] + + self.password = credentials['password'] + + except KeyError: + DatabaseArangodbConnectionException('Cannot find the database credentials') + + def _prepare_configs(self): + self.database_config = Config('database/arangodb') + +class DatabaseArangodbConnectionException(Exception): + + pass diff --git a/src/piracyshield_data_storage/database/arangodb/document.py b/src/piracyshield_data_storage/database/arangodb/document.py new file mode 100644 index 0000000..602b4ef --- /dev/null +++ b/src/piracyshield_data_storage/database/arangodb/document.py @@ -0,0 +1,27 @@ +from piracyshield_data_storage.database.arangodb.connection import DatabaseArangodbConnection + +from arango.exceptions import AQLQueryExecuteError + +class DatabaseArangodbDocument(DatabaseArangodbConnection): + + def collection(self, collection): + try: + return self.instance.collection(collection) + + except: + raise DatabaseArangodbCollectionNotFoundException() + + def query(self, aql, **kwargs): + try: + return self.instance.aql.execute(aql, **kwargs) + + except AQLQueryExecuteError: + raise DatabaseArangodbQueryException() + +class DatabaseArangodbCollectionNotFoundException(Exception): + + pass + +class DatabaseArangodbQueryException(Exception): + + pass diff --git a/src/piracyshield_data_storage/database/redis/__init__.py b/src/piracyshield_data_storage/database/redis/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/database/redis/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/database/redis/connection.py b/src/piracyshield_data_storage/database/redis/connection.py new file mode 100644 index 0000000..5ac7f04 --- /dev/null +++ b/src/piracyshield_data_storage/database/redis/connection.py @@ -0,0 +1,31 @@ +from piracyshield_component.config import Config + +from redis import Redis + +class DatabaseRedisConnection: + + instance = None + + database_config = None + + def __init__(self) -> None: + self._prepare_configs() + + def establish(self, database: str) -> None: + try: + self.instance = Redis( + host = self.database_config['host'], + port = self.database_config['port'], + db = database, + decode_responses = True + ) + + except: + raise DatabaseRedisConnectionException() + + def _prepare_configs(self) -> None: + self.database_config = Config('database/redis').get('connection') + +class DatabaseRedisConnectionException(Exception): + + pass diff --git a/src/piracyshield_data_storage/database/redis/document.py b/src/piracyshield_data_storage/database/redis/document.py new file mode 100644 index 0000000..8f4e8e9 --- /dev/null +++ b/src/piracyshield_data_storage/database/redis/document.py @@ -0,0 +1,34 @@ +from piracyshield_data_storage.database.redis.connection import DatabaseRedisConnection + +class DatabaseRedisDocument(DatabaseRedisConnection): + + def set_with_expiry(self, key: str, value: any, expiry: int) -> bool | Exception: + if self.instance.set(key, value, ex = expiry) == True: + return True + + raise DatabaseRedisSetException() + + def incr(self, key: str, amount: int = 1) -> bool | Exception: + return self.instance.incr(name = key, amount = amount) + + def get(self, key: str) -> any: + return self.instance.get(key) + + def delete(self, key: str) -> any: + return self.instance.delete(key) + +class DatabaseRedisSetException(Exception): + + """ + Cannot set the data. + """ + + pass + +class DatabaseRedisGetException(Exception): + + """ + Cannot get the data. + """ + + pass diff --git a/src/piracyshield_data_storage/dda/__init__.py b/src/piracyshield_data_storage/dda/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/dda/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/dda/storage.py b/src/piracyshield_data_storage/dda/storage.py new file mode 100644 index 0000000..8b28001 --- /dev/null +++ b/src/piracyshield_data_storage/dda/storage.py @@ -0,0 +1,253 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class DDAStorage(DatabaseArangodbDocument): + + collection_name = 'dda_instances' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new DDA identifier. + + :param document: dictionary with the expected data model values. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise DDAStorageCreateException() + + def get_global(self) -> Cursor | Exception: + """ + Fetches all the DDA instances. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + // resolve account identifier + LET account_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.account_id + RETURN a['name'] + )[0] + + SORT document.instance ASC + + RETURN {{ + 'dda_id': document.dda_id, + 'description': document.description, + 'instance': document.instance, + 'account_id': document.account_id, + 'account_name': account_name, + 'is_active': document.is_active, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql) + + except: + raise DDAStorageGetException() + + def get_all_by_account(self, account_id: str) -> Cursor | Exception: + """ + Fetches all the DDA instances assigned to an account. + + :param account_id: a valid account identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.account_id == @account_id + + SORT document.instance ASC + + 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 = { + 'account_id': account_id + }) + + except: + raise DDAStorageGetException() + + def exists_by_instance(self, instance: str) -> Cursor | Exception: + """ + Searches for an instance. + + :param instance: a valid DDA instance. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.instance == @instance + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'instance': instance + }) + + except: + raise DDAStorageGetException() + + def is_assigned_to_account(self, dda_id: str, account_id: str) -> Cursor | Exception: + """ + Searches for a DDA identifier assigned to a specified account identifier. + + :param dda_id: a valid DDA identifier. + :param account_id: a valid account 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 == @account_id AND + document.is_active == True + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'dda_id': dda_id, + 'account_id': account_id + }) + + except: + raise DDAStorageGetException() + + def update_status(self, dda_id: str, status: bool) -> list | Exception: + """ + Sets the DDA identifier status. + + :param dda_id: a valid DDA identifier. + :param status: true/false if active/non active. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.dda_id == @dda_id + + UPDATE document WITH {{ + is_active: @status + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'dda_id': dda_id, + 'status': status + }, + count = True + ) + + return affected_rows + + except: + return DDAStorageUpdateException() + + def remove(self, dda_id: str) -> Cursor | Exception: + """ + Removes a DDA identifier. + + :param dda_id: DDA identifier. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.dda_id == @dda_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'dda_id': dda_id + }, + count = True + ) + + return affected_rows + + except: + raise DDAStorageRemoveException() + +class DDAStorageCreateException(Exception): + + """ + Cannot create the DDA item. + """ + + pass + +class DDAStorageGetException(Exception): + + """ + Cannot get the DDA item. + """ + + pass + +class DDAStorageUpdateException(Exception): + + """ + Cannot update the DDA item. + """ + + pass + +class DDAStorageRemoveException(Exception): + + """ + Cannot remove the DDA item. + """ + + pass diff --git a/src/piracyshield_data_storage/forensic/__init__.py b/src/piracyshield_data_storage/forensic/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/forensic/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/forensic/storage.py b/src/piracyshield_data_storage/forensic/storage.py new file mode 100644 index 0000000..19475d2 --- /dev/null +++ b/src/piracyshield_data_storage/forensic/storage.py @@ -0,0 +1,268 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class ForensicStorage(DatabaseArangodbDocument): + + collection_name = 'forensics' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new ticket. + + :param document: dictionary with the expected ticket data model values. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise ForensicStorageCreateException() + + def exists_ticket_id(self, ticket_id: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + COLLECT WITH COUNT INTO length + + RETURN length + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise ForensicStorageGetException() + + def exists_hash_string(self, hash_string: str) -> Cursor | Exception: + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.hash_string == @hash_string + + COLLECT WITH COUNT INTO length + + RETURN length + """ + + try: + return self.query(aql, bind_vars = { + 'hash_string': hash_string + }) + + except: + raise ForensicStorageGetException() + + def get_by_ticket(self, ticket_id: str) -> Cursor | Exception: + """ + Returns forensic document by ticket identifier. + + :param ticket_id: a ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise ForensicStorageGetException() + + def get_by_ticket_for_reporter(self, ticket_id: str, reporter_id: str) -> Cursor | Exception: + """ + Returns forensic document by ticket identifier. + + :param ticket_id: a ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.metadata.created_by == @reporter_id + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'reporter_id': reporter_id + }) + + except: + raise ForensicStorageGetException() + + def update_archive_name(self, ticket_id: str, archive_name: str, status: str, updated_at: str) -> Cursor | Exception: + """ + Insert the forensic evidence archive name. + + :param ticket_id: ticket identifier. + :param archive_name: name of the file archive. + :param status: status of the analysis. + :param updated_at: date of this update. + :return: list of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + UPDATE document WITH {{ + archive_name: @archive_name, + status: @status, + metadata: {{ + updated_at: @updated_at + }} + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'archive_name': archive_name, + 'status': status, + 'updated_at': updated_at + }, + count = True + ) + + return affected_rows + + except: + raise ForensicStorageUpdateException() + + def update_archive_status(self, ticket_id: str, status: str, updated_at: str, reason: str = None) -> Cursor | Exception: + """ + Updates the status of a previously created record. + + :param ticket_id: ticket identifier. + :param status: status of the analysis. + :param reason: additional string message. + :param updated_at: date of this update. + :return: list of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + UPDATE document WITH {{ + status: @status, + reason: @reason, + metadata: {{ + updated_at: @updated_at + }} + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'status': status, + 'reason': reason, + 'updated_at': updated_at + }, + count = True + ) + + return affected_rows + + except: + raise ForensicStorageUpdateException() + + def remove_by_ticket(self, ticket_id: str) -> Cursor | Exception: + """ + Removes a forensic archive. + + :param ticket_id: ticket identifier. + :return: list of removed rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id + }, + count = True + ) + + return affected_rows + + except: + raise TicketStorageRemoveException() + +class ForensicStorageCreateException(Exception): + + """ + Cannot create the forensic archive. + """ + + pass + +class ForensicStorageGetException(Exception): + + """ + Cannot get the forensic archive. + """ + + pass + +class ForensicStorageUpdateException(Exception): + + """ + Cannot update the forensic archive. + """ + + pass + +class ForensicStorageRemoveException(Exception): + + """ + Cannot remove the forensic archive. + """ + + pass diff --git a/src/piracyshield_data_storage/guest/__init__.py b/src/piracyshield_data_storage/guest/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/guest/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/guest/storage.py b/src/piracyshield_data_storage/guest/storage.py new file mode 100644 index 0000000..80d3eda --- /dev/null +++ b/src/piracyshield_data_storage/guest/storage.py @@ -0,0 +1,10 @@ +from piracyshield_data_storage.account.storage import AccountStorage + +class GuestStorage(AccountStorage): + + COLLECTION = 'guests' + + collection_instance = None + + def __init__(self): + super().__init__(collection_name = self.COLLECTION) diff --git a/src/piracyshield_data_storage/internal/__init__.py b/src/piracyshield_data_storage/internal/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/internal/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/internal/storage.py b/src/piracyshield_data_storage/internal/storage.py new file mode 100644 index 0000000..77c0d39 --- /dev/null +++ b/src/piracyshield_data_storage/internal/storage.py @@ -0,0 +1,10 @@ +from piracyshield_data_storage.account.storage import AccountStorage + +class InternalStorage(AccountStorage): + + COLLECTION = 'internals' + + collection_instance = None + + def __init__(self): + super().__init__(collection_name = self.COLLECTION) diff --git a/src/piracyshield_data_storage/log/__init__.py b/src/piracyshield_data_storage/log/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/log/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/log/storage.py b/src/piracyshield_data_storage/log/storage.py new file mode 100644 index 0000000..c941041 --- /dev/null +++ b/src/piracyshield_data_storage/log/storage.py @@ -0,0 +1,114 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class LogStorage(DatabaseArangodbDocument): + + collection_name = 'logs' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict: + """ + Adds a new log record. + + :param document: dictionary with the values to insert. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise LogStorageCreateException() + + def get(self, identifier: str) -> Cursor: + """ + Gets log records for a specific identifier. + + :param identifier: generic identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.identifier == @identifier + + RETURN {{ + 'log_id': document.log_id, + 'time': document.identifier, + 'message': document.message, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'identifier': identifier + }) + + except: + raise LogStorageGetException() + + def remove_by_ticket_id(self, ticket_id: str) -> dict | bool: + """ + Removes all the logs related to a ticket. + + :param ticket_id: ticket identifier. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.identifier == @ticket_id + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id + }, + count = True + ) + + return affected_rows + + except: + raise LogStorageRemoveException() + +class LogStorageCreateException(Exception): + + """ + Cannot create the log. + """ + + pass + +class LogStorageGetException(Exception): + + """ + Cannot get the log. + """ + + pass + +class LogStorageRemoveException(Exception): + + """ + Cannot remove the log. + """ + + pass diff --git a/src/piracyshield_data_storage/log/ticket/__init__.py b/src/piracyshield_data_storage/log/ticket/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/log/ticket/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/log/ticket/item/__init__.py b/src/piracyshield_data_storage/log/ticket/item/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/log/ticket/item/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/log/ticket/item/storage.py b/src/piracyshield_data_storage/log/ticket/item/storage.py new file mode 100644 index 0000000..e009a3a --- /dev/null +++ b/src/piracyshield_data_storage/log/ticket/item/storage.py @@ -0,0 +1,114 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class LogTicketItemStorage(DatabaseArangodbDocument): + + collection_name = 'log_ticket_blocking_items' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new log record. + + :param document: dictionary with the values to insert. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise LogTicketItemStorageCreateException() + + def get_all(self, ticket_item_id: str) -> Cursor | Exception: + """ + Gets all records for a specific ticket item identifier. + + :param ticket_item_id: a valid ticket item identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_item_id == @ticket_item_id + + RETURN {{ + 'ticket_item_id': document.ticket_item_id, + 'message': document.message, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_item_id': ticket_item_id + }) + + except: + raise LogTicketItemStorageGetException() + + def remove_all(self, ticket_item_id: str) -> Cursor | Exception: + """ + Removes all the records related to a ticket item identifier. + + :param ticket_item_id: a valid ticket item identifier. + :return: number of affected records. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_item_id == @ticket_item_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_item_id': ticket_item_id + }, + count = True + ) + + return affected_rows + + except: + raise LogTicketItemStorageRemoveException() + +class LogTicketItemStorageCreateException(Exception): + + """ + Cannot create the log. + """ + + pass + +class LogTicketItemStorageGetException(Exception): + + """ + Cannot get the log. + """ + + pass + +class LogTicketItemStorageRemoveException(Exception): + + """ + Cannot remove the log. + """ + + pass diff --git a/src/piracyshield_data_storage/log/ticket/storage.py b/src/piracyshield_data_storage/log/ticket/storage.py new file mode 100644 index 0000000..1a34304 --- /dev/null +++ b/src/piracyshield_data_storage/log/ticket/storage.py @@ -0,0 +1,114 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class LogTicketStorage(DatabaseArangodbDocument): + + collection_name = 'log_ticket_blockings' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new log record. + + :param document: dictionary with the values to insert. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise LogTicketStorageCreateException() + + def get_all(self, ticket_id: str) -> Cursor | Exception: + """ + Gets all records for a specific ticket identifier. + + :param ticket_id: a valid ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + RETURN {{ + 'ticket_id': document.ticket_id, + 'message': document.message, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise LogTicketStorageGetException() + + def remove_all(self, ticket_id: str) -> Cursor | Exception: + """ + Removes all the records related to a ticket identifier. + + :param ticket_id: a valid ticket identifier. + :return: number of affected records. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id + }, + count = True + ) + + return affected_rows + + except: + raise LogTicketStorageRemoveException() + +class LogTicketStorageCreateException(Exception): + + """ + Cannot create the log. + """ + + pass + +class LogTicketStorageGetException(Exception): + + """ + Cannot get the log. + """ + + pass + +class LogTicketStorageRemoveException(Exception): + + """ + Cannot remove the log. + """ + + pass diff --git a/src/piracyshield_data_storage/provider/__init__.py b/src/piracyshield_data_storage/provider/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/provider/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/provider/storage.py b/src/piracyshield_data_storage/provider/storage.py new file mode 100644 index 0000000..f2102ce --- /dev/null +++ b/src/piracyshield_data_storage/provider/storage.py @@ -0,0 +1,10 @@ +from piracyshield_data_storage.account.storage import AccountStorage + +class ProviderStorage(AccountStorage): + + COLLECTION = 'providers' + + collection_instance = None + + def __init__(self): + super().__init__(collection_name = self.COLLECTION) diff --git a/src/piracyshield_data_storage/reporter/__init__.py b/src/piracyshield_data_storage/reporter/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/reporter/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/reporter/storage.py b/src/piracyshield_data_storage/reporter/storage.py new file mode 100644 index 0000000..a3f96ba --- /dev/null +++ b/src/piracyshield_data_storage/reporter/storage.py @@ -0,0 +1,10 @@ +from piracyshield_data_storage.account.storage import AccountStorage + +class ReporterStorage(AccountStorage): + + COLLECTION = 'reporters' + + collection_instance = None + + def __init__(self): + super().__init__(collection_name = self.COLLECTION) diff --git a/src/piracyshield_data_storage/security/__init__.py b/src/piracyshield_data_storage/security/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/security/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/security/anti_brute_force/__init__.py b/src/piracyshield_data_storage/security/anti_brute_force/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/security/anti_brute_force/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/security/anti_brute_force/memory.py b/src/piracyshield_data_storage/security/anti_brute_force/memory.py new file mode 100644 index 0000000..ce6f882 --- /dev/null +++ b/src/piracyshield_data_storage/security/anti_brute_force/memory.py @@ -0,0 +1,98 @@ +from piracyshield_data_storage.database.redis.document import DatabaseRedisDocument, DatabaseRedisSetException, DatabaseRedisGetException + +class SecurityAntiBruteForceMemory(DatabaseRedisDocument): + + def __init__(self, database: int): + super().__init__() + + self.establish(database) + + def set_login_attempts(self, email: str, timeframe: int, attempts: int = 1) -> bool | Exception: + """ + Sets the current login attempts in a given timeframe. + + :param email: a valid e-mail address. + :param timeframe: a timeframe in seconds. + :param attempts: sets number of attempts, default: 1. + :return: true if the value has been stored. + """ + + try: + return self.set_with_expiry( + key = email, + value = attempts, # store the correct types + expiry = timeframe + ) + + except DatabaseRedisSetException: + raise SecurityMemorySetException() + + def increment_login_attempts(self, email: str, amount: int = 1) -> bool | Exception: + """ + Increments the attempts by preserving expiry time. + + :param email: a valid e-mail address. + :param amount: increment by a number, default: 1. + :return: true if successfully executed. + """ + + try: + return self.incr( + key = email, + amount = amount + ) + + except DatabaseRedisSetException: + raise SecurityMemorySetException() + + def get_login_attempts(self, email: str) -> int | Exception: + """ + Gets the current account login attempts. + + :param email: a valid e-mail address. + :return: the number of attempts. + """ + + try: + response = self.get(key = email) + + if response: + # make sure we're dealing with a true int and not a string + return int(response) + + return response + + except DatabaseRedisGetException: + raise SecurityMemoryGetException() + + def reset_login_attempts(self, email: str) -> bool | Exception: + """ + Unsets the login attempts count. + + :param email: a valid e-mail address. + :return: true if successfully executed. + """ + + try: + return self.delete( + key = email + ) + + except DatabaseRedisSetException: + raise SecurityMemorySetException() + +class SecurityAntiBruteForceMemorySetException(Exception): + + """ + Cannot set the value. + """ + + pass + +class SecurityAntiBruteForceMemoryGetException(Exception): + + """ + Cannot get the value. + """ + + pass diff --git a/src/piracyshield_data_storage/security/blacklist/__init__.py b/src/piracyshield_data_storage/security/blacklist/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/security/blacklist/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/security/blacklist/memory.py b/src/piracyshield_data_storage/security/blacklist/memory.py new file mode 100644 index 0000000..dd9b499 --- /dev/null +++ b/src/piracyshield_data_storage/security/blacklist/memory.py @@ -0,0 +1,82 @@ +from piracyshield_data_storage.database.redis.document import DatabaseRedisDocument, DatabaseRedisSetException, DatabaseRedisGetException + +class SecurityBlacklistMemory(DatabaseRedisDocument): + + ip_address_prefix = 'ip_address' + + token_prefix = 'token' + + def __init__(self, database: int): + super().__init__() + + self.establish(database) + + def add_ip_address(self, ip_address: str, duration: int = 60) -> bool | Exception: + """ + Blacklists an IP address. + + :param ip_address: a valid IP address. + :param duration: duration of the blacklist in seconds. + :return: true if the item has been stored. + """ + + try: + return self.set_with_expiry( + key = f'{self.ip_address_prefix}:{ip_address}', + value = '1', + expiry = duration + ) + + except DatabaseRedisSetException: + raise SecurityBlacklistMemorySetException() + + def exists_by_ip_address(self, ip_address: str) -> bool | Exception: + """ + Verifies if an IP address is in the blacklist. + + :param item: a valid IP address. + :return: returns the TTL of the item. + """ + + try: + response = self.get(key = f'{self.ip_address_prefix}:{ip_address}') + + if response: + return True + + return False + + except DatabaseRedisGetException: + raise SecurityBlacklistMemoryGetException() + + def remove_ip_address(self, ip_address: str) -> bool | Exception: + """ + Removes an IP address from the blacklist. + + :param item: a valid IP address. + :return: true if the item has been removed. + """ + + try: + return self.delete( + key = f'{self.ip_address_prefix}:{ip_address}' + ) + + except DatabaseRedisSetException: + raise SecurityBlacklistMemorySetException() + +class SecurityBlacklistMemorySetException(Exception): + + """ + Cannot set the value. + """ + + pass + +class SecurityBlacklistMemoryGetException(Exception): + + """ + Cannot get the value. + """ + + pass diff --git a/src/piracyshield_data_storage/ticket/__init__.py b/src/piracyshield_data_storage/ticket/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/ticket/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/ticket/error/__init__.py b/src/piracyshield_data_storage/ticket/error/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/ticket/error/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/ticket/error/storage.py b/src/piracyshield_data_storage/ticket/error/storage.py new file mode 100644 index 0000000..18bba6a --- /dev/null +++ b/src/piracyshield_data_storage/ticket/error/storage.py @@ -0,0 +1,154 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class TicketErrorStorage(DatabaseArangodbDocument): + + collection_name = 'ticket_errors' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new error ticket. + + :param document: dictionary with the expected ticket data model values. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise TicketErrorStorageCreateException() + + def get(self, ticket_error_id: str) -> Cursor | Exception: + """ + Gets error ticket by its identifier. + + :param ticket_error_id: a valid ticket error identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_error_id == @ticket_error_id + + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a.name + )[0] + + RETURN {{ + 'ticket_error_id': document.ticket_error_id, + 'ticket_id': document.ticket_id, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_error_id': ticket_error_id + }) + + except: + raise TicketErrorStorageGetException() + + def get_by_reporter(self, ticket_error_id: str, reporter_id: str) -> Cursor | Exception: + """ + Gets error ticket by its identifier and creator account. + + :param ticket_error_id: a valid ticket error identifier. + :param reporter_id: a valid reporter account identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_error_id == @ticket_error_id AND + document.metadata.created_by == @reporter_id + + RETURN {{ + 'ticket_error_id': document.ticket_error_id, + 'ticket_id': document.ticket_id, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_error_id': ticket_error_id, + 'reporter_id': reporter_id + }) + + except: + raise TicketErrorStorageGetException() + + def get_by_ticket(self, ticket_id: str) -> Cursor | Exception: + """ + Gets error tickets by ticket identifier. + + :param ticket_id: a valid ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + RETURN {{ + 'ticket_error_id': document.ticket_error_id, + 'fqdns': document.fqdn or [], + 'ipv4': document.ipv4 or [], + 'ipv6': document.ipv6 or [], + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise TicketErrorStorageGetException() + +class TicketErrorStorageCreateException(Exception): + + """ + Cannot create the error ticket. + """ + + pass + +class TicketErrorStorageGetException(Exception): + + """ + Cannot get the error ticket. + """ + + pass diff --git a/src/piracyshield_data_storage/ticket/item/__init__.py b/src/piracyshield_data_storage/ticket/item/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/ticket/item/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/ticket/item/storage.py b/src/piracyshield_data_storage/ticket/item/storage.py new file mode 100644 index 0000000..e4ecd5c --- /dev/null +++ b/src/piracyshield_data_storage/ticket/item/storage.py @@ -0,0 +1,795 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class TicketItemStorage(DatabaseArangodbDocument): + + ticket_collection_name = 'ticket_blockings' + + collection_name = 'ticket_blocking_items' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new ticket item. + + :param document: dictionary ticket item data model structure. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise TicketItemStorageCreateException() + + def get_all_items_with_genre(self, genre: str) -> Cursor: + """ + Gets all ticket items. Values only. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.genre == @genre AND + document.is_active == true + + RETURN DISTINCT document.value + """ + + try: + return self.query(aql, bind_vars = { + 'genre': genre + }) + + except: + raise TicketItemStorageGetException() + + def get_all_items_with_genre_by_provider(self, genre: str, provider_id: str) -> Cursor | Exception: + """ + Gets all ticket items assigned to this provider. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.genre == @genre AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false AND + document.provider_id == @provider_id + + // 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.ticket != 'created' + + RETURN DISTINCT document.value + """ + + try: + return self.query(aql, bind_vars = { + 'genre': genre, + 'provider_id': provider_id + }) + + except: + raise TicketItemStorageGetException() + + def get_all_items_with_genre_by_ticket(self, ticket_id: str, genre: str) -> Cursor | Exception: + """ + Gets all ticket items. Values only, by ticket_id. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.genre == @genre AND + document.is_active == true + + RETURN DISTINCT document.value + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'genre': genre + }) + + except: + raise TicketItemStorageGetException() + + def get_all_items_with_genre_by_ticket_for_reporter(self, ticket_id: str, genre: str, reporter_id: str) -> Cursor | Exception: + """ + Gets items of a ticket created by a reporter account. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.genre == @genre + + FOR ticket IN {self.ticket_collection_name} + + FILTER + ticket.ticket_id == @ticket_id AND + ticket.metadata.created_by == @reporter_id + + RETURN DISTINCT {{ + 'value': document.value, + 'genre': document.genre, + 'is_duplicate': document.is_duplicate, + 'is_whitelisted': document.is_whitelisted, + 'is_error': document.is_error + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'genre': genre, + 'reporter_id': reporter_id + }) + + except: + raise TicketItemStorageGetException() + + def get_all_items_with_genre_by_ticket_for_provider(self, ticket_id: str, genre: str, provider_id: str) -> Cursor | Exception: + """ + Gets all ticket items. Values only, by ticket_id. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.genre == @genre AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false AND + document.provider_id == @provider_id + + // 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.ticket != 'created' + + RETURN DISTINCT document.value + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'genre': genre, + 'provider_id': provider_id + }) + + except: + raise TicketItemStorageGetException() + + def get_all_items_available_by_ticket(self, ticket_id: str, account_id: str) -> Cursor | Exception: + """ + Gets all the ticket items that aren't a duplicate, whitelisted or errors. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false + + FOR ticket IN {self.ticket_collection_name} + + FILTER + ticket.ticket_id == @ticket_id AND + ticket.metadata.created_by == @account_id + + RETURN DISTINCT document.value + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'account_id': account_id + }) + + except: + raise TicketItemStorageGetException() + + def get(self, ticket_id: str, value: str) -> Cursor | Exception: + """ + Gets a single item for a ticket. + + :param ticket_id: ticket identifier. + :param value: item value. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.value == @value + + RETURN {{ + 'ticket_id': document.ticket_id, + 'value': document.value, + 'genre': document.genre, + 'status': document.status, + 'reason': document.reason, + 'provider_id': document.provider_id, + 'created_at': document.created_at, + 'updated_at': document.updated_at + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'value': value + }) + + except: + raise TicketItemStorageGetException() + + def get_by_value(self, provider_id: str, value: str) -> Cursor | Exception: + """ + Gets a single item by its value. + + :param ticket_id: ticket identifier. + :param value: item value. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.provider_id == @provider_id AND + document.value == @value AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false + + RETURN {{ + 'value': document.value, + 'genre': document.genre, + 'metadata': {{ + 'created_at': document.metadata.created_at + }}, + 'settings': document.settings + }} + """ + + try: + return self.query(aql, bind_vars = { + 'provider_id': provider_id, + 'value': value + }) + + except: + raise TicketItemStorageGetException() + + def get_all_ticket_items_by(self, ticket_id: str, account_id: str, genre: str, status: str) -> Cursor | Exception: + """ + Gets all the items filtered by genre and associated account identifier. + + :param ticket_id: ticket identifier. + :param account_id: provider identifier. + :param genre: genre of the ticket item. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + account_id == @account_id AND + genre == @genre AND + status == @status + + RETURN {{ + 'ticket_id': document.ticket_id, + 'value': document.value, + 'genre': document.genre, + 'status': document.status, + 'reason': document.reason, + 'provider_id': document.provider_id, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'account_id': account_id, + 'genre': genre, + 'status': status + }) + + except: + raise TicketItemStorageGetException() + + def get_all(self, ticket_id: str) -> Cursor | Exception: + """ + Gets all the items for a ticket. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + RETURN {{ + 'ticket_id': document.ticket_id, + 'ticket_item_id': document.ticket_item_id, + 'value': document.value, + 'genre': document.genre, + 'status': document.status, + 'reason': document.reason, + 'provider_id': document.provider_id, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise TicketItemStorageGetException() + + def get_all_by_provider(self, provider_id: str, ticket_id: str) -> Cursor | Exception: + """ + Gets all the items for a ticket by a specific provider_id. + + :param provider_id: the provider_id. + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.provider_id == @provider_id AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false + + RETURN {{ + 'value': document.value, + 'status': document.status, + 'timestamp': document.timestamp, + 'note': document.note, + 'reason': document.reason + }} + """ + + try: + return self.query(aql, bind_vars = { + 'provider_id': provider_id, + 'ticket_id': ticket_id + }) + + except: + raise TicketItemStorageGetException() + + def get_details(self, ticket_id: str, ticket_item_id: str) -> Cursor | Exception: + """ + Gets all the items for a ticket. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.ticket_item_id == @ticket_item_id + + LET provider_details = ( + FOR a IN accounts_view + + FILTER a.account_id == document.provider_id + + RETURN {{ + 'account_id': a.account_id, + 'name': a.name + }} + )[0] + + RETURN {{ + 'ticket_id': document.ticket_id, + 'ticket_item_id': document.ticket_item_id, + 'value': document.value, + 'genre': document.genre, + 'status': document.status, + 'reason': document.reason, + 'is_active': document.is_active, + 'is_duplicate': document.is_duplicate, + 'is_whitelisted': document.is_whitelisted, + 'is_error': document.is_error, + 'provider': {{ + 'account_id': provider_details.account_id, + 'name': provider_details.name + }}, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'ticket_item_id': ticket_item_id + }) + + except: + raise TicketItemStorageGetException() + + def get_all_status(self, ticket_id: str) -> Cursor | Exception: + """ + Gets all ticket items statuses. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + RETURN DISTINCT {{ + 'ticket_item_id': document.ticket_item_id, + 'value': document.value, + 'genre': document.genre, + 'is_active': document.is_active, + 'is_whitelisted': document.is_whitelisted, + 'is_error': document.is_error + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise TicketItemStorageGetException() + + def exists_by_value(self, genre: str, value: str) -> Cursor | Exception: + """ + Searches for a duplicate. + + :param genre: the genre of the ticket item. + :param genre: the value of the ticket item. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.genre == @genre AND + document.value == @value AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false + + RETURN DISTINCT {{ + 'ticket_id': document.ticket_id, + 'genre': document.genre, + 'value': document.value + }} + """ + + try: + return self.query(aql, bind_vars = { + 'genre': genre, + 'value': value + }) + + except: + raise TicketItemStorageGetException() + + def update_status_by_value(self, + provider_id: str, + value: str, + status: str, + updated_at: str, + status_reason: str = None, + timestamp: str = None, + note: str = None + ) -> list | Exception: + """ + Sets the item status by its value. + + :param provider_id: the id of the provider account. + :param value: item value. + :param status: item status value. + :param updated_at: a timestamp of the update date. + :param status_reason: a valid predefined reason for unprocessed items. + :param timestamp: a timestamp of the update date set by the provider account. + :param note: a generic text set by the provider account. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.value == @value AND + document.provider_id == @provider_id AND + document.is_active == true AND + document.is_duplicate == false AND + document.is_whitelisted == false AND + document.is_error == false + + // prevent editing a ticket item with a non workable ticket + FOR parent_ticket IN {self.ticket_collection_name} + FILTER + parent_ticket.ticket_id == document.ticket_id AND + parent_ticket.status.ticket != 'created' + + UPDATE document WITH {{ + status: @status, + reason: @status_reason, + timestamp: @timestamp, + note: @note, + metadata: {{ + updated_at: @updated_at + }} + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'provider_id': provider_id, + 'value': value, + 'status': status, + 'updated_at': updated_at, + 'status_reason': status_reason, + 'timestamp': timestamp, + 'note': note + }, + count = True + ) + + return affected_rows + + except: + raise TicketItemStorageUpdateException() + + def set_flag_active(self, value: str, status: str) -> list | Exception: + """ + Sets the item activity status. + + :param value: item value. + :param status: item status value. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.value == @value + + UPDATE document WITH {{ + is_active: @status + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'value': value, + 'status': status + }, + count = True + ) + + return affected_rows + + except: + raise TicketItemStorageUpdateException() + + def set_flag_error(self, ticket_id: str, value: str, status: bool) -> list | Exception: + """ + Sets the error flag. + + :param ticket_id: a valid ticket identifier. + :param value: item value. + :param status: true or false if status is error or not. + :return: a list of affected rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.value == @value + + UPDATE document WITH {{ + is_error: @status + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'value': value, + 'status': status + }, + count = True + ) + + return affected_rows + + except: + raise TicketItemStorageUpdateException() + + def remove(self, ticket_id: str, value: str) -> list | Exception: + """ + Removes a ticket item. + + :param ticket_id: ticket identifier. + :param value: item to remove. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.value == @value + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'value': value + }, + count = True + ) + + return affected_rows + + except: + raise TicketItemStorageRemoveException() + + def remove_all(self, ticket_id: str) -> list | Exception: + """ + Removes all the ticket items. + + :param ticket_id: ticket identifier. + :param value: item to remove. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id + }, + count = True + ) + + return affected_rows + + except: + raise TicketItemStorageRemoveException() + +class TicketItemStorageCreateException(Exception): + + """ + Cannot create the ticket item. + """ + + pass + +class TicketItemStorageGetException(Exception): + + """ + Cannot get the ticket item. + """ + + pass + +class TicketItemStorageUpdateException(Exception): + + """ + Cannot update the ticket item. + """ + + pass + +class TicketItemStorageRemoveException(Exception): + + """ + Cannot remove the ticket item. + """ + + pass diff --git a/src/piracyshield_data_storage/ticket/storage.py b/src/piracyshield_data_storage/ticket/storage.py new file mode 100644 index 0000000..c5998f5 --- /dev/null +++ b/src/piracyshield_data_storage/ticket/storage.py @@ -0,0 +1,529 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class TicketStorage(DatabaseArangodbDocument): + + ticket_item_collection_name = 'ticket_blocking_items' + + collection_name = 'ticket_blockings' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new ticket. + + :param document: dictionary with the expected ticket data model values. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise TicketStorageCreateException() + + def get(self, ticket_id: str) -> Cursor | Exception: + """ + Gets a single ticket document. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a.name + )[0] + + LET assigned_to_details = ( + FOR i IN document.assigned_to + + FOR a IN accounts_view + + FILTER a.account_id == i + + RETURN {{ + 'account_id': a.account_id, + 'name': a.name + }} + ) + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'assigned_to': assigned_to_details, + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }}, + 'settings': document.settings, + 'tasks': document.tasks + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise TicketStorageGetException() + + def has_dda_id(self, dda_id: str) -> Cursor | Exception: + """ + Finds a ticket with the specified DDA 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 + + LIMIT 1 + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'dda_id': dda_id + }) + + except: + raise TicketStorageGetException() + + def get_reporter(self, ticket_id: str, account_id: str) -> Cursor | Exception: + """ + Gets a single ticket document with limitations. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.metadata.created_by == @account_id + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at + }}, + 'settings': {{ + 'revoke_time': document.settings.revoke_time, + 'autoclose_time': document.settings.autoclose_time + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'account_id': account_id + }) + + except: + raise TicketStorageGetException() + + def get_provider(self, ticket_id: str, account_id: str) -> Cursor | Exception: + """ + Gets a single ticket document with limitations. + + :param ticket_id: ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.ticket_id == @ticket_id AND + document.status != 'created' AND + POSITION(document.assigned_to, @account_id) == true + + LET ticket_items = ( + FOR ticket_item in {self.ticket_item_collection_name} + + FILTER + ticket_item.ticket_id == document.ticket_id AND + ticket_item.is_active == true AND + ticket_item.is_duplicate == false AND + ticket_item.is_whitelisted == false AND + ticket_item.is_error == false + + COLLECT genre = ticket_item.genre INTO groupedItems + FILTER genre IN ["fqdn", "ipv4", "ipv6"] + RETURN {{ + "genre": genre, + "items": UNIQUE(groupedItems[*].ticket_item.value) + }} + ) + + FILTER LENGTH(ticket_items) != 0 + + LET fqdn_items = (FOR item IN ticket_items FILTER item.genre == 'fqdn' 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) + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': fqdn_items[0] or [], + 'ipv4': ipv4_items[0] or [], + 'ipv6': ipv6_items[0] or [], + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id, + 'account_id': account_id + }) + + except: + raise TicketStorageGetException() + + def get_all(self) -> Cursor | Exception: + """ + Fetches all the tickets. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + // resolve account identifier + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a['name'] + )[0] + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'assigned_to': document.assigned_to, + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }}, + 'settings': document.settings, + 'tasks': document.tasks + }} + """ + + try: + return self.query(aql) + + except: + raise TicketStorageGetException() + + def get_all_reporter(self, account_id: str) -> Cursor | Exception: + """ + Fetches all the tickets for a single account. + + :param account_id: the identifier of the ticket's creator. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.metadata.created_by == @account_id + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': document.fqdn, + 'ipv4': document.ipv4, + 'ipv6': document.ipv6, + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at + }}, + 'settings': {{ + 'revoke_time': document.settings.revoke_time + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise TicketStorageGetException() + + def get_all_provider(self, account_id) -> Cursor | Exception: + """ + Fetches all the tickets with limited parameters. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.status != 'created' AND + POSITION(document.assigned_to, @account_id) == true + + LET ticket_items = ( + FOR ticket_item in {self.ticket_item_collection_name} + + FILTER + ticket_item.ticket_id == document.ticket_id AND + ticket_item.is_active == true AND + ticket_item.is_duplicate == false AND + ticket_item.is_whitelisted == false AND + ticket_item.is_error == false + + COLLECT genre = ticket_item.genre INTO groupedItems + FILTER genre IN ["fqdn", "ipv4", "ipv6"] + RETURN {{ + "genre": genre, + "items": UNIQUE(groupedItems[*].ticket_item.value) + }} + ) + + FILTER LENGTH(ticket_items) != 0 + + LET fqdn_items = (FOR item IN ticket_items FILTER item.genre == 'fqdn' 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) + + RETURN {{ + 'ticket_id': document.ticket_id, + 'description': document.description, + 'fqdn': fqdn_items[0] or [], + 'ipv4': ipv4_items[0] or [], + 'ipv6': ipv6_items[0] or [], + 'status': document.status, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise TicketStorageGetException() + + def get_total(self) -> int | Exception: + """ + Total documents in the tickets collection. + + :return: total number of documents. + """ + + try: + return self.collection_instance.count() + + except: + raise TicketStorageGetException() + + def exists_by_identifier(self, ticket_id: str) -> Cursor | Exception: + """ + Checks if a ticket with this identifier is in the collection. + + :param value: a valid ticket identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + LIMIT 1 + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'ticket_id': ticket_id + }) + + except: + raise TicketStorageGetException() + + def update_task_list(self, ticket_id: str, task_id: str) -> Cursor | Exception: + """ + Appends a new task. + + :param ticket_id: ticket identifier. + :param task_id: task identifier. + :return: list of updated rows. + """ + + aql = f""" + 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} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'task_id': task_id + }, + count = True + ) + + return affected_rows + + except: + raise TicketStorageUpdateException() + + def update_status(self, ticket_id: str, ticket_status: str) -> Cursor | Exception: + """ + Sets the ticket status. + Used by tasks to update from `CREATED` to `OPEN` and from `OPEN` to `CLOSED` once each time expires. + + :param ticket_id: ticket identifier. + :param ticket_status: ticket status value. + :return: list of updated rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + UPDATE document WITH {{ + 'status': @ticket_status + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id, + 'ticket_status': ticket_status + }, + count = True + ) + + return affected_rows + + except: + raise TicketStorageUpdateException() + + def remove(self, ticket_id: str) -> Cursor | Exception: + """ + Removes a ticket. + + :param ticket_id: ticket identifier. + :return: list of removed rows. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.ticket_id == @ticket_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'ticket_id': ticket_id + }, + count = True + ) + + return affected_rows + + except: + raise TicketStorageRemoveException() + +class TicketStorageCreateException(Exception): + + """ + Cannot create the ticket. + """ + + pass + +class TicketStorageGetException(Exception): + + """ + Cannot get the ticket. + """ + + pass + +class TicketStorageUpdateException(Exception): + + """ + Cannot update the ticket. + """ + + pass + +class TicketStorageRemoveException(Exception): + + """ + Cannot remove the ticket. + """ + + pass diff --git a/src/piracyshield_data_storage/whitelist/__init__.py b/src/piracyshield_data_storage/whitelist/__init__.py new file mode 100644 index 0000000..8d1c8b6 --- /dev/null +++ b/src/piracyshield_data_storage/whitelist/__init__.py @@ -0,0 +1 @@ + diff --git a/src/piracyshield_data_storage/whitelist/storage.py b/src/piracyshield_data_storage/whitelist/storage.py new file mode 100644 index 0000000..19fe7e5 --- /dev/null +++ b/src/piracyshield_data_storage/whitelist/storage.py @@ -0,0 +1,227 @@ +from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument + +from arango.cursor import Cursor + +class WhitelistStorage(DatabaseArangodbDocument): + + collection_name = 'whitelist' + + collection_instance = None + + def __init__(self): + super().__init__() + + self.collection_instance = self.collection(self.collection_name) + + def insert(self, document: dict) -> dict | Exception: + """ + Adds a new whitelist item. + + :param document: dictionary with the expected data model values. + :return: cursor with the inserted data. + """ + + try: + return self.collection_instance.insert(document) + + except: + raise WhitelistStorageCreateException() + + def get_all(self, account_id: str) -> Cursor | Exception: + """ + Fetches all the whitelist items created by an account. + + :param account_id: a valid account identifier. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.metadata.created_by == @account_id + + SORT document.value ASC + + RETURN {{ + 'genre': document.genre, + 'value': document.value, + 'is_active': document.is_active, + 'metadata': {{ + 'created_at': document.metadata.created_at + }} + }} + """ + + try: + return self.query(aql, bind_vars = { + 'account_id': account_id + }) + + except: + raise WhitelistStorageGetException() + + def get_global(self) -> Cursor | Exception: + """ + Fetches the whitelist items of all the accounts. + + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + // resolve account identifier + LET created_by_name = ( + FOR a IN accounts_view + FILTER a.account_id == document.metadata.created_by + RETURN a['name'] + )[0] + + SORT document.value ASC + + RETURN {{ + 'genre': document.genre, + 'value': document.value, + 'is_active': document.is_active, + 'metadata': {{ + 'created_at': document.metadata.created_at, + 'updated_at': document.metadata.updated_at, + 'created_by': document.metadata.created_by, + 'created_by_name': created_by_name + }} + }} + """ + + try: + return self.query(aql) + + except: + raise WhitelistStorageGetException() + + def exists_by_value(self, value: str) -> Cursor | Exception: + """ + Searches for an item. + + :param value: item value. + :return: cursor with the requested data. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.value == @value + + RETURN document + """ + + try: + return self.query(aql, bind_vars = { + 'value': value + }) + + except: + raise WhitelistStorageGetException() + + def update_status(self, value: str, status: bool) -> list | Exception: + """ + Sets the item status. + + :param value: item value. + :param status: true/false if active/non active. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER document.value == @value + + UPDATE document WITH {{ + is_active: @status + }} IN {self.collection_name} + + RETURN NEW + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'value': value, + 'status': status + }, + count = True + ) + + return affected_rows + + except: + return WhitelistStorageUpdateException() + + def remove(self, value: str, account_id: str) -> Cursor | Exception: + """ + Removes a whitelist item created by an account. + + :param value: item value. + :param account_id: a valid account identifier. + :return: true if the query has been processed successfully. + """ + + aql = f""" + FOR document IN {self.collection_name} + + FILTER + document.value == @value AND + document.metadata.created_by == @account_id + + REMOVE document IN {self.collection_name} + + RETURN OLD + """ + + try: + affected_rows = self.query( + aql, + bind_vars = { + 'value': value, + 'account_id': account_id + }, + count = True + ) + + return affected_rows + + except: + raise WhitelistStorageRemoveException() + +class WhitelistStorageCreateException(Exception): + + """ + Cannot create the whitelist item. + """ + + pass + +class WhitelistStorageGetException(Exception): + + """ + Cannot get the whitelist item. + """ + + pass + +class WhitelistStorageUpdateException(Exception): + + """ + Cannot update the whitelist item. + """ + + pass + +class WhitelistStorageRemoveException(Exception): + + """ + Cannot remove the whitelist item. + """ + + pass