Skip to content

Commit 41ac5e8

Browse files
phrockerphrocker
committed
commit
1 parent e736a17 commit 41ac5e8

File tree

16 files changed

+405
-31
lines changed

16 files changed

+405
-31
lines changed

‎.gcp.env‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
SENTRIUS_VERSION=1.0.34
22
SENTRIUS_SSH_VERSION=1.0.3
3-
SENTRIUS_KEYCLOAK_VERSION=1.0.4
3+
SENTRIUS_KEYCLOAK_VERSION=1.0.5
44
SENTRIUS_AGENT_VERSION=1.0.15

‎analyagents/src/main/java/io/sentrius/agent/analysis/agents/sessions/SessionAnalyticsAgent.java‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,35 @@
11
package io.sentrius.agent.analysis.agents.sessions;
22

3+
import java.nio.charset.StandardCharsets;
34
import java.sql.Timestamp;
45
import java.util.ArrayList;
5-
import java.util.Arrays;
6+
import java.util.Base64;
67
import java.util.List;
7-
import java.util.Map;
88
import java.util.Set;
99
import java.util.regex.Matcher;
1010
import java.util.regex.Pattern;
1111
import java.util.stream.Collectors;
12-
import com.fasterxml.jackson.core.JsonProcessingException;
12+
import io.sentrius.sso.core.model.categorization.CommandCategory;
1313
import io.sentrius.sso.core.model.metadata.AnalyticsTracking;
1414
import io.sentrius.sso.core.model.metadata.TerminalBehaviorMetrics;
1515
import io.sentrius.sso.core.model.metadata.TerminalCommand;
1616
import io.sentrius.sso.core.model.metadata.TerminalRiskIndicator;
1717
import io.sentrius.sso.core.model.metadata.TerminalSessionMetadata;
1818
import io.sentrius.sso.core.model.metadata.UserExperienceMetrics;
19-
import io.sentrius.sso.core.model.security.IntegrationSecurityToken;
2019
import io.sentrius.sso.core.model.sessions.TerminalLogs;
2120
import io.sentrius.sso.core.repository.AnalyticsTrackingRepository;
2221
import io.sentrius.sso.core.services.IntegrationSecurityTokenService;
23-
import io.sentrius.sso.core.services.PluggableServices;
2422
import io.sentrius.sso.core.services.SessionService;
2523
import io.sentrius.sso.core.services.metadata.TerminalBehaviorMetricsService;
2624
import io.sentrius.sso.core.services.metadata.TerminalCommandService;
2725
import io.sentrius.sso.core.services.metadata.TerminalRiskIndicatorService;
2826
import io.sentrius.sso.core.services.metadata.TerminalSessionMetadataService;
2927
import io.sentrius.sso.core.services.metadata.UserExperienceMetricsService;
30-
import io.sentrius.sso.core.utils.JsonUtil;
31-
import io.sentrius.sso.integrations.external.ExternalIntegrationDTO;
32-
import io.sentrius.sso.security.ApiKey;
28+
import io.sentrius.sso.core.services.openai.categorization.CommandCategorizer;
3329
import jakarta.transaction.Transactional;
3430
import lombok.RequiredArgsConstructor;
3531
import lombok.extern.slf4j.Slf4j;
3632
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
37-
import org.springframework.context.ApplicationContext;
3833
import org.springframework.scheduling.annotation.Scheduled;
3934
import org.springframework.stereotype.Component;
4035

