Skip to content

Commit cbf562e

Browse files
ayushaga14ayushaga14
authored
Merge pull request #2609 from akto-api-security/fix/threat-be-apis
Threat API Optimizations
2 parents a6e56e4 + c386db2 commit cbf562e

File tree

11 files changed

+131
-81
lines changed

11 files changed

+131
-81
lines changed

‎apps/dashboard/src/main/java/com/akto/action/threat_detection/ThreatActorAction.java

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,13 +79,27 @@ public ThreatActorAction() {
7979
}
8080

8181
public String getActorsCountPerCounty() {
82-
HttpGet get =
83-
new HttpGet(
84-
String.format("%s/api/dashboard/get_actors_count_per_country", this.getBackendUrl()));
85-
get.addHeader("Authorization", "Bearer " + this.getApiToken());
86-
get.addHeader("Content-Type", "application/json");
82+
HttpPost post = new HttpPost(String.format("%s/api/dashboard/get_actors_count_per_country", this.getBackendUrl()));
83+
post.addHeader("Authorization", "Bearer " + this.getApiToken());
84+
post.addHeader("Content-Type", "application/json");
8785

88-
try (CloseableHttpResponse resp = this.httpClient.execute(get)) {
86+
if(startTs == 0 || endTs == 0) {
87+
startTs = Context.now() - 1 * 24 * 60 * 60; // default to last 1 day
88+
endTs = Context.now();
89+
}
90+
91+
Map<String, Object> body = new HashMap<String, Object>() {
92+
{
93+
put("start_ts", startTs);
94+
put("end_ts", endTs);
95+
}
96+
};
97+
String msg = objectMapper.valueToTree(body).toString();
98+
99+
StringEntity requestEntity = new StringEntity(msg, ContentType.APPLICATION_JSON);
100+
post.setEntity(requestEntity);
101+
102+
try (CloseableHttpResponse resp = this.httpClient.execute(post)) {
89103
String responseBody = EntityUtils.toString(resp.getEntity());
90104

91105
ProtoMessageUtils.<ThreatActorByCountryResponse>toProtoMessage(

‎apps/dashboard/src/main/resources/struts.xml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8240,6 +8240,7 @@
82408240
<action name="api/getActorsCountPerCounty"
82418241
class="com.akto.action.threat_detection.ThreatActorAction"
82428242
method="getActorsCountPerCounty">
8243+
<interceptor-ref name="json" />
82438244
<interceptor-ref name="defaultStack" />
82448245
<interceptor-ref name="usageInterceptor">
82458246
<param name="featureLabel">THREAT_DETECTION</param>
@@ -8259,7 +8260,9 @@
82598260
<param
82608261
name="includeProperties">^actionErrors.*</param>
82618262
</result>
8262-
<result name="SUCCESS" type="json"> </result>
8263+
<result name="SUCCESS" type="json">
8264+
<param name="includeProperties">^actorsCountPerCountry.*</param>
8265+
</result>
82638266
<result name="ERROR" type="json">
82648267
<param name="statusCode">422</param>
82658268
<param name="ignoreHierarchy">false</param>

‎apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/threat_detection/ThreatActorPage.jsx

Lines changed: 5 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { ThreatSummary } from "./components/ThreatSummary";
1818
import ThreatActivityTimeline from "./components/ThreatActivityTimeline";
1919
import React from "react";
2020

21-
const ChartComponent = ({ mapData, loading, onSubCategoryClick, currDateRange }) => {
21+
const ChartComponent = ({ onSubCategoryClick, currDateRange }) => {
2222
return (
2323
<VerticalStack gap={4} columns={2}>
2424
<HorizontalGrid gap={4} columns={2}>
@@ -28,12 +28,12 @@ const ChartComponent = ({ mapData, loading, onSubCategoryClick, currDateRange })
2828
endTimestamp={parseInt(currDateRange.period.until.getTime()/1000)}
2929
/>
3030
<ThreatWorldMap
31-
data={mapData}
31+
startTimestamp={parseInt(currDateRange.period.since.getTime()/1000)}
32+
endTimestamp={parseInt(currDateRange.period.until.getTime()/1000)}
3233
style={{
3334
width: "100%",
3435
marginRight: "auto",
3536
}}
36-
loading={loading}
3737
key={"threat-actor-world-map"}
3838
/>
3939
</HorizontalGrid>
@@ -44,9 +44,6 @@ const ChartComponent = ({ mapData, loading, onSubCategoryClick, currDateRange })
4444
const MemoizedChartComponent = React.memo(ChartComponent);
4545

4646
function ThreatActorPage() {
47-
const [mapData, setMapData] = useState([]);
48-
const [loading, setLoading] = useState(false);
49-
const [subCategoryCount, setSubCategoryCount] = useState([]);
5047
const [actorDetails, setActorDetails] = useState(null);
5148
const [showActorDetails, setShowActorDetails] = useState(false);
5249

@@ -57,32 +54,6 @@ function ThreatActorPage() {
5754
);
5855

5956
useEffect(() => {
60-
const fetchActorsPerCountry = async () => {
61-
setLoading(true);
62-
const res = await api.getActorsCountPerCounty();
63-
if (res?.actorsCountPerCountry) {
64-
setMapData(
65-
res.actorsCountPerCountry.map((x) => {
66-
return {
67-
code: x.country,
68-
z: 100,
69-
count: x.count,
70-
};
71-
})
72-
);
73-
}
74-
setLoading(false);
75-
};
76-
const fetchThreatCategoryCount = async () => {
77-
setLoading(true);
78-
const res = await api.fetchThreatCategoryCount();
79-
const finalObj = threatDetectionFunc.getGraphsData(res);
80-
// setCategoryCount(finalObj.categoryCountRes);
81-
setSubCategoryCount(finalObj.subCategoryCount);
82-
setLoading(false);
83-
};
84-
fetchActorsPerCountry();
85-
fetchThreatCategoryCount();
8657
}, []);
8758

8859
const onSubCategoryClick = (subCategory) => {
@@ -96,16 +67,14 @@ function ThreatActorPage() {
9667

9768
const components = [
9869
<ThreatSummary startTimestamp={parseInt(currDateRange.period.since.getTime()/1000)} endTimestamp={parseInt(currDateRange.period.until.getTime()/1000)} />,
99-
<MemoizedChartComponent
100-
mapData={mapData}
101-
loading={loading}
70+
<MemoizedChartComponent
71+
key={"threat-actor-chart-component"}
10272
onSubCategoryClick={onSubCategoryClick}
10373
currDateRange={currDateRange}
10474
/>,
10575
<ThreatActorTable
10676
key={"threat-actor-data-table"}
10777
currDateRange={currDateRange}
108-
loading={loading}
10978
handleRowClick={onRowClick}
11079
/>,
11180
...(showActorDetails ? [<ActorDetails actorDetails={actorDetails} setShowActorDetails={setShowActorDetails} />] : [])

‎apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/threat_detection/api.js

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,11 +64,11 @@ const threatDetectionRequests = {
6464
}
6565
})
6666
},
67-
getActorsCountPerCounty() {
67+
getActorsCountPerCounty(startTs, endTs) {
6868
return request({
6969
url: '/api/getActorsCountPerCounty',
70-
method: 'get',
71-
data: {}
70+
method: 'post',
71+
data: {startTs, endTs}
7272
})
7373
},
7474
fetchThreatConfiguration() {

‎apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/threat_detection/components/ThreatActorsTable.jsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ function ThreatActorTable({ data, currDateRange, handleRowClick }) {
161161
...x,
162162
actor: x.id ? (
163163
<Text variant="bodyMd" fontWeight="medium">
164-
{x.id}
164+
{x.id?.length > 50 ? `${x.id.slice(0, 50)}...` : x.id}
165165
</Text>
166166
) : "-",
167167
latestIp: x.latestApiIp || "-",

‎apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/threat_detection/components/ThreatWorldMap.jsx

Lines changed: 46 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,39 @@
1-
import { useEffect } from "react";
1+
import { useEffect, useState } from "react";
22
import Highcharts from "highcharts/highmaps";
33
import Exporting from "highcharts/modules/exporting";
44
import ExportData from "highcharts/modules/export-data";
55
import FullScreen from "highcharts/modules/full-screen";
66
import InfoCard from "../../dashboard/new_components/InfoCard";
7+
import { Spinner } from "@shopify/polaris";
8+
import api from "../api";
9+
710
// Initialize modules
811
Exporting(Highcharts);
912
ExportData(Highcharts);
1013
FullScreen(Highcharts);
1114

12-
function ThreatWorldMap({ data, style, loading }) {
13-
useEffect(() => {
15+
function ThreatWorldMap({ startTimestamp, endTimestamp, style}) {
16+
17+
const [loading, setLoading] = useState(false);
18+
const [data, setData] = useState(false);
19+
20+
const fetchActorsPerCountry = async () => {
21+
// setLoading(true);
22+
const res = await api.getActorsCountPerCounty(startTimestamp, endTimestamp);
23+
if (res?.actorsCountPerCountry) {
24+
setData(
25+
res.actorsCountPerCountry.map((x) => {
26+
return {
27+
code: x.country,
28+
z: 100,
29+
count: x.count,
30+
};
31+
})
32+
);
33+
}
34+
setLoading(false);
35+
};
36+
1437
const fetchMapData = async () => {
1538
const topology = await fetch(
1639
"https://code.highcharts.com/mapdata/custom/world.topo.json"
@@ -104,14 +127,28 @@ function ThreatWorldMap({ data, style, loading }) {
104127
});
105128
};
106129

107-
fetchMapData();
130+
useEffect(() => {
131+
if (data) {
132+
fetchMapData();
133+
}
108134
}, [data]);
109135

110-
return <InfoCard
111-
title={"Threat Actor Map"}
112-
titleToolTip={"Threat Actor Map"}
113-
component={<div id="threat-world-map-container" style={style}></div>}
114-
/>;
136+
useEffect(() => {
137+
fetchActorsPerCountry();
138+
fetchMapData();
139+
}, [startTimestamp, endTimestamp]);
140+
141+
if (loading) {
142+
return <Spinner />;
143+
}
144+
145+
return (
146+
<InfoCard
147+
title={"Threat Actor Map"}
148+
titleToolTip={"Threat Actor Map"}
149+
component={<div id="threat-world-map-container" style={style}></div>}
150+
/>
151+
);
115152
}
116153

117154
export default ThreatWorldMap;

‎apps/threat-detection-backend/src/main/java/com/akto/threat/backend/router/DashboardRouter.java

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,12 +120,25 @@ public Router setup(Vertx vertx) {
120120
});
121121

122122
router
123-
.get("/get_actors_count_per_country")
123+
.post("/get_actors_count_per_country")
124124
.blockingHandler(ctx -> {
125+
RequestBody reqBody = ctx.body();
126+
ThreatActorByCountryRequest req = ProtoMessageUtils.<
127+
ThreatActorByCountryRequest
128+
>toProtoMessage(
129+
ThreatActorByCountryRequest.class,
130+
reqBody.asString()
131+
).orElse(null);
132+
133+
if (req == null) {
134+
ctx.response().setStatusCode(400).end("Invalid request");
135+
return;
136+
}
137+
125138
ProtoMessageUtils.toString(
126139
threatActorService.getThreatActorByCountry(
127140
ctx.get("accountId"),
128-
ThreatActorByCountryRequest.newBuilder().build()
141+
req
129142
)
130143
).ifPresent(s -> ctx.response().setStatusCode(200).end(s));
131144
});

‎apps/threat-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatActorService.java

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -419,30 +419,36 @@ public ThreatActorByCountryResponse getThreatActorByCountry(
419419
.getCollection(MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, Document.class);
420420

421421
List<Document> pipeline = new ArrayList<>();
422-
pipeline.add(
423-
new Document("$sort", new Document("country", 1).append("detectedAt", -1))); // sort
424-
pipeline.add(new Document("$limit", 10000));
425-
pipeline.add(
426-
new Document(
427-
"$group",
428-
new Document("_id", "$country")
429-
.append("distinctActors", new Document("$addToSet", "$actor"))));
430422

431-
pipeline.add(
432-
new Document(
433-
"$addFields", new Document("actorsCount", new Document("$size", "$distinctActors"))));
423+
// 1. Match on time range
424+
if (request.getStartTs() != 0 || request.getEndTs() != 0) {
425+
pipeline.add(new Document("$match",
426+
new Document("detectedAt",
427+
new Document("$gte", request.getStartTs())
428+
.append("$lte", request.getEndTs()))));
429+
}
430+
431+
// 2. Project only necessary fields
432+
pipeline.add(new Document("$project", new Document("country", 1).append("actor", 1)));
433+
434+
// 3. Group by country and collect distinct actors
435+
pipeline.add(new Document("$group",
436+
new Document("_id", "$country")
437+
.append("distinctActorsCount", new Document("$addToSet", "$actor"))));
434438

435-
pipeline.add(new Document("$sort", new Document("actorsCount", -1))); // sort
439+
// 4. Project the size of the distinct actors set
440+
pipeline.add(new Document("$project",
441+
new Document("distinctActorsCount", new Document("$size", "$distinctActorsCount"))));
436442

437443
List<ThreatActorByCountryResponse.CountryCount> actorsByCountryCount = new ArrayList<>();
438444

439-
try (MongoCursor<Document> cursor = coll.aggregate(pipeline).cursor()) {
445+
try (MongoCursor<Document> cursor = coll.aggregate(pipeline).batchSize(1000).cursor()) {
440446
while (cursor.hasNext()) {
441447
Document doc = cursor.next();
442448
actorsByCountryCount.add(
443449
ThreatActorByCountryResponse.CountryCount.newBuilder()
444450
.setCode(doc.getString("_id"))
445-
.setCount(doc.getInteger("actorsCount", 0))
451+
.setCount(doc.getInteger("distinctActorsCount", 0))
446452
.build());
447453
}
448454
}

‎apps/threat-detection/src/main/java/com/akto/threat/detection/tasks/ThreatConfigurationEvaluator.java

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -95,12 +95,11 @@ public boolean isHostNameMatching(HttpResponseParams responseParam, String patte
9595

9696
public String getActorId(HttpResponseParams responseParam) {
9797
getThreatConfiguration();
98-
String actor;
99-
String sourceIp = SourceIPActorGenerator.instance.generate(responseParam).orElse("");
100-
responseParam.setSourceIP(sourceIp);
98+
String actor = SourceIPActorGenerator.instance.generate(responseParam).orElse("");
99+
responseParam.setSourceIP(actor);
100+
logger.debugAndAddToDbCount("Actor ID generated: " + actor + " for response: " + responseParam.getOriginalMsg().get());
101101

102102
if (threatConfiguration == null) {
103-
actor = sourceIp;
104103
return actor;
105104
}
106105

@@ -116,18 +115,16 @@ public String getActorId(HttpResponseParams responseParam) {
116115
.get(actorId.getKey().toLowerCase());
117116
if (header != null && !header.isEmpty()) {
118117
actor = header.get(0);
118+
return actor;
119119
} else {
120120
logger.warn("Defaulting to source IP as actor id, header not found: "
121121
+ actorId.getKey());
122-
actor = sourceIp;
122+
return actor;
123123
}
124-
break;
125124
default:
126-
actor = sourceIp;
127125
break;
128126
}
129-
return actor;
130127
}
131-
return sourceIp;
128+
return actor;
132129
}
133130
}

‎libs/utils/src/main/java/com/akto/log/LoggerMaker.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class LoggerMaker {
4949
private final Class<?> aClass;
5050

5151
private static String slackWebhookUrl;
52+
private static int counter = 0;
5253

5354
public static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
5455
private static final DataActor dataActor = DataActorFactory.fetchInstance();
@@ -348,6 +349,14 @@ public void debugAndAddToDb(String message) {
348349
debugAndAddToDb(message, this.db);
349350
}
350351

352+
public void debugAndAddToDbCount(String message) {
353+
if(counter > 500){
354+
return;
355+
}
356+
counter++;
357+
debugAndAddToDb(message, this.db);
358+
}
359+
351360
public void debugAndAddToDb(String message, LogDb db) {
352361
String accountId = Context.accountId.get() != null ? Context.accountId.get().toString() : "NA";
353362
String debugMessage = "acc: " + accountId + ", " + message;

‎protobuf/threat_detection/service/dashboard_service/v1/service.proto

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,8 @@ message ListThreatApiResponse {
141141
}
142142

143143
message ThreatActorByCountryRequest {
144+
uint32 start_ts = 1;
145+
uint32 end_ts = 2;
144146
}
145147

146148
message ThreatActorByCountryResponse {

0 commit comments

Comments
 (0)