| 1 | from fastapi import APIRouter |
| 2 | from fastapi import Body |
| 3 | from fastapi import Depends |
| 4 | from fastapi import HTTPException |
| 5 | from fastapi import Security |
| 6 | from loguru import logger |
| 7 | from sqlalchemy.ext.asyncio import AsyncSession |
| 8 | from sqlalchemy.future import select |
| 9 | |
| 10 | from app.auth.utils import AuthHandler |
| 11 | from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest |
| 12 | from app.connectors.grafana.schema.dashboards import WazuhDashboard |
| 13 | from app.customer_provisioning.schema.provision import CustomerProvisionResponse |
| 14 | from app.customer_provisioning.schema.provision import CustomersMetaResponse |
| 15 | from app.customer_provisioning.schema.provision import CustomerSubsctipion |
| 16 | from app.customer_provisioning.schema.provision import GetDashboardsResponse |
| 17 | from app.customer_provisioning.schema.provision import GetSubscriptionsResponse |
| 18 | from app.customer_provisioning.schema.provision import ProvisionDashboardRequest |
| 19 | from app.customer_provisioning.schema.provision import ProvisionDashboardResponse |
| 20 | from app.customer_provisioning.schema.provision import ProvisionNewCustomer |
| 21 | from app.customer_provisioning.schema.provision import UpdateOffice365OrgIdRequest |
| 22 | from app.customer_provisioning.schema.provision import UpdateOffice365OrgIdResponse |
| 23 | from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerRequest |
| 24 | from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerResponse |
| 25 | from app.customer_provisioning.services.provision import provision_dashboards |
| 26 | from app.customer_provisioning.services.provision import provision_wazuh_customer |
| 27 | from app.customer_provisioning.services.provision import provision_wazuh_worker |
| 28 | from app.db.db_session import get_db |
| 29 | from app.db.universal_models import Customers |
| 30 | from app.db.universal_models import CustomersMeta |
| 31 | |
| 32 | customer_provisioning_router = APIRouter() |
| 33 | |
| 34 | |
| 35 | def get_available_dashboards(): |
| 36 | """ |
| 37 | Get a list of available dashboards. |
| 38 | |
| 39 | Returns: |
| 40 | list: A list of available dashboards. |
| 41 | |
| 42 | Raises: |
| 43 | HTTPException: If there is an error getting the available dashboards. |
| 44 | """ |
| 45 | try: |
| 46 | wazuh_dashboards = [dashboard.name for dashboard in WazuhDashboard] |
| 47 | # office365_dashboards = [dashboard.name for dashboard in Office365Dashboard] |
| 48 | # return wazuh_dashboards + office365_dashboards |
| 49 | return wazuh_dashboards |
| 50 | except Exception as e: |
| 51 | raise HTTPException( |
| 52 | status_code=500, |
| 53 | detail=f"Error getting available dashboards: {e}", |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def get_available_subscriptions(): |
| 58 | """ |
| 59 | Retrieves a list of available subscriptions. |
| 60 | |
| 61 | Returns: |
| 62 | list: A list of available subscriptions. |
| 63 | |
| 64 | Raises: |
| 65 | HTTPException: If there is an error getting the available subscriptions. |
| 66 | """ |
| 67 | try: |
| 68 | return [subscription.value for subscription in CustomerSubsctipion] |
| 69 | except Exception as e: |
| 70 | raise HTTPException( |
| 71 | status_code=500, |
| 72 | detail=f"Error getting available subscriptions: {e}", |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | async def check_customer_exists( |
| 77 | customer_code: str, |
| 78 | session: AsyncSession = Depends(get_db), |
| 79 | ) -> Customers: |
| 80 | """ |
| 81 | Check if a customer exists in the database. |
| 82 | |
| 83 | Args: |
| 84 | customer_code (str): The code of the customer to check. |
| 85 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 86 | |
| 87 | Returns: |
| 88 | Customers: The customer object if found. |
| 89 | |
| 90 | Raises: |
| 91 | HTTPException: If the customer is not found in the database. |
| 92 | """ |
| 93 | logger.info(f"Checking if customer {customer_code} exists") |
| 94 | result = await session.execute( |
| 95 | select(Customers).filter(Customers.customer_code == customer_code), |
| 96 | ) |
| 97 | customer = result.scalars().first() |
| 98 | |
| 99 | if not customer: |
| 100 | raise HTTPException( |
| 101 | status_code=404, |
| 102 | detail=f"Customer: {customer_code} not found. Please create the customer first.", |
| 103 | ) |
| 104 | |
| 105 | return customer |
| 106 | |
| 107 | |
| 108 | async def check_unique_ports(request: ProvisionNewCustomer, session: AsyncSession): |
| 109 | """ |
| 110 | Ensures that the ports specified in the request are unique. |
| 111 | |
| 112 | Args: |
| 113 | request: The request data containing the ports to check. |
| 114 | session: The database session. |
| 115 | |
| 116 | Raises: |
| 117 | HTTPException: If the ports are not unique. |
| 118 | """ |
| 119 | ports = { |
| 120 | "registration": request.wazuh_registration_port, |
| 121 | "logs": request.wazuh_logs_port, |
| 122 | "api": request.wazuh_api_port, |
| 123 | } |
| 124 | |
| 125 | for port_type, port_value in ports.items(): |
| 126 | customer_meta = await get_customer_meta_by_port(port_value, session) |
| 127 | if customer_meta: |
| 128 | raise HTTPException( |
| 129 | status_code=400, |
| 130 | detail=f"Ports must be unique. {port_type.capitalize()} port {port_value} is already in use for customer {customer_meta.customer_code}.", |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | async def get_customer_meta_by_port(port: int, session: AsyncSession): |
| 135 | """ |
| 136 | Retrieves customer metadata based on the provided port. |
| 137 | |
| 138 | Args: |
| 139 | port: The port to check. |
| 140 | session: The database session. |
| 141 | |
| 142 | Returns: |
| 143 | Customer metadata if a match is found, otherwise None. |
| 144 | """ |
| 145 | result = await session.execute( |
| 146 | select(CustomersMeta).filter( |
| 147 | (CustomersMeta.customer_meta_wazuh_registration_port == port) |
| 148 | | (CustomersMeta.customer_meta_wazuh_log_ingestion_port == port) |
| 149 | | (CustomersMeta.customer_meta_wazuh_api_port == port), |
| 150 | ), |
| 151 | ) |
| 152 | return result.scalars().first() |
| 153 | |
| 154 | |
| 155 | async def update_customer_meta_table( |
| 156 | request: ProvisionNewCustomer, |
| 157 | session: AsyncSession, |
| 158 | ): |
| 159 | """ |
| 160 | Update the customer meta table with the provided information. |
| 161 | |
| 162 | Args: |
| 163 | request (ProvisionNewCustomer): The request object containing customer information. |
| 164 | customer_meta (CustomerProvisionMeta): The customer meta object containing additional information. |
| 165 | session (AsyncSession): The database session. |
| 166 | |
| 167 | Returns: |
| 168 | CustomerProvisionMeta: The updated customer meta object. |
| 169 | """ |
| 170 | logger.info(f"Updating customer meta table for customer {request.customer_name}") |
| 171 | customer_meta = CustomersMeta( |
| 172 | customer_code=request.customer_code, |
| 173 | customer_name=request.customer_name, |
| 174 | customer_meta_graylog_index=request.graylog_index_id, |
| 175 | customer_meta_graylog_stream=request.graylog_stream_id, |
| 176 | customer_meta_grafana_org_id=request.grafana_org_id, |
| 177 | customer_meta_wazuh_group=request.customer_code, |
| 178 | customer_meta_index_retention=str(request.hot_data_retention), |
| 179 | customer_meta_wazuh_registration_port=request.wazuh_registration_port, |
| 180 | customer_meta_wazuh_log_ingestion_port=request.wazuh_logs_port, |
| 181 | customer_meta_wazuh_api_port=request.wazuh_api_port, |
| 182 | customer_meta_wazuh_auth_password=request.wazuh_auth_password, |
| 183 | ) |
| 184 | session.add(customer_meta) |
| 185 | await session.commit() |
| 186 | return customer_meta |
| 187 | |
| 188 | |
| 189 | @customer_provisioning_router.post( |
| 190 | "/provision", |
| 191 | response_model=CustomerProvisionResponse, |
| 192 | description="Provision New Customer", |
| 193 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 194 | ) |
| 195 | async def provision_customer_route( |
| 196 | request: ProvisionNewCustomer = Body(...), |
| 197 | _customer: Customers = Depends(check_customer_exists), |
| 198 | session: AsyncSession = Depends(get_db), |
| 199 | ): |
| 200 | """ |
| 201 | Provisions a new customer. |
| 202 | |
| 203 | Args: |
| 204 | request (ProvisionNewCustomer): The request data for provisioning a new customer. |
| 205 | _customer (Customers): The existing customer data. |
| 206 | session (AsyncSession): The database session. |
| 207 | |
| 208 | Returns: |
| 209 | CustomerProvisionResponse: The response data for the provisioned customer. |
| 210 | """ |
| 211 | if request.provision_wazuh_worker is True: |
| 212 | await check_unique_ports(request, session) |
| 213 | logger.info("Provisioning new customer") |
| 214 | if request.only_insert_into_db is True: |
| 215 | logger.info("Only inserting into the database") |
| 216 | customer_meta = await update_customer_meta_table(request, session=session) |
| 217 | return CustomerProvisionResponse( |
| 218 | success=True, |
| 219 | message="Customer inserted into the database successfully", |
| 220 | customer_meta=customer_meta, |
| 221 | wazuh_worker_provisioned=False, |
| 222 | ) |
| 223 | customer_provision = await provision_wazuh_customer(request, session=session) |
| 224 | return customer_provision |
| 225 | |
| 226 | |
| 227 | @customer_provisioning_router.post( |
| 228 | "/provision/wazuh_worker", |
| 229 | response_model=ProvisionWorkerResponse, |
| 230 | description="Provision Wazuh Worker", |
| 231 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 232 | ) |
| 233 | async def provision_wazuh_worker_route( |
| 234 | request: ProvisionWorkerRequest = Body(...), |
| 235 | session: AsyncSession = Depends(get_db), |
| 236 | ): |
| 237 | """ |
| 238 | Provisions a new Wazuh worker. |
| 239 | |
| 240 | Args: |
| 241 | request (ProvisionWorkerRequest): The request data for provisioning a new Wazuh worker. |
| 242 | session (AsyncSession): The database session. |
| 243 | |
| 244 | Returns: |
| 245 | ProvisionWorkerResponse: The response data for the provisioned Wazuh worker. |
| 246 | """ |
| 247 | logger.info("Provisioning Wazuh worker") |
| 248 | return await provision_wazuh_worker(request, session=session) |
| 249 | |
| 250 | |
| 251 | @customer_provisioning_router.get( |
| 252 | "/provision/dashboards", |
| 253 | response_model=GetDashboardsResponse, |
| 254 | description="Return the list of dashboards available for provisioning", |
| 255 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 256 | ) |
| 257 | async def get_dashboards_route(): |
| 258 | """ |
| 259 | Get the list of dashboards available for provisioning. |
| 260 | |
| 261 | Returns: |
| 262 | GetDashboardsResponse: The response containing the available dashboards. |
| 263 | """ |
| 264 | logger.info("Getting list of dashboards") |
| 265 | available_dashboards = get_available_dashboards() |
| 266 | return GetDashboardsResponse( |
| 267 | available_dashboards=available_dashboards, |
| 268 | success=True, |
| 269 | message="Dashboards retrieved successfully", |
| 270 | ) |
| 271 | |
| 272 | |
| 273 | @customer_provisioning_router.get( |
| 274 | "/provision/subscriptions", |
| 275 | response_model=GetSubscriptionsResponse, |
| 276 | description="Return the list of subscriptions available for provisioning", |
| 277 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 278 | ) |
| 279 | async def get_subscriptions_route(): |
| 280 | """ |
| 281 | Get the list of subscriptions available for provisioning. |
| 282 | |
| 283 | Returns: |
| 284 | GetSubscriptionsResponse: The response containing the available subscriptions. |
| 285 | """ |
| 286 | logger.info("Getting list of subscriptions") |
| 287 | available_subscriptions = get_available_subscriptions() |
| 288 | return GetSubscriptionsResponse( |
| 289 | available_subscriptions=available_subscriptions, |
| 290 | success=True, |
| 291 | message="Subscriptions retrieved successfully", |
| 292 | ) |
| 293 | |
| 294 | |
| 295 | # Get the customermeta based on the customer code |
| 296 | @customer_provisioning_router.get( |
| 297 | "/provision/{customer_code}", |
| 298 | response_model=CustomersMetaResponse, |
| 299 | description="Get Customer Meta", |
| 300 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 301 | ) |
| 302 | async def get_customer_meta( |
| 303 | customer_code: str, |
| 304 | session: AsyncSession = Depends(get_db), |
| 305 | ): |
| 306 | """ |
| 307 | Retrieve customer meta data for a given customer code. |
| 308 | |
| 309 | Args: |
| 310 | customer_code (str): The code of the customer to retrieve meta data for. |
| 311 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 312 | |
| 313 | Raises: |
| 314 | HTTPException: If customer meta data is not found for the given customer code. |
| 315 | |
| 316 | Returns: |
| 317 | CustomersMetaResponse: The response containing the customer meta data. |
| 318 | """ |
| 319 | logger.info(f"Getting customer meta for customer {customer_code}") |
| 320 | result = await session.execute( |
| 321 | select(CustomersMeta).filter(CustomersMeta.customer_code == customer_code), |
| 322 | ) |
| 323 | customer_meta = result.scalars().first() |
| 324 | |
| 325 | if not customer_meta: |
| 326 | raise HTTPException( |
| 327 | status_code=404, |
| 328 | detail=f"Customer meta not found for customer: {customer_code}. Please provision the customer first.", |
| 329 | ) |
| 330 | |
| 331 | return CustomersMetaResponse( |
| 332 | message="Customer meta retrieved successfully", |
| 333 | success=True, |
| 334 | customer_meta=customer_meta, |
| 335 | ) |
| 336 | |
| 337 | |
| 338 | @customer_provisioning_router.post( |
| 339 | "/provision/dashboards", |
| 340 | response_model=ProvisionDashboardResponse, |
| 341 | description="Return the list of dashboards available for provisioning", |
| 342 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 343 | ) |
| 344 | async def provision_dashboards_route( |
| 345 | request: ProvisionDashboardRequest = Body(...), |
| 346 | session: AsyncSession = Depends(get_db), |
| 347 | ): |
| 348 | """ |
| 349 | Provision dashboards for a customer. |
| 350 | |
| 351 | Args: |
| 352 | request (ProvisionDashboardsRequest): The request data for provisioning dashboards. |
| 353 | session (AsyncSession): The database session. |
| 354 | |
| 355 | Returns: |
| 356 | ProvisionDashboardsResponse: The response data for the provisioned dashboards. |
| 357 | """ |
| 358 | logger.info("Provisioning dashboards") |
| 359 | return await provision_dashboards( |
| 360 | DashboardProvisionRequest( |
| 361 | dashboards=request.dashboards_to_include.dashboards, |
| 362 | organizationId=request.grafana_org_id, |
| 363 | folderId=request.grafana_folder_id, |
| 364 | datasourceUid=request.grafana_datasource_uid, |
| 365 | grafana_url=request.grafana_url, |
| 366 | ), |
| 367 | ) |
| 368 | |
| 369 | |
| 370 | @customer_provisioning_router.put( |
| 371 | "/update/office365_org_id/{customer_code}", |
| 372 | response_model=UpdateOffice365OrgIdResponse, |
| 373 | description="Update Office 365 organization ID", |
| 374 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 375 | ) |
| 376 | async def update_office_365_org_id( |
| 377 | customer_code: str, |
| 378 | request: UpdateOffice365OrgIdRequest = Body(...), |
| 379 | session: AsyncSession = Depends(get_db), |
| 380 | ): |
| 381 | """ |
| 382 | Update Office 365 organization ID for a customer. |
| 383 | |
| 384 | Args: |
| 385 | customer_code (str): The code of the customer to update. |
| 386 | request (UpdateOffice365OrgIdRequest): The request data for updating Office 365 organization ID. |
| 387 | session (AsyncSession): The database session. |
| 388 | |
| 389 | Returns: |
| 390 | UpdateOffice365OrgIdResponse: The response data for the updated Office 365 organization ID. |
| 391 | """ |
| 392 | logger.info("Updating Office 365 organization ID") |
| 393 | # Update within the `CustomersMeta` table based on the customer code |
| 394 | stmt = select(CustomersMeta).where(CustomersMeta.customer_code == customer_code) |
| 395 | result = await session.execute(stmt) |
| 396 | customer_meta = result.scalars().first() |
| 397 | if not customer_meta: |
| 398 | raise HTTPException( |
| 399 | status_code=404, |
| 400 | detail=f"Customer meta not found for customer: {customer_code}. Please provision the customer first.", |
| 401 | ) |
| 402 | customer_meta.customer_meta_office365_organization_id = request.office365_org_id |
| 403 | await session.commit() |
| 404 | return UpdateOffice365OrgIdResponse( |
| 405 | message="Office 365 organization ID updated successfully", |
| 406 | success=True, |
| 407 | ) |