@cryptotaxi247 / CoPilot
Code Commits Issues Pulls Sessions CI/CD
main
main
py 145 lines 5.52 KB
Raw
1 from typing import Dict
2
3 from fastapi import APIRouter
4 from fastapi import Depends
5 from fastapi import HTTPException
6 from fastapi import Security
7 from sqlalchemy.ext.asyncio import AsyncSession
8
9 from app.auth.utils import AuthHandler
10 from app.connectors.graylog.utils.routing import GraylogContext
11 from app.connectors.graylog.utils.routing import clear_graylog_context
12 from app.connectors.graylog.utils.routing import set_graylog_context
13 from app.db.db_session import get_db
14 from app.network_connectors.routes import find_customer_network_connector
15 from app.network_connectors.routes import (
16 get_customer_network_connectors_by_customer_code,
17 )
18 from app.network_connectors.schema import CustomerNetworkConnectors
19 from app.network_connectors.schema import CustomerNetworkConnectorsResponse
20 from app.stack_provisioning.graylog.schema.fortinet import FortinetCustomerDetails
21 from app.stack_provisioning.graylog.schema.fortinet import ProvisionFortinetKeys
22 from app.stack_provisioning.graylog.schema.fortinet import ProvisionFortinetRequest
23 from app.stack_provisioning.graylog.schema.fortinet import ProvisionFortinetResponse
24 from app.stack_provisioning.graylog.services.fortinet import provision_fortinet
25
26 stack_provisioning_graylog_fortinet_router = APIRouter()
27
28
29 async def get_customer_integration_response(
30 customer_code: str,
31 session: AsyncSession,
32 ) -> CustomerNetworkConnectorsResponse:
33 """
34 Retrieves the integration response for a customer.
35
36 Args:
37 customer_code (str): The code of the customer.
38 session (AsyncSession): The async session object for database operations.
39
40 Returns:
41 CustomerIntegrationsResponse: The integration response for the customer.
42
43 Raises:
44 HTTPException: If the customer integration settings are not found.
45 """
46 customer_integration_response = await get_customer_network_connectors_by_customer_code(
47 customer_code,
48 session,
49 )
50 if customer_integration_response.available_network_connectors == []:
51 raise HTTPException(
52 status_code=404,
53 detail="Customer integration settings not found.",
54 )
55 return customer_integration_response
56
57
58 def extract_fortinet_keys(
59 customer_integration: CustomerNetworkConnectors,
60 ) -> Dict[str, str]:
61 """
62 Extracts the authentication keys for Office365 integration from the given customer integration.
63
64 Args:
65 customer_integration (CustomerIntegrations): The customer integration object.
66
67 Returns:
68 Dict[str, str]: A dictionary containing the authentication keys for Office365 integration.
69
70 Raises:
71 HTTPException: If no authentication keys are found for Office365 integration.
72 """
73 fortinet_keys = {}
74 for subscription in customer_integration.network_connectors_subscriptions:
75 if subscription.network_connectors_service.service_name == "Fortinet":
76 for auth_key in subscription.network_connectors_keys:
77 fortinet_keys[auth_key.auth_key_name] = auth_key.auth_value
78 if not fortinet_keys:
79 raise HTTPException(
80 status_code=404,
81 detail="No auth keys found for Fortinet integration. Please create auth keys for Fortinet network connector.",
82 )
83 return fortinet_keys
84
85
86 @stack_provisioning_graylog_fortinet_router.post(
87 "/graylog/provision/fortinet",
88 response_model=ProvisionFortinetResponse,
89 description="Provision Fortinet for the customer.",
90 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
91 )
92 async def provision_fortinet_route(
93 provision_fortinet_request: ProvisionFortinetRequest,
94 session: AsyncSession = Depends(get_db),
95 ) -> ProvisionFortinetResponse:
96 """
97 Provision Fortinet for the customer
98 """
99
100 # Set the Graylog context for this request - all downstream Graylog calls will use Graylog-Network
101 set_graylog_context(GraylogContext.NETWORK)
102
103 try:
104 customer_integration_response = await get_customer_integration_response(
105 provision_fortinet_request.customer_code,
106 session,
107 )
108
109 customer_integration = await find_customer_network_connector(
110 provision_fortinet_request.customer_code,
111 provision_fortinet_request.integration_name,
112 customer_integration_response,
113 )
114
115 fortinet_keys = extract_fortinet_keys(customer_integration)
116
117 if provision_fortinet_request.tcp_enabled and provision_fortinet_request.udp_enabled:
118 raise HTTPException(
119 status_code=400,
120 detail="Both TCP and UDP are enabled. Please choose one of them.",
121 )
122 elif provision_fortinet_request.tcp_enabled:
123 protocol_type = "TCP"
124 elif provision_fortinet_request.udp_enabled:
125 protocol_type = "UDP"
126 else:
127 raise HTTPException(
128 status_code=400,
129 detail="Either TCP or UDP should be enabled.",
130 )
131
132 return await provision_fortinet(
133 customer_details=FortinetCustomerDetails(
134 customer_code=provision_fortinet_request.customer_code,
135 customer_name=customer_integration.customer_name,
136 protocal_type=protocol_type,
137 syslog_port=int(fortinet_keys["SYSLOG_PORT"]),
138 hot_data_retention=provision_fortinet_request.hot_data_retention,
139 index_replicas=provision_fortinet_request.index_replicas,
140 ),
141 keys=ProvisionFortinetKeys(**fortinet_keys),
142 session=session,
143 )
144 finally:
145 clear_graylog_context()