feat(#3339): Add schema for custom table

wazuh · Jan 9, 2023 · 970c8c0 · 970c8c0
1 parent 0fdddf3
commit 970c8c0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git 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
@@ -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
 }