@cryptotaxi247 / CoPilot / commits / 06b6789a

Add ai analyst to comment (#354)

* Add alert_id to CreateAlertRequestRoute and implement current_time function for timestamping comments * feat: update threat_intel/ai api * refactor: rename request variables to ai_request and add comment creation for SOCFortress AI analysis --------- Co-authored-by: Davide Di Modica <webmaster.ddm@gmail.com>

taylor_socfortress committed Dec 5, 2024 at 15:41 UTC 06b6789a47ac9ff8eadaa9c98df536dc1d8f42e0
7 files changed +78 -21
backend/app/incidents/schema/incident_alert.py
+4
@@ -32,6 +32,10 @@ class CreateAlertRequestRoute(BaseModel):
32 None,
33 description="The agent id.",
34 )
35 + alert_id: Optional[int] = Field(
36 + None,
37 + description="The alert id.",
38 + )
39
40
41 class CreateAlertResponse(BaseModel):
backend/app/threat_intel/routes/socfortress.py
+43 -7
@@ -4,11 +4,13 @@ from fastapi import HTTPException
4 from fastapi import Security
5 from loguru import logger
6 from sqlalchemy.ext.asyncio import AsyncSession
7 -
7 +from datetime import datetime
8 from app.agents.services.status import get_agent_os_by_id
9 from app.auth.utils import AuthHandler
10 from app.connectors.velociraptor.services.artifacts import get_artifacts
11 from app.db.db_session import get_db
12 +from app.incidents.schema.db_operations import CommentCreate
13 +from app.incidents.services.db_operations import create_comment
14 from app.incidents.schema.incident_alert import CreateAlertRequest
15 from app.incidents.schema.incident_alert import CreateAlertRequestRoute
16 from app.incidents.schema.incident_alert import GenericAlertModel
@@ -198,6 +200,9 @@ async def process_name_intel_socfortress(
200 )
201 return socfortress_lookup
202
203 +async def current_time():
204 + return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
205 +
206
207 @threat_intel_socfortress_router.post(
208 "/ai/analyze-alert",
@@ -214,15 +219,26 @@ async def ai_anaylze_alert_socfortress(
219
220 assert isinstance(alert_details, GenericAlertModel)
221
217 - request = SocfortressAiAlertRequest(
222 + ai_request = SocfortressAiAlertRequest(
223 integration="SOCFORTRESS AI",
224 alert_payload=alert_details._source.dict(),
225 )
226
227 socfortress_lookup = await socfortress_ai_alert_lookup(
228 lincense_key=(await get_license(session)).license_key,
224 - request=request,
229 + request=ai_request,
230 + )
231 +
232 + await create_comment(
233 + CommentCreate(
234 + alert_id=request.alert_id,
235 + comment=f"SOCFortress AI Analysis: {socfortress_lookup.analysis}",
236 + user_name="admin",
237 + created_at=datetime.now(),
238 + ),
239 + db=session,
240 )
241 +
242 return socfortress_lookup
243
244
@@ -241,7 +257,7 @@ async def ai_wazuh_exclusion_rule_socfortress(
257
258 assert isinstance(alert_details, GenericAlertModel)
259
244 - request = SocfortressAiAlertRequest(
260 + ai_request = SocfortressAiAlertRequest(
261 integration="SOCFORTRESS AI",
262 alert_payload=alert_details._source.dict(),
263 )
@@ -250,7 +266,17 @@ async def ai_wazuh_exclusion_rule_socfortress(
266
267 socfortress_lookup = await socfortress_wazuh_exclusion_rule_lookup(
268 lincense_key=(await get_license(session)).license_key,
253 - request=request,
269 + request=ai_request,
270 + )
271 +
272 + await create_comment(
273 + CommentCreate(
274 + alert_id=request.alert_id,
275 + comment=f"SOCFortress AI Analysis: {socfortress_lookup.wazuh_exclusion_rule}/n/n{socfortress_lookup.wazuh_exclusion_rule_justification}",
276 + user_name="admin",
277 + created_at=datetime.now(),
278 + ),
279 + db=session,
280 )
281 return socfortress_lookup
282
@@ -344,7 +370,7 @@ async def ai_velociraptor_artifact_recommendation_socfortress(
370
371 os = await fetch_agent_os(request.agent_id, session)
372
347 - request = VelociraptorArtifactRecommendationRequest(
373 + ai_request = VelociraptorArtifactRecommendationRequest(
374 integration="SOCFORTRESS AI",
375 alert_payload=alert_payload._source.dict(),
376 os=os,
@@ -353,6 +379,16 @@ async def ai_velociraptor_artifact_recommendation_socfortress(
379
380 socfortress_lookup = await socfortress_velociraptor_recommendation_lookup(
381 lincense_key=(await get_license(session)).license_key,
356 - request=request,
382 + request=ai_request,
383 + )
384 +
385 + await create_comment(
386 + CommentCreate(
387 + alert_id=request.alert_id,
388 + comment=f"SOCFortress AI Analysis: {socfortress_lookup.artifact_recommendations}\n\n{socfortress_lookup.general_thoughts}",
389 + user_name="admin",
390 + created_at=datetime.now(),
391 + ),
392 + db=session,
393 )
394 return socfortress_lookup
frontend/src/api/endpoints/threatIntel.ts
+11 -6
@@ -32,36 +32,41 @@ export default {
32 body
33 )
34 },
35 - aiAlertAnalysis({ indexId, indexName }: { indexId: string; indexName: string }) {
35 + aiAlertAnalysis({ indexId, indexName, alertId }: { indexId: string; indexName: string; alertId: number }) {
36 return HttpClient.post<FlaskBaseResponse & AiAnalysisResponse>(`/threat_intel/ai/analyze-alert`, {
37 index_name: indexName,
38 - index_id: indexId
38 + index_id: indexId,
39 + alert_id: alertId
40 })
41 },
41 - aiWazuhExclusionRule({ indexId, indexName }: { indexId: string; indexName: string }) {
42 + aiWazuhExclusionRule({ indexId, indexName, alertId }: { indexId: string; indexName: string; alertId: number }) {
43 return HttpClient.post<FlaskBaseResponse & AiWazuhExclusionRuleResponse>(
44 `/threat_intel/ai/wazuh-exclusion-rule`,
45 {
46 index_name: indexName,
46 - index_id: indexId
47 + index_id: indexId,
48 + alert_id: alertId
49 }
50 )
51 },
52 aiVelociraptorArtifactRecommendation({
53 indexId,
54 indexName,
53 - agentId
55 + agentId,
56 + alertId
57 }: {
58 indexId: string
59 indexName: string
60 agentId: string
61 + alertId: number
62 }) {
63 return HttpClient.post<FlaskBaseResponse & AiVelociraptorArtifactRecommendationResponse>(
64 `/threat_intel/ai/velociraptor-artifact-recommendation`,
65 {
66 index_name: indexName,
67 index_id: indexId,
64 - agent_id: agentId
68 + agent_id: agentId,
69 + alert_id: alertId
70 }
71 )
72 },
frontend/src/components/incidentManagement/alerts/AlertAsset.vue
+11 -2
@@ -57,9 +57,18 @@
57 :index-id="asset.index_id"
58 :index-name="asset.index_name"
59 :agent-id="asset.agent_id"
60 + :alert-id="asset.alert_linked"
61 + />
62 + <AIWazuhExclusionRuleButton
63 + :index-id="asset.index_id"
64 + :index-name="asset.index_name"
65 + :alert-id="asset.alert_linked"
66 + />
67 + <AIAnalystButton
68 + :index-id="asset.index_id"
69 + :index-name="asset.index_name"
70 + :alert-id="asset.alert_linked"
71 />
61 - <AIWazuhExclusionRuleButton :index-id="asset.index_id" :index-name="asset.index_name" />
62 - <AIAnalystButton :index-id="asset.index_id" :index-name="asset.index_name" />
72 </div>
73 <n-divider class="!my-0" />
74 <n-tabs type="line" animated :tabs-padding="24">
frontend/src/components/threatIntel/AIAnalystButton.vue
+3 -2
@@ -104,9 +104,10 @@ import LicenseFeatureCheck from "@/components/license/LicenseFeatureCheck.vue"
104 import { NButton, NModal, NTabPane, NTabs, useMessage } from "naive-ui"
105 import { defineAsyncComponent, ref } from "vue"
106
107 -const { indexName, indexId, size } = defineProps<{
107 +const { indexName, indexId, alertId, size } = defineProps<{
108 indexName: string
109 indexId: string
110 + alertId: number
111 size?: Size
112 }>()
113
@@ -131,7 +132,7 @@ function analysis() {
132 loading.value = true
133
134 Api.threatIntel
134 - .aiAlertAnalysis({ indexName, indexId })
135 + .aiAlertAnalysis({ indexName, indexId, alertId })
136 .then(res => {
137 if (res.data.success) {
138 analysisResponse.value = res.data
frontend/src/components/threatIntel/AIVelociraptorArtifactRecommendationButton.vue
+3 -2
@@ -74,10 +74,11 @@ import LicenseFeatureCheck from "@/components/license/LicenseFeatureCheck.vue"
74 import { NButton, NEmpty, NModal, useMessage } from "naive-ui"
75 import { ref } from "vue"
76
77 -const { indexName, indexId, agentId, size } = defineProps<{
77 +const { indexName, indexId, agentId, alertId, size } = defineProps<{
78 indexName: string
79 indexId: string
80 agentId: string
81 + alertId: number
82 size?: Size
83 }>()
84
@@ -99,7 +100,7 @@ function analysis() {
100 loading.value = true
101
102 Api.threatIntel
102 - .aiVelociraptorArtifactRecommendation({ indexName, indexId, agentId })
103 + .aiVelociraptorArtifactRecommendation({ indexName, indexId, agentId, alertId })
104 .then(res => {
105 if (res.data.success) {
106 analysisResponse.value = res.data
frontend/src/components/threatIntel/AIWazuhExclusionRuleButton.vue
+3 -2
@@ -66,9 +66,10 @@ import LicenseFeatureCheck from "@/components/license/LicenseFeatureCheck.vue"
66 import { NButton, NEmpty, NModal, useMessage } from "naive-ui"
67 import { defineAsyncComponent, ref } from "vue"
68
69 -const { indexName, indexId, size } = defineProps<{
69 +const { indexName, indexId, alertId, size } = defineProps<{
70 indexName: string
71 indexId: string
72 + alertId: number
73 size?: Size
74 }>()
75
@@ -93,7 +94,7 @@ function analysis() {
94 loading.value = true
95
96 Api.threatIntel
96 - .aiWazuhExclusionRule({ indexName, indexId })
97 + .aiWazuhExclusionRule({ indexName, indexId, alertId })
98 .then(res => {
99 if (res.data.success) {
100 analysisResponse.value = res.data