Skip to content

Commit eb5c5b7

Browse files
phrockerphrocker
committed
commit
1 parent 2465501 commit eb5c5b7

File tree

78 files changed

+7894
-66
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

78 files changed

+7894
-66
lines changed

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ dataplane/target/**
4444
dataplane/target/
4545
llm-proxy/target/**
4646
llm-proxy/target/
47+
llm-dataplane/target/**
48+
llm-dataplane/target/
4749
node/*
4850
node_modules/*
4951
api/node_modules/*

‎ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ public class AgentVerbs {
4848
final VerbRegistry verbRegistry;
4949
final AgentClientService agentClientService;
5050

51+
5152
@Value("${agent.ai.config}")
5253
private String agentConfigFile;
5354

‎ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/TerminalVerbs.java‎

Lines changed: 62 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,9 +99,10 @@ public List<ObjectNode> fetchTerminalOutput(TokenDTO token, List<HostSystemDTO>
9999
}
100100
}
101101

102-
@Verb(name = "kill_session", description = "Kills a terminal session.", requiresTokenManagement = true,
102+
@Verb(name = "kill_session_with_assessment", description = "Kills a terminal session using a terminal assessment.",
103+
requiresTokenManagement = true,
103104
outputInterpreter = TerminalOutputInterpreter.class, inputInterpreter = ObjectListInterpreter.class)
104-
public List<ObjectNode> killTerminalSession(AgentExecution execution, List<Object> dtos) {
105+
public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution execution, List<Object> dtos) {
105106
try {
106107
List<ObjectNode> responses = new ArrayList<>();
107108
log.info("Terminal list response: {}", dtos);
@@ -154,4 +155,63 @@ public List<ObjectNode> killTerminalSession(AgentExecution execution, List<Objec
154155
throw new RuntimeException("Failed to retrieve terminal list", e);
155156
}
156157
}
158+
159+
@Verb(name = "kill_session_id" , description = "Kills a terminal session by session id.",
160+
requiresTokenManagement = true,
161+
outputInterpreter = TerminalOutputInterpreter.class, inputInterpreter = ObjectListInterpreter.class)
162+
public List<ObjectNode> killTerminalBySessionId(AgentExecution execution, List<Object> dtos) {
163+
try {
164+
List<ObjectNode> responses = new ArrayList<>();
165+
log.info("Terminal list response: {}", dtos);
166+
for (Object dto : dtos) {
167+
// submit the kill
168+
ObjectNode node = JsonUtil.MAPPER.readValue(dto.toString(), ObjectNode.class);
169+
ArrayNode assessments = node.get("assessments").deepCopy();
170+
for(int i = 0; i < assessments.size(); i++) {
171+
var assessment = assessments.get(i);
172+
var risk = assessment.get("risk");
173+
var description = assessment.get("description");
174+
if (null != risk && null != description) {
175+
switch(risk.asText()) {
176+
case "low":
177+
// skip and do nothing
178+
continue;
179+
case "medium":
180+
case "high":
181+
// kill the session
182+
log.info("Killing terminal session: {}", assessment.get("sessionId").asText());
183+
break;
184+
default:
185+
throw new RuntimeException("Unknown risk level: " + risk.asText());
186+
}
187+
try {
188+
var response = zeroTrustClientService.callGetOnApi(
189+
execution, "/ssh/terminal/kill",
190+
Maps.immutableEntry("sessionId", List.of(assessment.get("sessionId").asText()))
191+
);
192+
if (response != null) {
193+
// Successfully retrieved logs
194+
log.info("Terminal output response: {}", response);
195+
var obj = JsonUtil.MAPPER.createObjectNode();
196+
obj.put("id", assessment.get("hostConnection").asText());
197+
obj.put("terminalOutput", response);
198+
responses.add(obj);
199+
}
200+
}catch (ZtatException e) {
201+
var ztatRequest = e.getZtatRequestId();
202+
agentVerbs.justifyAgent(execution, ztatRequest, description.asText());
203+
}
204+
}
205+
}
206+
207+
log.info("Terminal output response: {}", dto);
208+
209+
}
210+
return responses;
211+
} catch (Exception | ZtatException e) {
212+
throw new RuntimeException("Failed to retrieve terminal list", e);
213+
}
214+
}
215+
216+
157217
}

‎ai-agent/src/main/java/io/sentrius/agent/analysis/api/AgentController.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ public class AgentController {
1818

1919
KeycloakService keycloakService;
2020

21-
@GetMapping("/status")
21+
@GetMapping("/ping")
2222
public ResponseEntity<AgentStatus> getStatus() {
2323
String hostName = "unknown";
2424
try {
@@ -43,6 +43,7 @@ public ResponseEntity<AgentStatus> getStatus() {
4343
.freeMemory(runtime.freeMemory())
4444
.build();
4545

46+
log.info("Ping status: {}", status);
4647
return ResponseEntity.ok(status);
4748
}
4849

‎ai-agent/src/main/java/io/sentrius/agent/analysis/api/UserCommunicator.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package io.sentrius.agent.analysis.api;
2+
3+
import java.lang.management.ManagementFactory;
4+
import java.net.InetAddress;
5+
import java.net.UnknownHostException;
6+
import io.sentrius.sso.core.model.AgentStatus;
7+
import io.sentrius.sso.core.services.security.KeycloakService;
8+
import lombok.extern.slf4j.Slf4j;
9+
import org.springframework.http.ResponseEntity;
10+
import org.springframework.web.bind.annotation.GetMapping;
11+
import org.springframework.web.bind.annotation.RequestMapping;
12+
import org.springframework.web.bind.annotation.RestController;
13+
14+
@Slf4j
15+
@RestController
16+
@RequestMapping("/api/v1/agent/communicator")
17+
public class UserCommunicator {
18+
19+
KeycloakService keycloakService;
20+
21+
@GetMapping("/ping")
22+
public ResponseEntity<AgentStatus> getStatus() {
23+
String hostName = "unknown";
24+
try {
25+
hostName = InetAddress.getLocalHost().getHostName();
26+
} catch (UnknownHostException e) {
27+
log.warn("Unable to resolve hostname", e);
28+
}
29+
30+
long uptimeMillis = ManagementFactory.getRuntimeMXBean().getUptime();
31+
Runtime runtime = Runtime.getRuntime();
32+
33+
AgentStatus status = AgentStatus.builder()
34+
.status("UP")
35+
.version("1.0.0")
36+
.health("OK")
37+
.osName(System.getProperty("os.name"))
38+
.osArch(System.getProperty("os.arch"))
39+
.osVersion(System.getProperty("os.version"))
40+
.hostName(hostName)
41+
.uptimeMillis(uptimeMillis)
42+
.totalMemory(runtime.totalMemory())
43+
.freeMemory(runtime.freeMemory())
44+
.build();
45+
46+
log.info("Ping status: {}", status);
47+
return ResponseEntity.ok(status);
48+
}
49+
50+
}

‎ai-agent/src/main/java/io/sentrius/agent/analysis/model/UserCommunication.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package io.sentrius.agent.analysis.model;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import io.sentrius.sso.genai.Message;
6+
import lombok.AllArgsConstructor;
7+
import lombok.Builder;
8+
import lombok.Data;
9+
import lombok.Getter;
10+
import lombok.NoArgsConstructor;
11+
12+
@Data
13+
@Builder
14+
@Getter
15+
@NoArgsConstructor
16+
@AllArgsConstructor
17+
public class UserCommunication {
18+
@Builder.Default
19+
List<Message> conversations = new ArrayList<>();
20+
21+
String currentMessage;
22+
23+
}

‎ai-agent/src/main/java/io/sentrius/agent/config/AgentKeycloakUserSyncFilter.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package io.sentrius.agent.config;
2+
3+
import java.io.IOException;
4+
import io.sentrius.sso.core.services.security.KeycloakService;
5+
import jakarta.servlet.Filter;
6+
import jakarta.servlet.FilterChain;
7+
import jakarta.servlet.ServletException;
8+
import jakarta.servlet.ServletRequest;
9+
import jakarta.servlet.ServletResponse;
10+
import lombok.extern.slf4j.Slf4j;
11+
import org.keycloak.adapters.springsecurity.token.KeycloakAuthenticationToken;
12+
import org.springframework.security.core.Authentication;
13+
import org.springframework.security.core.context.SecurityContextHolder;
14+
import org.springframework.stereotype.Component;
15+
16+
@Component
17+
@Slf4j
18+
public class AgentKeycloakUserSyncFilter implements Filter {
19+
20+
private final KeycloakService keycloakService;
21+
22+
public AgentKeycloakUserSyncFilter(KeycloakService keycloakService) {
23+
this.keycloakService = keycloakService;
24+
}
25+
26+
@Override
27+
public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
28+
throws IOException, ServletException {
29+
30+
Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
31+
32+
if (authentication instanceof KeycloakAuthenticationToken) {
33+
KeycloakAuthenticationToken keycloakAuth = (KeycloakAuthenticationToken) authentication;
34+
String userId = keycloakAuth.getAccount().getKeycloakSecurityContext().getToken().getSubject();
35+
log.info("Syncing user attributes for user: {}", userId);
36+
37+
38+
}
39+
40+
chain.doFilter(request, response);
41+
}
42+
}

‎ai-agent/src/main/java/io/sentrius/agent/config/GlobalExceptionHandler.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package io.sentrius.agent.config;
2+
3+
import java.util.HashMap;
4+
import java.util.Map;
5+
import com.fasterxml.jackson.databind.node.ObjectNode;
6+
import io.sentrius.sso.core.utils.JsonUtil;
7+
import io.sentrius.sso.core.utils.ZTATUtils;
8+
import lombok.RequiredArgsConstructor;
9+
import org.springframework.http.HttpStatus;
10+
import org.springframework.http.ResponseEntity;
11+
import org.springframework.web.bind.annotation.ControllerAdvice;
12+
import org.springframework.web.bind.annotation.ExceptionHandler;
13+
import org.springframework.web.server.ResponseStatusException;
14+
15+
@ControllerAdvice
16+
@RequiredArgsConstructor
17+
public class GlobalExceptionHandler {
18+
19+
20+
public static String createErrorHash(StackTraceElement[] trace, String t) {
21+
StringBuilder sb = new StringBuilder();
22+
for (StackTraceElement element : trace) {
23+
sb.append(element.toString());
24+
}
25+
sb.append(t);
26+
return ZTATUtils.getCommandHash(sb.toString());
27+
}
28+
29+
@ExceptionHandler(Throwable.class)
30+
public ResponseEntity<Object> handleAllExceptions(Throwable ex) {
31+
ex.printStackTrace();
32+
33+
if (ex instanceof ResponseStatusException rse) {
34+
HttpStatus status = (HttpStatus) rse.getStatusCode();
35+
Map<String, Object> error = new HashMap<>();
36+
error.put("status", status.value());
37+
error.put("error", status.getReasonPhrase());
38+
try{
39+
40+
ObjectNode node = (ObjectNode) JsonUtil.MAPPER.readTree(rse.getReason());
41+
error.put("message", node);
42+
}catch(Exception e){
43+
error.put("message", rse.getReason());
44+
}
45+
46+
return new ResponseEntity<>(error, status);
47+
}
48+
49+
// Default fallback
50+
Map<String, Object> fallback = new HashMap<>();
51+
fallback.put("status", 500);
52+
fallback.put("error", "Internal Server Error");
53+
fallback.put("message", ex.getMessage());
54+
return new ResponseEntity<>(fallback, HttpStatus.INTERNAL_SERVER_ERROR);
55+
}
56+
57+
58+
}

‎ai-agent/src/main/java/io/sentrius/agent/config/SecurityConfig.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package io.sentrius.agent.config;
2+
3+
import java.util.Collection;
4+
import java.util.Collections;
5+
import java.util.List;
6+
import java.util.Map;
7+
import java.util.Optional;
8+
9+
import lombok.RequiredArgsConstructor;
10+
import lombok.extern.slf4j.Slf4j;
11+
import org.springframework.beans.factory.annotation.Value;
12+
import org.springframework.context.annotation.Bean;
13+
import org.springframework.context.annotation.Configuration;
14+
import org.springframework.security.config.Customizer;
15+
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
16+
import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
17+
import org.springframework.security.core.GrantedAuthority;
18+
import org.springframework.security.oauth2.jwt.Jwt;
19+
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationConverter;
20+
import org.springframework.security.oauth2.server.resource.authentication.JwtGrantedAuthoritiesConverter;
21+
import org.springframework.security.web.SecurityFilterChain;
22+
import org.springframework.security.web.authentication.SimpleUrlAuthenticationSuccessHandler;
23+
24+
@Slf4j
25+
@Configuration
26+
@EnableWebSecurity
27+
@RequiredArgsConstructor
28+
public class SecurityConfig {
29+
30+
31+
@Value("${https.required:false}") // Default is false
32+
private boolean httpsRequired;
33+
34+
35+
@Bean
36+
public SecurityFilterChain securityFilterChain(HttpSecurity http) throws Exception {
37+
38+
http
39+
.authorizeHttpRequests(auth -> auth.
40+
requestMatchers("/actuator/**").permitAll() // Public endpoints
41+
.requestMatchers("/**").fullyAuthenticated())
42+
.oauth2ResourceServer(oauth2 -> oauth2
43+
.jwt(jwt -> jwt.jwtAuthenticationConverter(jwtAuthenticationConverterForKeycloak()))
44+
)
45+
.oauth2Login(oauth2 -> oauth2
46+
.loginPage("/oauth2/authorization/keycloak")
47+
.successHandler(new SimpleUrlAuthenticationSuccessHandler())
48+
)
49+
.cors(Customizer.withDefaults());
50+
51+
if (httpsRequired) {
52+
http.requiresChannel(channel -> channel
53+
.requestMatchers("/actuator/**").requiresInsecure() // Allow HTTP for Actuator
54+
.anyRequest().requiresSecure() // Force HTTPS for all other requests
55+
);
56+
}
57+
58+
59+
return http.build();
60+
}
61+
62+
@Bean
63+
public JwtAuthenticationConverter jwtAuthenticationConverterForKeycloak() {
64+
JwtAuthenticationConverter converter = new JwtAuthenticationConverter();
65+
log.info("**** Initializing JwtAuthenticationConverter");
66+
67+
converter.setJwtGrantedAuthoritiesConverter(jwt -> {
68+
log.info("**** Jwt Authentication Converter invoked");
69+
Collection<GrantedAuthority> authorities = new JwtGrantedAuthoritiesConverter().convert(jwt);
70+
log.info("JWT Claims: {}", jwt.getClaims());
71+
72+
String userId = jwt.getClaimAsString("sub");
73+
String username = jwt.getClaimAsString("preferred_username");
74+
String email = jwt.getClaimAsString("email");
75+
76+
log.info("Extracted User Info: userId={}, username={}, email={}", userId, username, email);
77+
78+
79+
return authorities;
80+
});
81+
82+
return converter;
83+
}
84+
85+
86+
private List<String> getRoles(Jwt jwt) {
87+
return Optional.ofNullable(jwt.getClaimAsMap("resource_access"))
88+
.map(resourceAccess -> (Map<String, Object>) resourceAccess.get("sentrius-api"))
89+
.map(client -> (List<String>) ((Map<String, Object>) client).get("roles"))
90+
.orElse(Collections.emptyList());
91+
}
92+
93+
}

0 commit comments

Comments
 (0)