Refine OS artifact filtering and enhance error handling for large messages (#353)
taylor_socfortress committed
Dec 5, 2024 at 12:47 UTC
595d0a1f5c0d867d9e6532bd98f4a9b7a4daa695
2 files changed
+8
-3
backend/app/threat_intel/routes/socfortress.py
+3
-3
@@ -293,9 +293,9 @@ async def fetch_agent_os(agent_id: str, session: AsyncSession) -> str:
293
294
295
async def filter_artifacts_by_os(artifacts, os):
296
- # Only get the artifacts that start with `Windows.`, `Linux.`, `MacOS.`, or `Generic.`
297
- os_artifacts = ["Windows", "Linux", "MacOS", "Generic"]
298
- os_artifacts = [os_artifact for os_artifact in os_artifacts if os_artifact in os or os_artifact == "Generic"]
296
+ # Only get the artifacts that start with `Windows.`, `Linux.`, or `MacOS.`
297
+ os_artifacts = ["Windows", "Linux", "MacOS"]
298
+ os_artifacts = [os_artifact for os_artifact in os_artifacts if os_artifact in os]
299
300
# Artifacts to be stripped out
301
excluded_artifacts = {
backend/app/threat_intel/services/socfortress.py
+5
@@ -444,6 +444,11 @@ async def get_velociraptor_artifact_recommendation_response(
444
status_code=429,
445
detail="Message is too large. Please try again with a smaller message.",
446
)
447
+ elif "too large" in response_data.get("message", ""):
448
+ raise HTTPException(
449
+ status_code=429,
450
+ detail="Message is too large. Please try again with a smaller message.",
451
+ )
452
453
return VelociraptorArtifactRecommendationResponse(**response_data)
454