feat: enhance YAML file upload endpoint to accept connector name or ID (#424)
taylor_socfortress committed
Mar 13, 2025 at 15:22 UTC
ed35bcc3bc4da6f6c95cad1f5f208de7d00a3d03
2 files changed
+121
-12
backend/app/connectors/routes.py
+92
-12
@@ -187,39 +187,119 @@ async def update_connector(
187
)
188
189
190
+# @connector_router.post(
191
+# "/upload/{connector_id}",
192
+# description="Upload a YAML file for a specific connector",
193
+# dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
194
+# )
195
+# async def upload_yaml_file(
196
+# connector_id: int,
197
+# file: UploadFile = File(...),
198
+# session: AsyncSession = Depends(get_db),
199
+# ) -> dict:
200
+# """
201
+# Upload a YAML file for a specific connector ID.
202
+
203
+# This endpoint allows you to upload a `.yaml` file for a specific connector
204
+# identified by `connector_id`.
205
+
206
+# Args:
207
+# connector_id (int): The unique identifier for the connector.
208
+# file (UploadFile): The `.yaml` file to be uploaded.
209
+
210
+# Returns:
211
+# dict: A dictionary with a success message and other information.
212
+
213
+# Raises:
214
+# HTTPException: An exception with a 400 status code is raised if the file format is incorrect or connector ID is not 6.
215
+# """
216
+# if connector_id not in [5, 6]:
217
+# raise HTTPException(
218
+# status_code=400,
219
+# detail="Only the Velociraptor or another specific connector is allowed for YAML file uploads.",
220
+# )
221
+# if not file.filename.endswith(".yaml"):
222
+# raise HTTPException(status_code=400, detail="Only .yaml files are allowed.")
223
+# try:
224
+# save_file_result = await ConnectorServices.save_file(file, connector_id, session=session)
225
+# if save_file_result:
226
+# await ConnectorServices.verify_connector_by_id(
227
+# connector_id,
228
+# session=session,
229
+# )
230
+# return {"success": True, "message": "File uploaded successfully"}
231
+# else:
232
+# raise HTTPException(status_code=500, detail="Failed to upload file")
233
+# except Exception as e:
234
+# logger.error(f"Failed to upload file: {e}")
235
+# raise HTTPException(status_code=500, detail="Failed to upload file")
236
+
237
@connector_router.post(
191
- "/upload/{connector_id}",
238
+ "/upload/{connector_identifier}",
239
description="Upload a YAML file for a specific connector",
240
dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
241
)
242
async def upload_yaml_file(
196
- connector_id: int,
243
+ connector_identifier: str,
244
file: UploadFile = File(...),
245
session: AsyncSession = Depends(get_db),
246
) -> dict:
247
"""
201
- Upload a YAML file for a specific connector ID.
248
+ Upload a YAML file for a specific connector identified by ID or name.
249
203
- This endpoint allows you to upload a `.yaml` file for a specific connector
204
- identified by `connector_id`.
250
+ This endpoint allows you to upload a `.yaml` file for a connector
251
+ identified by either its ID (number) or name (string).
252
253
Args:
207
- connector_id (int): The unique identifier for the connector.
254
+ connector_identifier (str): The unique identifier (ID or name) for the connector.
255
file (UploadFile): The `.yaml` file to be uploaded.
256
257
Returns:
258
dict: A dictionary with a success message and other information.
259
260
Raises:
214
- HTTPException: An exception with a 400 status code is raised if the file format is incorrect or connector ID is not 6.
261
+ HTTPException: An exception is raised if the file format is incorrect or connector is not found/supported.
262
"""
216
- if connector_id not in [5, 6]:
217
- raise HTTPException(
218
- status_code=400,
219
- detail="Only the Velociraptor or another specific connector is allowed for YAML file uploads.",
220
- )
263
+ # Try to parse as integer for connector_id
264
+ connector_id = None
265
+ try:
266
+ connector_id = int(connector_identifier)
267
+ # Check if this is a valid connector ID for YAML uploads
268
+ if connector_id not in [5, 6]:
269
+ raise HTTPException(
270
+ status_code=400,
271
+ detail="Only the Velociraptor or another specific connector is allowed for YAML file uploads.",
272
+ )
273
+ except ValueError:
274
+ # If not an integer, treat as connector name
275
+ connector_name = connector_identifier
276
+ try:
277
+ # Look up the connector ID from the name
278
+ connector = await ConnectorServices.fetch_connector_by_name(connector_name, session=session)
279
+ if not connector:
280
+ raise HTTPException(
281
+ status_code=404,
282
+ detail=f"No connector found with name: {connector_name}",
283
+ )
284
+ connector_id = connector["id"]
285
+ # Check if this is a valid connector for YAML uploads
286
+ if connector_id not in [5, 6]:
287
+ raise HTTPException(
288
+ status_code=400,
289
+ detail=f"Connector '{connector_name}' does not support YAML file uploads.",
290
+ )
291
+ except Exception as e:
292
+ logger.error(f"Error finding connector by name: {str(e)}")
293
+ raise HTTPException(
294
+ status_code=404,
295
+ detail=f"Could not find connector with name: {connector_name}",
296
+ )
297
+
298
+ # Check file format
299
if not file.filename.endswith(".yaml"):
300
raise HTTPException(status_code=400, detail="Only .yaml files are allowed.")
301
+
302
+ # Process the file upload
303
try:
304
save_file_result = await ConnectorServices.save_file(file, connector_id, session=session)
305
if save_file_result:
backend/app/connectors/services.py
+29
@@ -271,6 +271,34 @@ class ConnectorServices:
271
return ConnectorResponse.from_orm(connector)
272
return None
273
274
+ @classmethod
275
+ async def fetch_connector_by_name(
276
+ cls,
277
+ connector_name: str,
278
+ session: AsyncSession,
279
+ ) -> Optional[ConnectorResponse]:
280
+ """
281
+ Fetches a connector by its name from the database.
282
+
283
+ Args:
284
+ connector_name (str): The name of the connector to fetch.
285
+ session (AsyncSession): The database session.
286
+
287
+ Returns:
288
+ Optional[ConnectorResponse]: The fetched connector, or None if not found.
289
+ """
290
+ try:
291
+ result = await session.execute(
292
+ select(Connectors).where(Connectors.connector_name == connector_name),
293
+ )
294
+ connector = result.scalar_one_or_none()
295
+ if connector:
296
+ return ConnectorResponse.from_orm(connector)
297
+ return None
298
+ except Exception as e:
299
+ logger.error(f"Error fetching connector by name '{connector_name}': {e}")
300
+ return None
301
+
302
@classmethod
303
async def verify_connector_by_id(
304
cls,
@@ -441,3 +469,4 @@ class ConnectorServices:
469
return False
470
else:
471
return False
472
+