@@ -51,6 +46,7 @@ public class SessionAnalyticsAgent {
5146
private final UserExperienceMetricsService experienceMetricsService;
5247
private final AnalyticsTrackingRepository trackingRepository;
5348
private final SessionService sessionService;
49+
private final CommandCategorizer commandCategorizer;
5450
final IntegrationSecurityTokenService integrationSecurityTokenService;
5551

5652

@@ -61,6 +57,7 @@ public void processSessions() {
6157

6258
// Fetch already processed session IDs in bulk
6359
Set<Long> processedSessionIds = trackingRepository.findAllSessionIds();
60+
log.info("Found {} processed sessions", processedSessionIds.size());
6461
List<TerminalSessionMetadata> unprocessedSessions = sessionMetadataService.getSessionsByState("CLOSED").stream()
6562
.filter(session -> !processedSessionIds.contains(session.getId()))
6663
.collect(Collectors.toList());
@@ -69,13 +66,13 @@ public void processSessions() {
6966
try {
7067
processSession(session);
7168
// ACTIVE -> INACTIVE -> CLOSED -> PROCESSED
72-
// saveToTracking(session.getId(), "PROCESSED");
69+
saveToTracking(session.getId(), "PROCESSED");
7370
} catch (Exception e) {
7471
log.error("Error processing session {}: {}", session.getId(), e.getMessage(), e);
7572
saveToTracking(session.getId(), "ERROR");
7673
}
77-
// session.setSessionStatus("PROCESSED");
78-
// sessionMetadataService.saveSession(session);
74+
session.setSessionStatus("PROCESSED");
75+
sessionMetadataService.saveSession(session);
7976
}
8077

8178
log.info("Finished processing sessions");
@@ -185,7 +182,6 @@ public static String extractCommand(TerminalLogs previousLog, String logLine) {
185182
return matcher.group(1).trim();
186183
} else {
187184
if (null != previousLog) {
188-
log.info("Previous log: {}", previousLog.getOutput());
189185
// it could be that we are at the beginning of the log set.
190186
String lastLogLine = getLastLogLine(previousLog);
191187
if (!lastLogLine.isEmpty()) {
@@ -215,26 +211,20 @@ private static String getLastLogLine(TerminalLogs previousLog) {
215211
}
216212

217213
private TerminalCommand createTerminalCommand(String command, TerminalLogs terminalLog, TerminalSessionMetadata sessionMetadata) {
214+
String encodedString = Base64.getEncoder().encodeToString(command.trim().getBytes(StandardCharsets.UTF_8));
215+
218216
TerminalCommand terminalCommand = new TerminalCommand();
219-
terminalCommand.setCommand(command.trim());
217+
terminalCommand.setCommand(encodedString);
220218
terminalCommand.setSession(sessionMetadata);
221219
terminalCommand.setExecutionTime(new Timestamp(System.currentTimeMillis()));
222220
terminalCommand.setExecutionStatus("SUCCESS");
223221
terminalCommand.setOutput(""); // Assume no output initially
224-
terminalCommand.setCommandCategory(categorizeCommand(command));
222+
terminalCommand.setCommandCategory(categorizeCommand(command).getCategoryName());
225223

226224
return terminalCommand;
227225
}
228226

229-
private String categorizeCommand(String command) {
230-
// probably need to define externally
231-
if (command.startsWith("sudo")) {
232-
return "PRIVILEGED";
233-
} else if (command.contains("rm")) {
234-
return "DESTRUCTIVE";
235-
} else if (command.contains("ls") || command.contains("cat")) {
236-
return "INFORMATIONAL";
237-
}
238-
return "GENERAL";
227+
private CommandCategory categorizeCommand(String command) {
228+
return commandCategorizer.categorizeCommand(command);
239229
}
240230
}

‎analyagents/src/main/resources/application.properties‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,4 +61,5 @@ spring.security.oauth2.client.provider.keycloak.issuer-uri=http://192.168.1.162:
6161
# for testing analytics agents
6262
agents.session-analytics.enabled=true
6363
management.endpoints.web.exposure.include=health
64-
management.endpoint.health.show-details=always
64+
management.endpoint.health.show-details=always
65+

‎api/src/main/resources/db/migration/V13__command_categorizer.sql‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
CREATE TABLE command_categories (
2+
id SERIAL PRIMARY KEY,
3+
category_name VARCHAR(50) NOT NULL,
4+
pattern TEXT NOT NULL, -- Store regex patterns
5+
priority INT NOT NULL DEFAULT 0 -- Optional: for matching precedence
6+
);
7+
8+
9+
CREATE INDEX idx_pattern ON command_categories (pattern);

‎core/pom.xml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,11 @@
139139
<artifactId>quartz</artifactId>
140140
<version>${quartz-version}</version>
141141
</dependency>
142-
142+
<dependency>
143+
<groupId>com.github.ben-manes.caffeine</groupId>
144+
<artifactId>caffeine</artifactId>
145+
<version>${caffeine-version}</version>
146+
</dependency>
143147
<dependency>
144148
<groupId>com.google.protobuf</groupId>
145149
<artifactId>protobuf-java</artifactId>

‎core/src/main/java/io/sentrius/sso/core/model/categorization/CommandCategory.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.sentrius.sso.core.model.categorization;
2+
3+
import jakarta.persistence.Entity;
4+
import jakarta.persistence.GeneratedValue;
5+
import jakarta.persistence.GenerationType;
6+
import jakarta.persistence.Id;
7+
import jakarta.persistence.Table;
8+
import lombok.AllArgsConstructor;
9+
import lombok.Builder;
10+
import lombok.Getter;
11+
import lombok.NoArgsConstructor;
12+
import lombok.ToString;
13+
14+
@Getter
15+
@Builder
16+
@NoArgsConstructor
17+
@ToString
18+
@AllArgsConstructor
19+
@Entity
20+
@Table(name = "command_categories")
21+
public class CommandCategory {
22+
@Id
23+
@GeneratedValue(strategy = GenerationType.IDENTITY)
24+
private Long id;
25+
private String categoryName;
26+
private String pattern;
27+
private int priority;
28+
29+
}

‎core/src/main/java/io/sentrius/sso/core/repository/CommandCategoryRepository.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package io.sentrius.sso.core.repository;
2+
3+
import java.util.List;
4+
import io.sentrius.sso.core.model.categorization.CommandCategory;
5+
import org.springframework.data.jpa.repository.JpaRepository;
6+
import org.springframework.data.jpa.repository.Query;
7+
import org.springframework.stereotype.Repository;
8+
9+
@Repository
10+
public interface CommandCategoryRepository extends JpaRepository<CommandCategory, Long> {
11+
@Query("SELECT c FROM CommandCategory c ORDER BY c.priority ASC")
12+
List<CommandCategory> findAllOrderedByPriority();
13+
14+
List<CommandCategory> findByPattern(String pattern);
15+
}

‎core/src/main/java/io/sentrius/sso/core/services/openai/categorization/CommandCategorizer.java‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package io.sentrius.sso.core.services.openai.categorization;
2+
3+
import java.util.List;
4+
import java.util.concurrent.TimeUnit;
5+
import java.util.regex.Pattern;
6+
import com.fasterxml.jackson.core.JsonProcessingException;
7+
import com.github.benmanes.caffeine.cache.Cache;
8+
import com.github.benmanes.caffeine.cache.Caffeine;
9+
import io.sentrius.sso.core.model.categorization.CommandCategory;
10+
import io.sentrius.sso.core.repository.CommandCategoryRepository;
11+
import io.sentrius.sso.core.services.IntegrationSecurityTokenService;
12+
import io.sentrius.sso.core.utils.JsonUtil;
13+
import io.sentrius.sso.genai.GenerativeAPI;
14+
import io.sentrius.sso.genai.GeneratorConfiguration;
15+
import io.sentrius.sso.genai.LLMCommandCategorizer;
16+
import io.sentrius.sso.integrations.external.ExternalIntegrationDTO;
17+
import io.sentrius.sso.security.ApiKey;
18+
import jakarta.annotation.PostConstruct;
19+
import jakarta.transaction.Transactional;
20+
import jdk.jfr.TransitionTo;
21+
import lombok.AllArgsConstructor;
22+
import lombok.RequiredArgsConstructor;
23+
import lombok.extern.slf4j.Slf4j;
24+
import org.springframework.stereotype.Service;
25+
26+
@Slf4j
27+
@Service
28+
@RequiredArgsConstructor
29+
public class CommandCategorizer {
30+
31+
private final IntegrationSecurityTokenService integrationSecurityTokenService;
32+
33+
private final CommandCategoryRepository commandCategoryRepository;
34+
35+
36+
private final CommandTrie commandTrie = new CommandTrie();
37+
38+
private final Cache<String, CommandCategory> commandCache = Caffeine.newBuilder()
39+
.maximumSize(10000)
40+
.expireAfterWrite(24, TimeUnit.HOURS)
41+
.build();
42+
43+
@PostConstruct
44+
public void initializeTrie() {
45+
List<CommandCategory> categories = commandCategoryRepository.findAll();
46+
for (CommandCategory category : categories) {
47+
log.info("Adding command category {} to trie", category);
48+
commandTrie.insert(category.getPattern(), category);
49+
}
50+
}
51+
52+
@Transactional
53+
public CommandCategory categorizeCommand(String command) {
54+
return commandCache.get(command, this::categorizeWithRulesOrML);
55+
}
56+
57+
58+
protected List<CommandCategory> getDBCommandCategory(String command){
59+
return commandCategoryRepository.findByPattern(command);
60+
}
61+
62+
@Transactional
63+
protected void addCommandCategory(String command, CommandCategory category) {
64+
commandTrie.insert(command, category);
65+
commandCategoryRepository.save(category);
66+
}
67+
68+
69+
@Transactional
70+
protected CommandCategory categorizeWithRulesOrML(String command) {
71+
CommandCategory category = commandTrie.searchByPrefix(command);
72+
if (category != null) {
73+
log.info("Found command category {} for {} ", category, command);
74+
return category;
75+
}
76+
77+
var openaiService = integrationSecurityTokenService.findByConnectionType("openai").stream().findFirst().orElse(null);
78+
79+
if (null != openaiService){
80+
log.info("OpenAI service is available");
81+
ExternalIntegrationDTO externalIntegrationDTO = null;
82+
try {
83+
externalIntegrationDTO = JsonUtil.MAPPER.readValue(openaiService.getConnectionInfo(),
84+
ExternalIntegrationDTO.class);
85+
} catch (JsonProcessingException e) {
86+
throw new RuntimeException(e);
87+
}
88+
ApiKey key =
89+
ApiKey.builder().apiKey(externalIntegrationDTO.getApiToken()).principal(externalIntegrationDTO.getUsername()).build();
90+
91+
var commandCategorizer = new LLMCommandCategorizer(key, new GenerativeAPI(key), GeneratorConfiguration.builder().build());
92+
93+
try {
94+
category = commandCategorizer.generate(command);
95+
96+
addCommandCategory(category.getPattern(), category);
97+
98+
log.info("Categorized command: {}", category);
99+
return category;
100+
} catch (Exception e) {
101+
e.printStackTrace();
102+
log.error("Error categorizing command", e);
103+
}
104+
105+
} else {
106+
log.info("OpenAI service is not enabled");
107+
}
108+
109+
log.info("Finished processing terminal commands");
110+
111+
return CommandCategory.builder().build();
112+
}
113+
}

‎core/src/main/java/io/sentrius/sso/core/services/openai/categorization/CommandTrie.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package io.sentrius.sso.core.services.openai.categorization;
2+
3+
import io.sentrius.sso.core.model.categorization.CommandCategory;
4+
import lombok.extern.slf4j.Slf4j;
5+
6+
@Slf4j
7+
public class CommandTrie {
8+
private final TrieNode root = new TrieNode();
9+
10+
private String normalizePath(String path) {
11+
return path.replaceAll("/+$", ""); // Remove trailing slashes
12+
}
13+
14+
public void insert(String command, CommandCategory category) {
15+
String[] parts = command.split(" ");
16+
TrieNode current = root;
17+
for (String part : parts) {
18+
part = normalizePath(part); // Normalize each part
19+
current = current.children.computeIfAbsent(part, k -> new TrieNode());
20+
}
21+
current.isEndOfCommand = true;
22+
if (category.getPattern().endsWith("*")) {
23+
current.isWildcard = true;
24+
}
25+
else {
26+
current.isWildcard = false;
27+
}
28+
current.commandCategory = category;
29+
}
30+
31+
public CommandCategory search(String command) {
32+
String[] parts = command.split(" ");
33+
TrieNode current = root;
34+
for (String part : parts) {
35+
current = current.children.get(part);
36+
if (current == null) {
37+
return null; // Command not found
38+
}
39+
}
40+
return current.isEndOfCommand ? current.commandCategory : null;
41+
}
42+
43+
public CommandCategory searchByPrefix(String command) {
44+
String[] parts = command.split(" ");
45+
TrieNode current = root;
46+
CommandCategory lastCategory = null;
47+
48+
for (String part : parts) {
49+
part = normalizePath(part); // Normalize each part
50+
log.info("Searching for part: {}", part);
51+
52+
// Check if the part exists in the children
53+
if (current.children.containsKey(part)) {
54+
current = current.children.get(part);
55+
if (current.isEndOfCommand) {
56+
lastCategory = current.commandCategory;
57+
}
58+
} else {
59+
// If no exact match, check for wildcard match (e.g., "cat /etc/")
60+
if (current.isEndOfCommand) {
61+
log.info("Partial match found at: {}", part);
62+
lastCategory = current.commandCategory;
63+
}
64+
break; // No further match possible
65+
}
66+
}
67+
68+
return lastCategory;
69+
}
70+
}

0 commit comments

Comments
 (0)