@cryptotaxi247 / CoPilot / commits / 495a2d76

Asynccontextmanager (#110)

* setup dbsession * converted all routes to establish a db session with `get_bd` instead of using `get_session`

taylor_socfortress committed Dec 22, 2023 at 13:05 UTC 495a2d76cb085f846356a17007a6e16bc5b3380b

16 files changed +113 -99

backend/app/agents/routes/agents.py

+9 -9

@@ -25,7 +25,7 @@ from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilit

26 # App specific imports

27 from app.auth.routes.auth import AuthHandler

28 -from app.db.db_session import get_session

28 +from app.db.db_session import get_session, get_db

30 # App specific imports

31 # from app.db.db_session import session

@@ -78,7 +78,7 @@ async def delete_agent_from_database(db: AsyncSession, agent_id: str):

78 description="Get all agents currently synced to the database",

79 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

80 )

81 -async def get_agents(db: AsyncSession = Depends(get_session)) -> AgentsResponse:

81 +async def get_agents(db: AsyncSession = Depends(get_db)) -> AgentsResponse:

82 logger.info("Fetching all agents")

83 try:

84 # agents = session.query(Agents).all()

@@ -96,7 +96,7 @@ async def get_agents(db: AsyncSession = Depends(get_session)) -> AgentsResponse:

96 description="Get agent by agent_id",

97 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

98 )

99 -async def get_agent(agent_id: str, db: AsyncSession = Depends(get_session)) -> AgentsResponse:

99 +async def get_agent(agent_id: str, db: AsyncSession = Depends(get_db)) -> AgentsResponse:

100 logger.info(f"Fetching agent with agent_id: {agent_id}")

101 try:

102 result = await db.execute(select(Agents).filter(Agents.agent_id == agent_id))

@@ -116,7 +116,7 @@ async def get_agent(agent_id: str, db: AsyncSession = Depends(get_session)) -> A

116 description="Get agent by hostname",

117 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

118 )

119 -async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_session)) -> AgentsResponse:

119 +async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_db)) -> AgentsResponse:

120 logger.info(f"Fetching agent with hostname: {hostname}")

121 try:

122 result = await db.execute(select(Agents).filter(Agents.hostname == hostname))

@@ -137,7 +137,7 @@ async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_se

137 description="Sync agents from Wazuh Manager",

138 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler"))],

139 )

140 -async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSession = Depends(get_session)) -> SyncedAgentsResponse:

140 +async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSession = Depends(get_db)) -> SyncedAgentsResponse:

141 logger.info("Syncing agents from Wazuh Manager")

142 backgroud_tasks.add_task(sync_agents, session)

143 # return sync_agents()

@@ -150,7 +150,7 @@ async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSessio

150 description="Mark agent as critical",

151 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

152 )

153 -async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(get_session)) -> AgentModifyResponse:

153 +async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(get_db)) -> AgentModifyResponse:

154 logger.info(f"Marking agent {agent_id} as critical")

155 # return mark_agent_criticality(agent_id, True)

156 try:

@@ -176,7 +176,7 @@ async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(

176 description="Mark agent as not critical",

177 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

178 )

179 -async def mark_agent_as_not_critical(agent_id: str, session: AsyncSession = Depends(get_session)) -> AgentModifyResponse:

179 +async def mark_agent_as_not_critical(agent_id: str, session: AsyncSession = Depends(get_db)) -> AgentModifyResponse:

180 logger.info(f"Marking agent {agent_id} as not critical")

181 try:

182 # Asynchronously fetch the agent by id

@@ -212,7 +212,7 @@ async def get_agent_vulnerabilities(agent_id: str) -> WazuhAgentVulnerabilitiesR

212 description="Get all outdated Wazuh agents",

213 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

214 )

215 -async def get_outdated_wazuh_agents(session: AsyncSession = Depends(get_session)) -> OutdatedWazuhAgentsResponse:

215 +async def get_outdated_wazuh_agents(session: AsyncSession = Depends(get_db)) -> OutdatedWazuhAgentsResponse:

216 logger.info("Fetching all outdated Wazuh agents")

217 return await get_outdated_agents_wazuh(session)

218

@@ -223,7 +223,7 @@ async def get_outdated_wazuh_agents(session: AsyncSession = Depends(get_session)

223 description="Get all outdated Velociraptor agents",

224 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

225 )

226 -async def get_outdated_velociraptor_agents(session: AsyncSession = Depends(get_session)) -> OutdatedVelociraptorAgentsResponse:

226 +async def get_outdated_velociraptor_agents(session: AsyncSession = Depends(get_db)) -> OutdatedVelociraptorAgentsResponse:

227 logger.info("Fetching all outdated Velociraptor agents")

228 return await get_outdated_agents_velociraptor(session)

229

backend/app/auth/routes/auth.py

+3 -3

@@ -22,7 +22,7 @@ from app.auth.schema.user import UserBaseResponse

22 from app.auth.services.universal import find_user

23 from app.auth.services.universal import select_all_users

24 from app.auth.utils import AuthHandler

