From 970c8c0cb528a6c24868dea594e2fb136b7cb9c2 Mon Sep 17 00:00:00 2001 From: Nicolas Stefani Date: Thu, 29 Dec 2022 11:09:17 -0300 Subject: [PATCH] feat(#3339): Add schema for custom table --- .../wazuh_testing/modules/aws/constants.py | 1 + .../wazuh_testing/modules/aws/db_utils.py | 15 +++++++++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/deps/wazuh_testing/wazuh_testing/modules/aws/constants.py b/deps/wazuh_testing/wazuh_testing/modules/aws/constants.py index db184a813d..84dfbedd5f 100644 --- a/deps/wazuh_testing/wazuh_testing/modules/aws/constants.py +++ b/deps/wazuh_testing/wazuh_testing/modules/aws/constants.py @@ -23,6 +23,7 @@ # Bucket types CLOUD_TRAIL_TYPE = 'cloudtrail' VPC_FLOW_TYPE = 'vpcflow' +ALB_TYPE = 'alb' KMS_TYPE = 'kms' TRUSTED_ADVISOR_TYPE = 'trusted' CUSTOM_TYPE = 'custom' diff --git a/deps/wazuh_testing/wazuh_testing/modules/aws/db_utils.py b/deps/wazuh_testing/wazuh_testing/modules/aws/db_utils.py index 177f3f9784..16972407ad 100644 --- a/deps/wazuh_testing/wazuh_testing/modules/aws/db_utils.py +++ b/deps/wazuh_testing/wazuh_testing/modules/aws/db_utils.py @@ -3,7 +3,13 @@ from pathlib import Path from typing import Iterator, Type -from .constants import S3_CLOUDTRAIL_DB_PATH, CLOUD_TRAIL_TYPE, ALB_TYPE +from .constants import ( + ALB_TYPE, + CLOUD_TRAIL_TYPE, + CUSTOM_TYPE, + KMS_TYPE, + S3_CLOUDTRAIL_DB_PATH, +) SELECT_QUERY_TEMPLATE = 'SELECT * FROM {table_name}' @@ -15,9 +21,14 @@ 'S3ALBFlowRow', 'bucket_path aws_account_id log_key processed_date created_date' ) +S3CustomRow = namedtuple( + 'S3CustomRow', 'bucket_path aws_account_id log_key processed_date created_date' +) + s3_rows_map = { CLOUD_TRAIL_TYPE: S3CloudTrailRow, - ALB_TYPE: S3ALBFlowRow + ALB_TYPE: S3ALBFlowRow, + CUSTOM_TYPE: S3CustomRow }