Add customer meta table update functionality*** (#157)
* Add customer meta table update functionality*** * precommit fixes
taylor_socfortress committed
Feb 15, 2024 at 07:57 UTC
e86e4161229acc19e8213ad04a01345729bd38ca
2 files changed
+60
backend/app/customer_provisioning/routes/provision.py
+44
@@ -143,6 +143,41 @@ async def get_customer_meta_by_port(port: int, session: AsyncSession):
143
return result.scalars().first()
144
145
146
+async def update_customer_meta_table(
147
+ request: ProvisionNewCustomer,
148
+ session: AsyncSession,
149
+):
150
+ """
151
+ Update the customer meta table with the provided information.
152
+
153
+ Args:
154
+ request (ProvisionNewCustomer): The request object containing customer information.
155
+ customer_meta (CustomerProvisionMeta): The customer meta object containing additional information.
156
+ session (AsyncSession): The database session.
157
+
158
+ Returns:
159
+ CustomerProvisionMeta: The updated customer meta object.
160
+ """
161
+ logger.info(f"Updating customer meta table for customer {request.customer_name}")
162
+ customer_meta = CustomersMeta(
163
+ customer_code=request.customer_code,
164
+ customer_name=request.customer_name,
165
+ customer_meta_graylog_index=request.graylog_index_id,
166
+ customer_meta_graylog_stream=request.graylog_stream_id,
167
+ customer_meta_grafana_org_id=request.dashboards_to_include.organizationId,
168
+ customer_meta_wazuh_group=request.customer_code,
169
+ customer_meta_index_retention=str(request.hot_data_retention),
170
+ customer_meta_wazuh_registration_port=request.wazuh_registration_port,
171
+ customer_meta_wazuh_log_ingestion_port=request.wazuh_logs_port,
172
+ customer_meta_wazuh_api_port=request.wazuh_api_port,
173
+ customer_meta_wazuh_auth_password=request.wazuh_auth_password,
174
+ customer_meta_iris_customer_id=request.dfir_iris_id,
175
+ )
176
+ session.add(customer_meta)
177
+ await session.commit()
178
+ return customer_meta
179
+
180
+
181
@customer_provisioning_router.post(
182
"/provision",
183
response_model=CustomerProvisionResponse,
@@ -167,6 +202,15 @@ async def provision_customer_route(
202
"""
203
await check_unique_ports(request, session)
204
logger.info("Provisioning new customer")
205
+ if request.only_insert_into_db is True:
206
+ logger.info("Only inserting into the database")
207
+ customer_meta = await update_customer_meta_table(request, session=session)
208
+ return CustomerProvisionResponse(
209
+ success=True,
210
+ message="Customer inserted into the database successfully",
211
+ customer_meta=customer_meta,
212
+ wazuh_worker_provisioned=False,
213
+ )
214
customer_provision = await provision_wazuh_customer(request, session=session)
215
return customer_provision
216
backend/app/customer_provisioning/schema/provision.py
+16
@@ -72,6 +72,22 @@ class ProvisionNewCustomer(BaseModel):
72
wazuh_cluster_key: str = Field(..., description="Password for the Wazuh cluster")
73
wazuh_master_ip: str = Field(..., description="IP address of the Wazuh master")
74
grafana_url: str = Field(..., description="URL of the Grafana instance")
75
+ only_insert_into_db: Optional[bool] = Field(
76
+ False,
77
+ description="Whether to only insert the customer into the database without provisioning any services",
78
+ )
79
+ dfir_iris_id: Optional[int] = Field(
80
+ None,
81
+ description="ID of the DFIR Iris customer",
82
+ )
83
+ graylog_index_id: Optional[str] = Field(
84
+ None,
85
+ description="ID of the Graylog index set",
86
+ )
87
+ graylog_stream_id: Optional[str] = Field(
88
+ None,
89
+ description="ID of the Graylog stream",
90
+ )
91
92
@validator("customer_index_name")
93
def validate_customer_index_name(cls, v):