Skip to content

Commit b3fd68b

Browse files
phrockerphrocker
committed
commit
1 parent 8f19f46 commit b3fd68b

23 files changed

Lines changed: 592 additions & 102 deletions

File tree

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

Lines changed: 51 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
package io.sentrius.agent.analysis.agents.agents;
22

33
import java.io.IOException;
4-
import java.util.List;
5-
import java.util.Map;
4+
import java.util.concurrent.TimeUnit;
65
import com.fasterxml.jackson.core.JsonProcessingException;
7-
import io.jsonwebtoken.lang.Maps;
6+
import com.fasterxml.jackson.databind.node.ArrayNode;
87
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
8+
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
99
import io.sentrius.sso.core.exceptions.ZtatException;
1010
import io.sentrius.sso.core.model.security.Ztat;
11-
import io.sentrius.sso.core.services.agents.LLMService;
1211
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
1312
import io.sentrius.sso.core.dto.UserDTO;
1413
import io.sentrius.sso.core.utils.JsonUtil;
@@ -18,6 +17,7 @@
1817
import org.springframework.boot.context.event.ApplicationReadyEvent;
1918
import org.springframework.context.ApplicationListener;
2019
import org.springframework.stereotype.Component;
20+
import org.springframework.web.client.HttpClientErrorException;
2121

2222
@Slf4j
2323
@Component
@@ -33,6 +33,33 @@ public class RegisteredAgent implements ApplicationListener<ApplicationReadyEven
3333

3434

3535

36+
public ArrayNode promptAgent(UserDTO user) {
37+
ArrayNode response = null;
38+
while(null== response || response.isEmpty()){
39+
try {
40+
log.info("Prompting agent...");
41+
response = agentVerbs.promptAgent(null);
42+
return response;
43+
} catch (ZtatException e) {
44+
log.info("Mechanisms {}" , e.getMechanisms());
45+
var endpoint = zeroTrustClientService.createEndPoingRequest("prompt_agent", e.getEndpoint());
46+
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
47+
.user(user)
48+
.command(endpoint.toString())
49+
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
50+
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
51+
.build();
52+
var request = zeroTrustClientService.requestZtatToken(user,ztatRequestDTO);
53+
54+
var token = zeroTrustClientService.awaitZtatToken(user, request, 60, TimeUnit.MINUTES);
55+
zeroTrustClientService.setZtat(token);
56+
} catch (Exception e) {
57+
e.printStackTrace();
58+
throw new RuntimeException(e);
59+
}
60+
}
61+
return response;
62+
}
3663

3764

3865
@Override
@@ -58,28 +85,35 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
5885
log.info("Registering v1.0.2 agent...");
5986

6087
// register
88+
6189
var register = zeroTrustClientService.registerAgent(user);
6290
log.info("Registered agent is running {} ", register);
6391

6492
var ztat = JsonUtil.MAPPER.readValue(register, Ztat.class);
6593

66-
zeroTrustClientService.setZtat(ztat.getZtatToken());
6794
//while(true){
68-
95+
// get ztat token
6996
try {
7097

98+
99+
100+
zeroTrustClientService.setZtat(ztat.getZtatToken());
101+
71102
// this phase is called "prompting"
72-
var response = agentVerbs.promptAgent(null);
103+
var response = promptAgent(user);
104+
for(var node : response){
105+
if (node.get("verb") != null){
106+
var verb = node.get("verb").asText();
107+
log.info("executing verb is {}", verb);
108+
verbRegistry.execute(verb, null);
109+
}
110+
log.info("node {}", node);
111+
}
73112
log.info("got " + response);
74-
} catch (ZtatException e) {
75-
e.printStackTrace();
76-
//zeroTrustClientService.requestZtatToken()
77-
// we have been requested to get a ztat
78-
} catch (JsonProcessingException e) {
79-
throw new RuntimeException(e);
80-
} catch (IOException e) {
81-
throw new RuntimeException(e);
113+
} catch (HttpClientErrorException e){
114+
log.info("oh boy");
82115
}
116+
83117
// execute command
84118
// var result = verbRegistry.execute(command, null);
85119
//log.info("Command executed: {}", result);
@@ -92,6 +126,8 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
92126

