From 0b7cb8088ca7c8702e67b94cb2363c7ae6774185 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B6rg=20Stucke?= Date: Mon, 16 Dec 2024 16:21:39 +0100 Subject: [PATCH] feat: add db view unpacking --- src/init_postgres.py | 3 +- src/storage/db_setup.py | 35 +++++++++++++++++++++-- src/storage/graphql/hasura/init_hasura.py | 34 +++++++++++++++++++++- src/test/common_helper.py | 2 ++ 4 files changed, 70 insertions(+), 4 deletions(-) diff --git a/src/init_postgres.py b/src/init_postgres.py index 7f8b6bf2c..8ebe82570 100755 --- a/src/init_postgres.py +++ b/src/init_postgres.py @@ -86,10 +86,11 @@ def main(command_line_options=None, config_path: str | None = None, skip_user_cr if not db_setup.table_exists('file_object'): logging.info('Creating FACT tables...') db_setup.connection.create_tables() - db_setup.set_table_privileges() set_alembic_revision() else: logging.info('Skipping creation of FACT tables: already exist') + db_setup.create_unpacking_view() + db_setup.set_table_privileges() return 0 diff --git a/src/storage/db_setup.py b/src/storage/db_setup.py index 2d1a536ee..26e2699da 100644 --- a/src/storage/db_setup.py +++ b/src/storage/db_setup.py @@ -6,6 +6,18 @@ from storage.db_connection import AdminConnection, DbConnection from storage.db_interface_base import ReadWriteDbInterface +CREATE_VIEW_CMD = text( + """ +CREATE VIEW unpacking AS + SELECT + uid, + (result->>'entropy')::float AS entropy, + (result->>'number_of_unpacked_files')::integer AS number_of_unpacked_files + FROM analysis + WHERE plugin = 'unpacker'; +""" +) + class Privileges: SELECT = 'SELECT' @@ -38,6 +50,10 @@ def table_exists(self, table_name: str): with self.connection.engine.connect() as db, db.engine.begin() as connection: return inspect(connection).has_table(table_name, None) + def view_exists(self, view_name: str): + with self.connection.engine.connect() as db, db.engine.begin() as connection: + return view_name in inspect(connection).get_view_names(schema='public') + def database_exists(self, db_name: str) -> bool: with self.get_read_only_session() as session: return bool(session.execute(text(f"SELECT 1 FROM pg_database WHERE datname = '{db_name}'")).scalar()) @@ -63,8 +79,23 @@ def set_table_privileges(self): ]: user = getattr(config.backend.postgres, f'{key}_user') for privilege in privileges: - self.grant_privilege(user, privilege) + self.grant_table_privilege(user, privilege) + self.grant_view_privilege(user, privilege) - def grant_privilege(self, user_name: str, privilege: str): + def grant_table_privilege(self, user_name: str, privilege: str): with self.get_read_write_session() as session: session.execute(text(f'GRANT {privilege} ON ALL TABLES IN SCHEMA public TO {user_name};')) + + def create_unpacking_view(self): + if not self.view_exists('unpacking'): + with self.get_read_write_session() as session: + session.execute(CREATE_VIEW_CMD) + + def grant_view_privilege(self, user_name: str, privilege: str): + with self.get_read_write_session() as session: + session.execute(text(f'GRANT {privilege} ON unpacking TO {user_name};')) + + def delete_unpacking_view(self): + if self.view_exists('unpacking'): + with self.get_read_write_session() as session: + session.execute(text('DROP VIEW IF EXISTS unpacking;')) diff --git a/src/storage/graphql/hasura/init_hasura.py b/src/storage/graphql/hasura/init_hasura.py index 49ecb7086..9a8ef7eba 100644 --- a/src/storage/graphql/hasura/init_hasura.py +++ b/src/storage/graphql/hasura/init_hasura.py @@ -2,6 +2,7 @@ import logging import sys +from itertools import permutations from pathlib import Path import requests @@ -17,7 +18,15 @@ HTML_OK = 200 HTML_BAD_REQUEST = 400 -TRACKED_TABLES = ('analysis', 'file_object', 'firmware', 'fw_files', 'included_files', 'virtual_file_path') +TRACKED_TABLES = ( + 'analysis', + 'file_object', + 'firmware', + 'fw_files', + 'included_files', + 'virtual_file_path', + 'unpacking', +) RELATIONSHIPS = { 'pg_create_object_relationship': [ # table, name, constraint @@ -64,6 +73,7 @@ def init_hasura(self, db_args: dict | None = None): self._add_database(db_args) self._track_tables() self._add_relationships() + self._add_view_relationship() self._add_ro_user_role_to_tables() logging.info('Hasura initialization successful') @@ -163,6 +173,28 @@ def _add_relationships(self): f'Failed to add constraint {name} on table {table}: {response.json().get("error")}' ) + def _add_view_relationship(self): + for source, target in permutations(['unpacking', 'file_object']): + query = { + 'type': 'pg_create_object_relationship', + 'args': { + 'table': {'name': source, 'schema': 'public'}, + 'name': target, + 'source': self.db_name, + 'using': { + 'manual_configuration': { + 'remote_table': target, + 'column_mapping': {'uid': 'uid'}, + } + }, + }, + } + response = requests.post(self.url, headers=self.headers, json=query) + if response.status_code != HTML_OK: + if _was_already_added(response): + continue + raise HasuraInitError(f'Failed to add relationships on view unpacking: {response.json().get("error")}') + def _db_was_already_added(self) -> bool: query = {'type': 'pg_get_source_tables', 'args': {'source': self.db_name}} response = requests.post(self.url, headers=self.headers, json=query) diff --git a/src/test/common_helper.py b/src/test/common_helper.py index e990a1dce..772306482 100644 --- a/src/test/common_helper.py +++ b/src/test/common_helper.py @@ -294,10 +294,12 @@ def store_binary_on_file_system(tmp_dir: str, test_object: FileObject | Firmware def setup_test_tables(db_setup): db_setup.connection.create_tables() + db_setup.create_unpacking_view() db_setup.set_table_privileges() def clear_test_tables(db_setup): + db_setup.delete_unpacking_view() db_setup.connection.base.metadata.drop_all(db_setup.connection.engine)