25 -from app.db.db_session import get_session

25 +from app.db.db_session import get_session, get_db

26 from app.db.db_session import session

28 ACCESS_TOKEN_EXPIRE_MINUTES = 1440

@@ -54,7 +54,7 @@ async def refresh_token(current_user: User = Depends(auth_handler.get_current_us

56 @auth_router.post("/register", response_model=UserResponse, status_code=201, description="Register new user")

57 -async def register(user: UserInput, session: AsyncSession = Depends(get_session)):

57 +async def register(user: UserInput, session: AsyncSession = Depends(get_db)):

58 # users = select_all_users()

59 users = await select_all_users()

60 if any(x.username == user.username for x in users):

@@ -82,7 +82,7 @@ async def login(user: UserLogin):

83 # Get all users

84 @auth_router.get("/users", response_model=UserBaseResponse, description="Get all users")

85 -async def get_users(session: AsyncSession = Depends(get_session)):

85 +async def get_users(session: AsyncSession = Depends(get_db)):

86 # users = select_all_users()

87 users = await select_all_users()

88 return UserBaseResponse(users=users, message="Users retrieved successfully", success=True)

backend/app/connectors/routes.py

+6 -6

@@ -17,7 +17,7 @@ from app.connectors.schema import ConnectorsListResponse

17 from app.connectors.schema import UpdateConnector

18 from app.connectors.schema import VerifyConnectorResponse

19 from app.connectors.services import ConnectorServices

20 -from app.db.db_session import get_session

20 +from app.db.db_session import get_session, get_db

22 connector_router = APIRouter()

@@ -28,7 +28,7 @@ connector_router = APIRouter()

28 description="Fetch all available connectors",

29 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

30 )

31 -async def get_connectors(session: AsyncSession = Depends(get_session)) -> ConnectorsListResponse:

31 +async def get_connectors(session: AsyncSession = Depends(get_db)) -> ConnectorsListResponse:

32 """

33 Fetch all available connectors from the database.

@@ -55,7 +55,7 @@ async def get_connectors(session: AsyncSession = Depends(get_session)) -> Connec

55 description="Fetch a specific connector",

56 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

57 )

58 -async def get_connector(connector_id: int, session: AsyncSession = Depends(get_session)) -> Union[ConnectorResponse, HTTPException]:

58 +async def get_connector(connector_id: int, session: AsyncSession = Depends(get_db)) -> Union[ConnectorResponse, HTTPException]:

59 """

60 Fetch a specific connector by its ID.

@@ -85,7 +85,7 @@ async def get_connector(connector_id: int, session: AsyncSession = Depends(get_s

85 )

86 async def verify_connector(

87 connector_id: int,

88 - session: AsyncSession = Depends(get_session),

88 + session: AsyncSession = Depends(get_db),

89 ) -> Union[VerifyConnectorResponse, HTTPException]:

90 """

91 Verify a connector by its ID.

@@ -118,7 +118,7 @@ async def verify_connector(

118 async def update_connector(

119 connector_id: int,

120 connector: UpdateConnector,

121 - session: AsyncSession = Depends(get_session),

121 + session: AsyncSession = Depends(get_db),

122 ) -> ConnectorListResponse:

123 """

124 Update a connector by its ID.

@@ -147,7 +147,7 @@ async def update_connector(

147 description="Upload a YAML file for a specific connector",

148 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

149 )

150 -async def upload_yaml_file(connector_id: int, file: UploadFile = File(...), session: AsyncSession = Depends(get_session)) -> dict:

150 +async def upload_yaml_file(connector_id: int, file: UploadFile = File(...), session: AsyncSession = Depends(get_db)) -> dict:

151 """

152 Upload a YAML file for a specific connector ID.

153

backend/app/connectors/sublime/routes/alerts.py

+3 -3

@@ -10,13 +10,13 @@ from app.connectors.sublime.schema.alerts import AlertResponseBody

10 from app.connectors.sublime.schema.alerts import SublimeAlertsResponse

11 from app.connectors.sublime.services.alerts import collect_alerts

12 from app.connectors.sublime.services.alerts import store_sublime_alert

13 -from app.db.db_session import get_session

13 +from app.db.db_session import get_session, get_db

15 sublime_alerts_router = APIRouter()

18 @sublime_alerts_router.post("/alert", description="Receive alert from Sublime and store it in the database")

19 -async def receive_sublime_alert(alert_request_body: AlertRequestBody, session: AsyncSession = Depends(get_session)) -> AlertResponseBody:

19 +async def receive_sublime_alert(alert_request_body: AlertRequestBody, session: AsyncSession = Depends(get_db)) -> AlertResponseBody:

20 """

21 Endpoint to store alert in the `sublimealerts` table.

22 Invoked by the Sublime alert webhook which is configured in the Sublime UI.

@@ -34,7 +34,7 @@ async def receive_sublime_alert(alert_request_body: AlertRequestBody, session: A

34 description="Get all alerts",

35 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

36 )

37 -async def get_sublime_alerts(session: AsyncSession = Depends(get_session)) -> SublimeAlertsResponse:

37 +async def get_sublime_alerts(session: AsyncSession = Depends(get_db)) -> SublimeAlertsResponse:

38 """

39 Endpoint to retrieve alerts from the `sublimealerts` table.

backend/app/connectors/velociraptor/routes/artifacts.py

+5 -5

@@ -22,7 +22,7 @@ from app.connectors.velociraptor.services.artifacts import get_artifacts

22 from app.connectors.velociraptor.services.artifacts import quarantine_host

23 from app.connectors.velociraptor.services.artifacts import run_artifact_collection

24 from app.connectors.velociraptor.services.artifacts import run_remote_command

25 -from app.db.db_session import get_session

25 +from app.db.db_session import get_session, get_db

26 from app.db.universal_models import Agents

28 # App specific imports

@@ -144,7 +144,7 @@ async def get_all_artifacts_for_os_prefix(os_prefix: str = Depends(verify_os_pre

144 description="Get all artifacts for a specific host's OS prefix",

145 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

146 )

147 -async def get_all_artifacts_for_hostname(hostname: str, session: AsyncSession = Depends(get_session)) -> ArtifactsResponse:

147 +async def get_all_artifacts_for_hostname(hostname: str, session: AsyncSession = Depends(get_db)) -> ArtifactsResponse:

148 logger.info(f"Fetching all artifacts for hostname {hostname}")

149

150 # Asynchronous query to find the agent

@@ -198,7 +198,7 @@ async def get_all_artifacts_for_hostname(hostname: str, session: AsyncSession =

198 )

199 async def collect_artifact(

200 collect_artifact_body: CollectArtifactBody,

201 - session: AsyncSession = Depends(get_session),

201 + session: AsyncSession = Depends(get_db),

202 ) -> CollectArtifactResponse:

203 logger.info(f"Received request to collect artifact {collect_artifact_body}")

204 result = await get_all_artifacts_for_hostname(collect_artifact_body.hostname, session)

@@ -222,7 +222,7 @@ async def collect_artifact(

222 description="Run a remote command",

223 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

224 )

225 -async def run_command(run_command_body: RunCommandBody, session: AsyncSession = Depends(get_session)) -> RunCommandResponse:

225 +async def run_command(run_command_body: RunCommandBody, session: AsyncSession = Depends(get_db)) -> RunCommandResponse:

226 logger.info(f"Received request to run command {run_command_body}")

227 result = await get_all_artifacts_for_hostname(run_command_body.hostname, session)

228 artifact_names = [artifact.name for artifact in result.artifacts]

@@ -243,7 +243,7 @@ async def run_command(run_command_body: RunCommandBody, session: AsyncSession =

243 description="Quarantine a host",

244 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

245 )

246 -async def quarantine(quarantine_body: QuarantineBody, session: AsyncSession = Depends(get_session)) -> QuarantineResponse:

246 +async def quarantine(quarantine_body: QuarantineBody, session: AsyncSession = Depends(get_db)) -> QuarantineResponse:

247 logger.info(f"Received request to quarantine host {quarantine_body}")

248 result = await get_all_artifacts_for_hostname(quarantine_body.hostname, session)

249 artifact_names = [artifact.name for artifact in result.artifacts]

backend/app/connectors/wazuh_manager/routes/rules.py

+4 -4

@@ -19,7 +19,7 @@ from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse

19 from app.connectors.wazuh_manager.services.rules import disable_rule

20 from app.connectors.wazuh_manager.services.rules import enable_rule

21 from app.connectors.wazuh_manager.services.rules import exclude_rule

22 -from app.db.db_session import get_session

22 +from app.db.db_session import get_session, get_db

24 NEW_LEVEL = "1"

25 wazuh_manager_rules_router = APIRouter()

@@ -36,7 +36,7 @@ def query_disabled_rule(rule_id: str):

36 description="Get all disabled rules",

37 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

38 )

39 -async def get_disabled_rules(session: AsyncSession = Depends(get_session)) -> AllDisabledRuleResponse:

39 +async def get_disabled_rules(session: AsyncSession = Depends(get_db)) -> AllDisabledRuleResponse:

40 result = await session.execute(select(DisabledRule))

41 disabled_rules = result.scalars().all()

42 return AllDisabledRuleResponse(disabled_rules=disabled_rules, success=True, message="Successfully fetched all disabled rules")

@@ -77,7 +77,7 @@ async def get_disabled_rules(session: AsyncSession = Depends(get_session)) -> Al

77 )

78 async def disable_wazuh_rule(

79 rule: RuleDisable,

80 - session: AsyncSession = Depends(get_session),

80 + session: AsyncSession = Depends(get_db),

81 username: str = Depends(AuthHandler().get_current_user),

82 ) -> RuleDisableResponse:

83 # Asynchronously check if the rule is already disabled

@@ -131,7 +131,7 @@ async def disable_wazuh_rule(

131 description="Enable a Wazuh Rule",

132 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

133 )

134 -async def enable_wazuh_rule(rule: RuleEnable, session: AsyncSession = Depends(get_session)) -> RuleEnableResponse:

134 +async def enable_wazuh_rule(rule: RuleEnable, session: AsyncSession = Depends(get_db)) -> RuleEnableResponse:

135 # Asynchronously fetch the disabled rule

136 logger.info(f"rule: {rule}")

137 result = await session.execute(select(DisabledRule).where(DisabledRule.rule_id == rule.rule_id))

backend/app/customer_provisioning/routes/decommission.py

+4 -4

@@ -10,7 +10,7 @@ from sqlalchemy.future import select

10 from app.auth.utils import AuthHandler

11 from app.customer_provisioning.schema.decommission import DecommissionCustomerResponse

12 from app.customer_provisioning.services.decommission import decomission_wazuh_customer

13 -from app.db.db_session import get_session

13 +from app.db.db_session import get_session, get_db

14 from app.db.universal_models import Customers

15 from app.db.universal_models import CustomersMeta

@@ -20,13 +20,13 @@ from app.db.universal_models import CustomersMeta

20 customer_decommissioning_router = APIRouter()

23 -async def check_customermeta_exists(customer_name: str, session: AsyncSession = Depends(get_session)) -> CustomersMeta:

23 +async def check_customermeta_exists(customer_name: str, session: AsyncSession = Depends(get_db)) -> CustomersMeta:

24 """

25 Check if a customer exists in the database.

27 Args:

28 customer_name (str): The name of the customer to check.

29 - session (AsyncSession, optional): The database session. Defaults to Depends(get_session).

29 + session (AsyncSession, optional): The database session. Defaults to Depends(get_db).

31 Returns:

32 CustomersMeta: The customer object if found.

@@ -52,7 +52,7 @@ async def check_customermeta_exists(customer_name: str, session: AsyncSession =

52 )

53 async def decommission_customer_route(

54 _customer: CustomersMeta = Depends(check_customermeta_exists),

55 - session: AsyncSession = Depends(get_session),

55 + session: AsyncSession = Depends(get_db),

56 ):

57 logger.info("Decommissioning customer")

58 customer_decommission = await decomission_wazuh_customer(_customer, session=session)

backend/app/customer_provisioning/routes/provision.py

+22 -22

@@ -17,7 +17,7 @@ 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 ProvisionNewCustomer

19 from app.customer_provisioning.services.provision import provision_wazuh_customer

20 -from app.db.db_session import get_session

20 +from app.db.db_session import get_session, get_db

21 from app.db.universal_models import Customers

22 from app.db.universal_models import CustomersMeta

@@ -40,7 +40,7 @@ def get_available_subscriptions():

40 raise HTTPException(status_code=500, detail=f"Error getting available subscriptions: {e}")

43 -async def check_customer_exists(customer_name: str, session: AsyncSession = Depends(get_session)) -> Customers:

43 +async def check_customer_exists(customer_name: str, session: AsyncSession = Depends(get_db)) -> Customers:

44 logger.info(f"Checking if customer {customer_name} exists")

45 result = await session.execute(select(Customers).filter(Customers.customer_name == customer_name))

46 customer = result.scalars().first()

@@ -51,25 +51,6 @@ async def check_customer_exists(customer_name: str, session: AsyncSession = Depe

51 return customer

54 -# Get the customermeta based on the customer name

55 -@customer_provisioning_router.get(

56 - "/provision/{customer_name}",

57 - response_model=CustomersMetaResponse,

58 - description="Get Customer Meta",

59 - dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

60 -)

61 -async def get_customer_meta(customer_name: str, session: AsyncSession = Depends(get_session)):

62 - logger.info(f"Getting customer meta for customer {customer_name}")

63 - result = await session.execute(select(CustomersMeta).filter(CustomersMeta.customer_name == customer_name))

64 - customer_meta = result.scalars().first()

65 -

66 - if not customer_meta:

67 - raise HTTPException(

68 - status_code=404, detail=f"Customer meta not found for customer: {customer_name}. Please provision the customer first.",

69 - )

70 -

71 - return CustomersMetaResponse(message="Customer meta retrieved successfully", success=True, customer_meta=customer_meta)

72 -

55 @customer_provisioning_router.post(

56 "/provision",

@@ -80,7 +61,7 @@ async def get_customer_meta(customer_name: str, session: AsyncSession = Depends(

61 async def provision_customer_route(

62 request: ProvisionNewCustomer = Body(...),

63 _customer: Customers = Depends(check_customer_exists),

83 - session: AsyncSession = Depends(get_session),

64 + session: AsyncSession = Depends(get_db),

65 ):

66 logger.info("Provisioning new customer")

67 customer_provision = await provision_wazuh_customer(request, session=session)

@@ -113,3 +94,22 @@ async def get_subscriptions_route():

94 success=True,

95 message="Subscriptions retrieved successfully",

96 )

97 +

98 +# Get the customermeta based on the customer name

99 +@customer_provisioning_router.get(

100 + "/provision/{customer_name}",

101 + response_model=CustomersMetaResponse,

102 + description="Get Customer Meta",

103 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

104 +)

105 +async def get_customer_meta(customer_name: str, session: AsyncSession = Depends(get_db)):

106 + logger.info(f"Getting customer meta for customer {customer_name}")

107 + result = await session.execute(select(CustomersMeta).filter(CustomersMeta.customer_name == customer_name))

108 + customer_meta = result.scalars().first()

109 +

110 + if not customer_meta:

111 + raise HTTPException(

112 + status_code=404, detail=f"Customer meta not found for customer: {customer_name}. Please provision the customer first.",

113 + )

114 +

115 + return CustomersMetaResponse(message="Customer meta retrieved successfully", success=True, customer_meta=customer_meta)

backend/app/customers/routes/customers.py

+18 -14

@@ -19,7 +19,7 @@ from app.customers.schema.customers import CustomerMetaResponse

19 from app.customers.schema.customers import CustomerRequestBody

20 from app.customers.schema.customers import CustomerResponse

21 from app.customers.schema.customers import CustomersResponse

22 -from app.db.db_session import get_session

22 +from app.db.db_session import get_session, get_db

23 from app.db.db_session import session

24 from app.db.universal_models import Agents

25 from app.db.universal_models import Customers

@@ -53,7 +53,7 @@ async def verify_unique_customer_code(session: AsyncSession, customer: CustomerR

53 description="Create a new customer",

54 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

55 )

56 -async def create_customer(customer: CustomerRequestBody, session: AsyncSession = Depends(get_session)) -> CustomerResponse:

56 +async def create_customer(customer: CustomerRequestBody, session: AsyncSession = Depends(get_db)) -> CustomerResponse:

57 await verify_unique_customer_code(session, customer)

58 logger.info(f"Creating new customer: {customer}")

59 new_customer = Customers(**customer.dict())

@@ -68,7 +68,7 @@ async def create_customer(customer: CustomerRequestBody, session: AsyncSession =

68 description="Get all customers",

69 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

70 )

71 -async def get_customers(session: AsyncSession = Depends(get_session)) -> CustomersResponse:

71 +async def get_customers(session: AsyncSession = Depends(get_db)) -> CustomersResponse:

72 logger.info("Fetching all customers")

74 # Asynchronous query to fetch all customers

@@ -86,7 +86,7 @@ async def get_customers(session: AsyncSession = Depends(get_session)) -> Custome

86 description="Get customer by customer_code",

87 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

88 )

89 -async def get_customer(customer_code: str, session: AsyncSession = Depends(get_session)) -> CustomerResponse:

89 +async def get_customer(customer_code: str, session: AsyncSession = Depends(get_db)) -> CustomerResponse:

90 logger.info(f"Fetching customer with customer_code: {customer_code}")

92 # Asynchronous query to fetch customer

@@ -110,7 +110,7 @@ async def get_customer(customer_code: str, session: AsyncSession = Depends(get_s

110 async def update_customer(

111 customer_code: str,

112 customer: CustomerRequestBody,

113 - session: AsyncSession = Depends(get_session),

113 + session: AsyncSession = Depends(get_db),

114 ) -> CustomerResponse:

115 logger.info(f"Updating customer with customer_code: {customer_code}")

116

@@ -141,7 +141,7 @@ async def update_customer(

141 description="Delete customer by customer_code",

142 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

143 )

144 -async def delete_customer(customer_code: str, session: AsyncSession = Depends(get_session)) -> CustomerResponse:

144 +async def delete_customer(customer_code: str, session: AsyncSession = Depends(get_db)) -> CustomerResponse:

145 logger.info(f"Deleting customer with customer_code: {customer_code}")

146

147 result = await session.execute(select(Customers).filter(Customers.customer_code == customer_code))

@@ -172,11 +172,12 @@ async def delete_customer(customer_code: str, session: AsyncSession = Depends(ge

172 response_model=CustomerMetaResponse,

173 description="Add new customer meta",

174 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

175 + deprecated=True,

176 )

177 async def add_customer_meta(

178 customer_code: str,

179 customer_meta: CustomerMetaRequestBody,

179 - session: AsyncSession = Depends(get_session),

180 + session: AsyncSession = Depends(get_db),

181 ) -> CustomerMetaResponse:

182 logger.info(f"Adding new customer meta: {customer_meta}")

183

@@ -202,8 +203,9 @@ async def add_customer_meta(

203 response_model=CustomerMetaResponse,

204 description="Get customer meta by customer_code",

205 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

206 + deprecated=True,

207 )

206 -async def get_customer_meta(customer_code: str, session: AsyncSession = Depends(get_session)) -> CustomerMetaResponse:

208 +async def get_customer_meta(customer_code: str, session: AsyncSession = Depends(get_db)) -> CustomerMetaResponse:

209 logger.info(f"Fetching customer meta with customer_code: {customer_code}")

210

211 result = await session.execute(select(CustomersMeta).filter(CustomersMeta.customer_code == customer_code))

@@ -226,11 +228,12 @@ async def get_customer_meta(customer_code: str, session: AsyncSession = Depends(

228 response_model=CustomerMetaResponse,

229 description="Update customer meta by customer_code",

230 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

231 + deprecated=True,

232 )

233 async def update_customer_meta(

234 customer_code: str,

235 customer_meta: CustomerMetaRequestBody,

233 - session: AsyncSession = Depends(get_session),

236 + session: AsyncSession = Depends(get_db),

237 ) -> CustomerMetaResponse:

238 logger.info(f"Updating customer meta with customer_code: {customer_code}")

239

@@ -260,8 +263,9 @@ async def update_customer_meta(

263 response_model=CustomerMetaResponse,

264 description="Delete customer meta by customer_code",

265 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

266 + deprecated=True,

267 )

264 -async def delete_customer_meta(customer_code: str, session: AsyncSession = Depends(get_session)) -> CustomerMetaResponse:

268 +async def delete_customer_meta(customer_code: str, session: AsyncSession = Depends(get_db)) -> CustomerMetaResponse:

269 logger.info(f"Deleting customer meta with customer_code: {customer_code}")

270

271 result = await session.execute(select(CustomersMeta).filter(CustomersMeta.customer_code == customer_code))

@@ -292,7 +296,7 @@ async def delete_customer_meta(customer_code: str, session: AsyncSession = Depen

296 description="Get customer and customer meta by customer_code",

297 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

298 )

295 -async def get_customer_full(customer_code: str, session: AsyncSession = Depends(get_session)) -> CustomerFullResponse:

299 +async def get_customer_full(customer_code: str, session: AsyncSession = Depends(get_db)) -> CustomerFullResponse:

300 logger.info(f"Fetching customer and customer meta with customer_code: {customer_code}")

301

302 customer_result = await session.execute(select(Customers).filter(Customers.customer_code == customer_code))

@@ -323,7 +327,7 @@ async def get_customer_full(customer_code: str, session: AsyncSession = Depends(

327 description="Get agents for the given customer_code",

328 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

329 )

326 -async def get_agents(customer_code: str, session: AsyncSession = Depends(get_session)) -> AgentsResponse:

330 +async def get_agents(customer_code: str, session: AsyncSession = Depends(get_db)) -> AgentsResponse:

331 logger.info(f"Fetching agents for customer_code: {customer_code}")

332

333 # Check if the customer exists

@@ -349,7 +353,7 @@ async def get_agents(customer_code: str, session: AsyncSession = Depends(get_ses

353 )

354 async def get_wazuh_agents_healthcheck(

355 customer_code: str,

352 - session: AsyncSession = Depends(get_session),

356 + session: AsyncSession = Depends(get_db),

357 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

358 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

359 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

@@ -381,7 +385,7 @@ async def get_wazuh_agents_healthcheck(

385 )

386 async def get_velociraptor_agents_healthcheck(

387 customer_code: str,

384 - session: AsyncSession = Depends(get_session),

388 + session: AsyncSession = Depends(get_db),

389 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

390 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

391 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

backend/app/db/db_session.py

+11 -1

@@ -66,7 +66,17 @@ def get_sync_db_session():

66 logger.info("Closing sync DB session")

67 session.close()

69 +# ! OLD CODE RELATING TO THE SESSION NOT CLOSING ! #

70 +# async def get_session():

71 +# async with get_db_session() as session:

72 +# return session

74 +# ! NEW CODE RELATING TO THE SESSION NOT CLOSING ! #

75 +@asynccontextmanager

76 async def get_session():

77 async with get_db_session() as session:

72 - return session

78 + yield session

79 +

80 +async def get_db():

81 + async with get_session() as session:

82 + yield session

backend/app/healthchecks/agents/routes/agents.py

+6 -6

@@ -9,7 +9,7 @@ from sqlalchemy.future import select

9 from starlette.status import HTTP_401_UNAUTHORIZED

11 from app.auth.utils import AuthHandler

12 -from app.db.db_session import get_session

12 +from app.db.db_session import get_session, get_db

13 from app.db.db_session import session

14 from app.db.universal_models import Agents

15 from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse

@@ -37,7 +37,7 @@ def verify_admin(user):

37 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

38 )

39 async def get_wazuh_agent_healthcheck(

40 - session: AsyncSession = Depends(get_session),

40 + session: AsyncSession = Depends(get_db),

41 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

42 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

43 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

@@ -58,7 +58,7 @@ async def get_wazuh_agent_healthcheck(

58 )

59 async def get_wazuh_agent_healthcheck_by_agent_id(

60 agent_id: str,

61 - session: AsyncSession = Depends(get_session),

61 + session: AsyncSession = Depends(get_db),

62 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

63 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

64 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

@@ -81,7 +81,7 @@ async def get_wazuh_agent_healthcheck_by_agent_id(

81 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

82 )

83 async def get_velociraptor_agent_healthcheck(

84 - session: AsyncSession = Depends(get_session),

84 + session: AsyncSession = Depends(get_db),

85 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

86 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

87 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

@@ -102,7 +102,7 @@ async def get_velociraptor_agent_healthcheck(

102 )

103 async def get_velociraptor_agent_healthcheck_by_agent_id(

104 agent_id: str,

105 - session: AsyncSession = Depends(get_session),

105 + session: AsyncSession = Depends(get_db),

106 minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),

107 hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),

108 days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),

@@ -123,7 +123,7 @@ async def get_velociraptor_agent_healthcheck_by_agent_id(

123 description="Get host logs",

124 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin", "analyst"])],

125 )

126 -async def get_host_logs(body: HostLogsSearchBody, session: AsyncSession = Depends(get_session)) -> HostLogsSearchResponse:

126 +async def get_host_logs(body: HostLogsSearchBody, session: AsyncSession = Depends(get_db)) -> HostLogsSearchResponse:

127 logger.info(f"Received request to get host logs for {body.agent_name}")

128

129 # Asynchronously verify the agent exists

backend/app/integrations/alert_creation/general/routes/alert.py

+2 -2

@@ -5,7 +5,7 @@ from loguru import logger

5 from sqlalchemy.ext.asyncio import AsyncSession

6 from sqlalchemy.future import select

8 -from app.db.db_session import get_session

8 +from app.db.db_session import get_session, get_db

9 from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest

10 from app.integrations.alert_creation.general.schema.alert import CreateAlertResponse

11 from app.integrations.alert_creation.general.services.alert import create_alert

@@ -51,7 +51,7 @@ async def is_customer_code_valid(create_alert_request: CreateAlertRequest, sessi

51 )

52 async def create_general_alert(

53 create_alert_request: CreateAlertRequest,

54 - session: AsyncSession = Depends(get_session),

54 + session: AsyncSession = Depends(get_db),

55 ):

56 logger.info(f"create_alert_request: {create_alert_request.dict()}")

backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py

+7 -7

@@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession

9 from sqlalchemy.future import select

10 from sqlalchemy.orm import joinedload

12 -from app.db.db_session import get_session

12 +from app.db.db_session import get_session, get_db

13 from app.integrations.alert_creation_settings.models.alert_creation_settings import (

14 AlertCreationEventConfig,

15 )

@@ -46,7 +46,7 @@ alert_creation_settings_router = APIRouter()

46 )

47 async def get_customer_event_configs(

48 customer_code: str,

49 - session: AsyncSession = Depends(get_session),

49 + session: AsyncSession = Depends(get_db),

50 ):

51 event_configs = await get_customer_alert_event_configs(customer_code, session)

@@ -63,7 +63,7 @@ async def get_customer_event_configs(

63 )

64 async def create_alert_creation_settings(

65 alert_creation_settings: AlertCreationSettingsCreate,

66 - session: AsyncSession = Depends(get_session),

66 + session: AsyncSession = Depends(get_db),

67 ):

68 logger.info(f"alert_creation_settings: {alert_creation_settings.dict()}")

@@ -100,7 +100,7 @@ async def create_alert_creation_settings(

100 )

101 async def get_alert_creation_settings(

102 customer_name: str,

103 - session: AsyncSession = Depends(get_session),

103 + session: AsyncSession = Depends(get_db),

104 ):

105 result = await session.execute(

106 select(AlertCreationSettings)

@@ -123,7 +123,7 @@ async def get_alert_creation_settings(

123 async def add_event_order(

124 customer_name: str,

125 event_order: EventOrderCreate,

126 - session: AsyncSession = Depends(get_session),

126 + session: AsyncSession = Depends(get_db),

127 ):

128 result = await session.execute(

129 select(AlertCreationSettings)

@@ -161,7 +161,7 @@ async def add_event_order(

161 async def update_event_orders(

162 customer_name: str,

163 event_orders: List[EventOrderCreate],

164 - session: AsyncSession = Depends(get_session),

164 + session: AsyncSession = Depends(get_db),

165 ):

166 result = await session.execute(

167 select(AlertCreationSettings)

@@ -200,7 +200,7 @@ async def update_event_orders(

200 async def delete_event_order(

201 customer_name: str,

202 order_label: str,

203 - session: AsyncSession = Depends(get_session),

203 + session: AsyncSession = Depends(get_db),

204 ):

205 result = await session.execute(

206 select(AlertCreationSettings)

backend/app/integrations/ask_socfortress/routes/ask_socfortress.py

+3 -3

@@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncSession

8 from sqlalchemy.future import select

10 from app.auth.utils import AuthHandler

11 -from app.db.db_session import get_session

11 +from app.db.db_session import get_session, get_db

12 from app.db.universal_models import Customers

13 from app.db.universal_models import CustomersMeta

14 from app.integrations.ask_socfortress.schema.ask_socfortress import (

@@ -30,7 +30,7 @@ from app.utils import get_connector_attribute

30 ask_socfortress_router = APIRouter()

33 -async def ensure_api_key_exists(session: AsyncSession = Depends(get_session)) -> bool:

33 +async def ensure_api_key_exists(session: AsyncSession = Depends(get_db)) -> bool:

34 """

35 Ensures that the Ask SocFortress API key exists in the database.

@@ -59,7 +59,7 @@ async def ensure_api_key_exists(session: AsyncSession = Depends(get_session)) ->

59 )

60 async def ask_socfortress_sigma(

61 alert: AskSocfortressRequest,

62 - session: AsyncSession = Depends(get_session),

62 + session: AsyncSession = Depends(get_db),

63 _key_exists: bool = Depends(ensure_api_key_exists),

64 ):

65 logger.info("Running Ask SOCFortress Sigma lookup.")

backend/app/threat_intel/routes/socfortress.py

+3 -3

@@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncSession

8 from sqlalchemy.future import select

10 from app.auth.utils import AuthHandler

11 -from app.db.db_session import get_session

11 +from app.db.db_session import get_session, get_db

12 from app.db.universal_models import Customers

13 from app.db.universal_models import CustomersMeta

14 from app.threat_intel.schema.socfortress import IoCResponse

@@ -21,7 +21,7 @@ from app.utils import get_connector_attribute

21 threat_intel_socfortress_router = APIRouter()

24 -async def ensure_api_key_exists(session: AsyncSession = Depends(get_session)) -> bool:

24 +async def ensure_api_key_exists(session: AsyncSession = Depends(get_db)) -> bool:

25 """

26 Ensures that the SocFortress API key exists in the database.

@@ -50,7 +50,7 @@ async def ensure_api_key_exists(session: AsyncSession = Depends(get_session)) ->

50 )

51 async def threat_intel_socfortress(

52 request: SocfortressThreatIntelRequest,

53 - session: AsyncSession = Depends(get_session),

53 + session: AsyncSession = Depends(get_db),

54 _key_exists: bool = Depends(ensure_api_key_exists),

55 ):

56 logger.info("Running SOCFortress Threat Intel")

backend/app/utils.py

+7 -7

@@ -29,7 +29,7 @@ from app.db.all_models import Connectors

29 from app.db.db_session import Session

30 from app.db.db_session import engine

31 from app.db.db_session import get_db_session

32 -from app.db.db_session import get_session

32 +from app.db.db_session import get_session, get_db

33 from app.db.universal_models import LogEntry

34 from app.integrations.alert_creation_settings.models.alert_creation_settings import (

35 AlertCreationEventConfig,

@@ -267,7 +267,7 @@ logs_router = APIRouter()

267 description="Fetch all logs",

268 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

269 )

270 -async def get_logs(session: AsyncSession = Depends(get_session)) -> LogsResponse:

270 +async def get_logs(session: AsyncSession = Depends(get_db)) -> LogsResponse:

271 """

272 Fetch all logs from the database.

273

@@ -296,7 +296,7 @@ async def get_logs(session: AsyncSession = Depends(get_session)) -> LogsResponse

296 description="Fetch logs by user ID",

297 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

298 )

299 -async def get_logs_by_user_id(user_id: int, session: AsyncSession = Depends(get_session)) -> LogsResponse:

299 +async def get_logs_by_user_id(user_id: int, session: AsyncSession = Depends(get_db)) -> LogsResponse:

300 """

301 Fetch all logs from the database where the user_id matches the provided user_id.

302

@@ -327,7 +327,7 @@ async def get_logs_by_user_id(user_id: int, session: AsyncSession = Depends(get_

327 description="Fetch logs by time range",

328 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

329 )

330 -async def get_logs_by_time_range(time_range: TimeRangeModel, session: AsyncSession = Depends(get_session)) -> LogsResponse:

330 +async def get_logs_by_time_range(time_range: TimeRangeModel, session: AsyncSession = Depends(get_db)) -> LogsResponse:

331 """

332 Fetch all logs from the database where the timestamp is within the provided time range.

333

@@ -367,7 +367,7 @@ async def get_logs_by_time_range(time_range: TimeRangeModel, session: AsyncSessi

367 )

368 async def get_logs_by_event_type(

369 event_type: EventType,

370 - session: AsyncSession = Depends(get_session),

370 + session: AsyncSession = Depends(get_db),

371 ) -> LogsResponse: # Update this line to use the new model

372 """

373 Fetch all logs from the database where the event_type matches the provided event_type.

@@ -399,7 +399,7 @@ async def get_logs_by_event_type(

399 description="Purge all logs",

400 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

401 )

402 -async def purge_logs(session: AsyncSession = Depends(get_session)) -> LogsResponse: # Update this line to use the new model

402 +async def purge_logs(session: AsyncSession = Depends(get_db)) -> LogsResponse: # Update this line to use the new model

403 """

404 Purge all logs from the database.

405

@@ -429,7 +429,7 @@ async def purge_logs(session: AsyncSession = Depends(get_session)) -> LogsRespon

429 description="Purge logs by time range",

430 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],

431 )

432 -async def purge_logs_by_time_range(time_range: TimeRangeModel, session: AsyncSession = Depends(get_session)) -> LogsResponse:

432 +async def purge_logs_by_time_range(time_range: TimeRangeModel, session: AsyncSession = Depends(get_db)) -> LogsResponse:

433 """

434 Purge all logs from the database where the timestamp is within the provided time range.

435