93127
} catch (InterruptedException | JsonProcessingException e) {
94128
throw new RuntimeException(e);
129+
} catch (ZtatException e) {
130+
throw new RuntimeException(e);
95131
}
96132
}
97133

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

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
import java.util.List;
77
import java.util.Map;
88
import com.fasterxml.jackson.core.JsonProcessingException;
9+
import com.fasterxml.jackson.databind.JsonNode;
910
import com.fasterxml.jackson.databind.ObjectMapper;
11+
import com.fasterxml.jackson.databind.node.ArrayNode;
1012
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
1113
import io.sentrius.agent.analysis.agents.agents.AgentConfig;
1214
import io.sentrius.agent.analysis.agents.agents.PromptBuilder;
@@ -15,7 +17,9 @@
1517
import io.sentrius.sso.core.model.verbs.Verb;
1618
import io.sentrius.sso.core.services.agents.LLMService;
1719
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
20+
import io.sentrius.sso.core.utils.JsonUtil;
1821
import io.sentrius.sso.genai.Message;
22+
import io.sentrius.sso.genai.Response;
1923
import io.sentrius.sso.genai.model.ChatRequest;
2024
import lombok.extern.slf4j.Slf4j;
2125
import org.springframework.beans.factory.annotation.Value;
@@ -47,7 +51,7 @@ public AgentVerbs(ZeroTrustClientService zeroTrustClientService, LLMService llmS
4751

4852

4953
@Verb(name = "prompt_agent", description = "Prompts for agent workload.", isAiCallable = false)
50-
public String promptAgent(Map<String, Object> args) throws ZtatException, IOException {
54+
public ArrayNode promptAgent(Map<String, Object> args) throws ZtatException, IOException {
5155
InputStream is = getClass().getClassLoader().getResourceAsStream("agent-config.yaml");
5256
if (is == null) {
5357
throw new RuntimeException("agent-config.yaml not found on classpath");
@@ -62,8 +66,24 @@ public String promptAgent(Map<String, Object> args) throws ZtatException, IOExce
6266
messages.add(Message.builder().role("system").content(prompt).build());
6367

6468
ChatRequest chatRequest = ChatRequest.builder().model("gpt-3.5-turbo").messages(messages).build();
65-
66-
return llmService.askQuestion(chatRequest);
69+
var resp = llmService.askQuestion(chatRequest);
70+
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
71+
log.info("Response is {}", resp);
72+
for (Response.Choice choice : response.getChoices()) {
73+
var content = choice.getMessage().getContent();
74+
log.info("content is {}", content);
75+
if (null != content && !content.isEmpty()){
76+
JsonNode node = JsonUtil.MAPPER.readTree(content);
77+
log.info("Node is {}", node);
78+
if (node.get("plan") != null){
79+
ArrayNode plan = (ArrayNode) node.get("plan");
80+
log.info("Plan is {}", plan);
81+
return plan;
82+
}
83+
}
84+
}
85+
log.info("ahhh");
86+
return JsonUtil.MAPPER.createArrayNode();
6787
}
6888

6989
@Verb(name = "justify_operations", description = "Chats with an agent to justify operations.", isAiCallable = false)

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@
44
import io.sentrius.sso.core.exceptions.ZtatException;
55
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
66
import io.sentrius.sso.core.model.verbs.Verb;
7+
import lombok.extern.slf4j.Slf4j;
78
import org.springframework.stereotype.Service;
89

10+
@Slf4j
911
@Service
1012
public class TerminalVerbs {
1113

@@ -34,6 +36,7 @@ public String listTerminals(Map<String, Object> args) {
3436
try {
3537

3638
var response = zeroTrustClientService.callGetOnApi("/ssh/terminal/list");
39+
log.info("Terminal list response: {}", response);
3740
/*
3841
URL url = new URL(apiUrl);
3942
HttpURLConnection conn = (HttpURLConnection) url.openConnection();

‎ai-agent/src/main/resources/application.properties‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,11 @@ server.port=8083
6666
agent.api.url=http://localhost:8080/
6767
agent.open.ai.endpoint=http://localhost:8084/
6868
agent.ai.config=agent-config.yaml
69+
otel.exporter.otlp.endpoint=${OTEL_EXPORTER_OTLP_ENDPOINT:http://home.guard.local:4317}
70+
otel.traces.exporter=otlp
71+
otel.exporter.otlp.protocol=grpc
72+
otel.metrics.exporter=none
73+
otel.logs.exporter=none
74+
otel.resource.attributes="service.name=ai-agent"
75+
otel.traces.sampler=always_on
76+
otel.exporter.otlp.timeout=10s

‎api/src/main/java/io/sentrius/sso/config/GlobalExceptionHandler.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,12 @@
55
import io.sentrius.sso.core.utils.MessagingUtil;
66
import io.sentrius.sso.core.utils.ZTATUtils;
77
import lombok.RequiredArgsConstructor;
8+
import lombok.extern.slf4j.Slf4j;
89
import org.springframework.web.bind.annotation.ControllerAdvice;
910
import org.springframework.web.bind.annotation.ExceptionHandler;
1011
import org.springframework.web.servlet.mvc.support.RedirectAttributes;
1112

13+
@Slf4j
1214
@ControllerAdvice
1315
@RequiredArgsConstructor
1416
public class GlobalExceptionHandler {
@@ -42,6 +44,8 @@ public String handleAllExceptions(Throwable ex, RedirectAttributes redirectAttri
4244
// Add messageId as a redirect attribute
4345
redirectAttributes.addAttribute("errorId", MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR));
4446

47+
ex.printStackTrace();
48+
log.info("ahhasldigjudaslkgj");
4549
// Redirect to "/mydashboard" with the messageId parameter
4650
return "redirect:/sso/v1/dashboard";
4751
}

‎api/src/main/java/io/sentrius/sso/controllers/api/HostApiController.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -155,11 +155,9 @@ public ResponseEntity<ObjectNode> deleteServer(HttpServletRequest request, HttpS
155155
// operating user
156156
var user = getOperatingUser(request, response);
157157

158-
var hostSystem = hostGroupService.getHostSystem(hostId);
159-
160-
hostGroupService.deleteHostSystem(user, hostSystem.get());
158+
hostGroupService.deleteHostSystem(user, hostId);
161159

162-
node.put("deletedSystemId", hostSystem.get().getId());
160+
node.put("deletedSystemId", hostId);
163161

164162
return ResponseEntity.ok(node);
165163
}

‎api/src/main/java/io/sentrius/sso/controllers/api/ZeroTrustATApiController.java‎

Lines changed: 83 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@
44
import java.sql.SQLException;
55
import java.util.List;
66
import java.util.Map;
7+
import io.sentrius.sso.core.annotations.LimitAccess;
78
import io.sentrius.sso.core.config.SystemOptions;
89
import io.sentrius.sso.core.controllers.BaseController;
910
import io.sentrius.sso.core.dto.JITTrackerDTO;
1011
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
12+
import io.sentrius.sso.core.model.security.enums.ApplicationAccessEnum;
1113
import io.sentrius.sso.core.model.users.User;
14+
import io.sentrius.sso.core.model.zt.ZeroTrustAccessTokenReason;
1215
import io.sentrius.sso.core.services.ErrorOutputService;
1316
import io.sentrius.sso.core.services.NotificationService;
1417
import io.sentrius.sso.core.services.UserService;
@@ -110,7 +113,7 @@ private void manageTerminalZtAt(User operatingUser, Long ztatId, String status)
110113
@PostMapping("/request")
111114
public ResponseEntity<?> requestZtat(
112115
@RequestHeader("Authorization") String token,
113-
@RequestBody ZtatRequestDTO request) {
116+
@RequestBody ZtatRequestDTO ztatRequest, HttpServletRequest request, HttpServletResponse response) {
114117

115118
String compactJwt = token.startsWith("Bearer ") ? token.substring(7) : token;
116119

@@ -120,17 +123,91 @@ public ResponseEntity<?> requestZtat(
120123
return ResponseEntity.status(HttpStatus.SC_UNAUTHORIZED).body("Invalid Keycloak token");
121124
}
122125

126+
// Extract agent identity from the JWT
127+
var operatingUser = getOperatingUser(request, response );
128+
123129
// Extract agent identity from the JWT
124130
String agentId = keycloakService.extractAgentId(compactJwt);
125131

132+
if (null == operatingUser) {
133+
log.warn("No operating user found for agent: {}", agentId);
134+
var username = keycloakService.extractUsername(compactJwt);
135+
operatingUser = userService.getUserWithDetails(username);
136+
137+
}
138+
126139
log.info("Received ZTAT request from agent: {}", agentId);
127140
// Store the request in the database
128-
//var ztatRequest = ztatService.createRequest(agentId, request.getCommand(), request.getCommandHash());
141+
ZeroTrustAccessTokenReason reason = ztatService.createReason(ztatRequest.getJustification(), "", ztatRequest.getCommand());
142+
var submittedZtatRequest = ztatService.createOpsRequest(ztatRequest.getCommand(), ztatRequest.getCommand(),
143+
reason, operatingUser);
144+
submittedZtatRequest = ztatService.addJITRequest(submittedZtatRequest);
145+
146+
return ResponseEntity.ok(Map.of("ztat_request", submittedZtatRequest.getId()));
147+
}
148+
149+
@GetMapping("/status/{type}")
150+
@LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_LOG_IN})
151+
public ResponseEntity<?> getRequest(HttpServletRequest request, HttpServletResponse response,
152+
@RequestHeader("Authorization") String token,
153+
@PathVariable("type") String type,
154+
@RequestParam("ztatId") Long ztatId) throws SQLException, GeneralSecurityException {
155+
String compactJwt = token.startsWith("Bearer ") ? token.substring(7) : token;
156+
157+
158+
log.info("Received ZTAT request from agent: {}", compactJwt);
159+
if (!keycloakService.validateJwt(compactJwt)) {
160+
log.warn("Invalid Keycloak token");
161+
return ResponseEntity.status(HttpStatus.SC_UNAUTHORIZED).body("Invalid Keycloak token");
162+
}
163+
164+
// Extract agent identity from the JWT
165+
var operatingUser = getOperatingUser(request, response );
166+
167+
// Extract agent identity from the JWT
168+
String agentId = keycloakService.extractAgentId(compactJwt);
169+
170+
if (null == operatingUser) {
171+
log.warn("No operating user found for agent: {}", agentId);
172+
var username = keycloakService.extractUsername(compactJwt);
173+
operatingUser = userService.getUserWithDetails(username);
174+
175+
}
176+
177+
if (null != type ){
178+
switch(type){
179+
case "terminal":
180+
var terminalJIT = ztatService.getZtatRequest(ztatId);
181+
if (terminalJIT.getUser().getId() == operatingUser.getId()){
182+
if ( terminalJIT.getApprovals().size() > 0 && terminalJIT.getApprovals().get(0).isApproved() ) {
183+
return ResponseEntity.ok(Map.of("status", "approved", "ztat_token", terminalJIT.getApprovals().get(0).getToken()));
184+
}
185+
else {
186+
log.info("User {} is not the owner of the request {}", operatingUser.getId(), ztatId);
187+
return ResponseEntity.ok(Map.of("status", "unknown"));
188+
}
189+
} else {
190+
log.info("User {} is not the owner of the request {}", operatingUser.getId(), ztatId);
191+
return ResponseEntity.ok(Map.of("status", "unknown"));
192+
}
193+
case "ops":
194+
var opsJit = ztatService.getOpsJITRequest(ztatId);
195+
if (opsJit.getUser().getId() == operatingUser.getId()){
196+
if ( ztatService.isApproved(opsJit) ) {
197+
return ResponseEntity.ok(Map.of("status", "approved", "ztat_token", opsJit.getApprovals().get(0).getToken()));
198+
}
199+
else {
200+
return ResponseEntity.ok(Map.of("status", "unknown"));
201+
}
202+
} else {
203+
return ResponseEntity.ok(Map.of("status", "unknown"));
204+
}
129205

130-
// Generate a Zero Trust Access Token (ZTAT)
131-
//String ztatToken = ztatService.generateZtatToken(ztatRequest);
132-
var ztatToken = "lskejtgsadlkjg";
206+
default:
133207

134-
return ResponseEntity.ok(Map.of("ztat_token", ztatToken));
208+
}
209+
}
210+
return ResponseEntity.ok(Map.of("status", "unknown"));
135211
}
212+
136213
}

‎api/src/main/resources/templates/sso/ztats/view_ztats.html‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,8 +136,10 @@ <h2 class="toast-header">Trust AT (TAT) Management</h2>
136136
<div id="approvedTats" class="tab-pane fade">
137137
<div id="approvedtats-table" class="tab-pane fade show active">
138138
<ul class="nav nav-pills mb-3">
139-
<li class="nav-item"><a class="nav-link active" data-bs-toggle="tab" href="#open-terminal">Terminal TATs</a></li>
140-
<li class="nav-item"><a class="nav-link" data-bs-toggle="tab" href="#open-ops">Operational TATs</a></li>
139+
<li class="nav-item"><a class="nav-link active" data-bs-toggle="tab"
140+
href="#approved-terminal">Terminal TATs</a></li>
141+
<li class="nav-item"><a class="nav-link" data-bs-toggle="tab"
142+
href="#approved-ops">Operational TATs</a></li>
141143
</ul>
142144
<div class="tab-content">
143145
<!-- Terminal TATs -->

0 commit comments

Comments
 (0)