Skip to content

Commit b12b7e6

Browse files
phrockerphrocker
committed
update
1 parent 1dc32b5 commit b12b7e6

Some content is hidden

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

41 files changed

+1140
-162
lines changed

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

Lines changed: 5 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.fasterxml.jackson.core.type.TypeReference;
88
import com.fasterxml.jackson.databind.node.ArrayNode;
99
import com.fasterxml.jackson.databind.node.ObjectNode;
10+
import io.sentrius.agent.analysis.model.AssessedTerminal;
1011
import io.sentrius.agent.analysis.model.Assessment;
1112
import io.sentrius.sso.core.dto.HostSystemDTO;
1213
import io.sentrius.sso.core.model.verbs.ListInterpreter;
@@ -16,7 +17,7 @@
1617
import lombok.extern.slf4j.Slf4j;
1718

1819
@Slf4j
19-
public class AsessmentListInterpreter extends ListInterpreter<Assessment> implements OutputInterpreterIfc {
20+
public class AsessmentListInterpreter extends ListInterpreter<AssessedTerminal> implements OutputInterpreterIfc {
2021

2122
@Override
2223
public Map<String, Object> interpret(VerbResponse input) throws Exception {
@@ -25,38 +26,18 @@ public Map<String, Object> interpret(VerbResponse input) throws Exception {
2526
Map<String,Object> responseMap = new HashMap<>();
2627
responseMap.put("verb.response.type", "list");
2728
responseMap.put("verb.response.map.key", "assessments");
28-
responseMap.put("verb.response.map.type", Assessment.class.getCanonicalName());
29+
responseMap.put("verb.response.map.type", AssessedTerminal.class.getCanonicalName());
2930

3031
if (input.getResponse() instanceof List<?> list) {
3132
log.info("AssessmentListInterpreter: interpret() called with input list");
32-
if (list.isEmpty() || list.get(0) instanceof Assessment) {
33+
if (list.isEmpty() || list.get(0) instanceof AssessedTerminal) {
3334

3435
responseMap.put("assessments", list);
3536
} else {
3637
throw new IllegalArgumentException("Input response is not a List of Assessment objects");
3738
}
3839
} else {
39-
40-
var str = input.getResponse().toString();
41-
log.info("AssessmentListInterpreter: interpret() called with input string {} " ,str);
42-
ArrayNode node = (ArrayNode) JsonUtil.MAPPER.readTree(str);
43-
if (node == null) {
44-
throw new IllegalArgumentException("Input response is not a valid JSON array");
45-
}
46-
List<Assessment> list = new ArrayList<>();
47-
for (int i = 0; i < node.size(); i++) {
48-
var item = node.get(i);
49-
if (item.has("sessionId") && item.has("risk") && item.has("description")) {
50-
Assessment hostSystemDTO = new Assessment();
51-
hostSystemDTO.setSessionId(item.get("sessionId").asText());
52-
hostSystemDTO.setDescription(item.get("description").asText());
53-
hostSystemDTO.setRisk(item.get("risk").asText());
54-
list.add(hostSystemDTO);
55-
} else {
56-
throw new IllegalArgumentException("Input response does not contain required fields");
57-
}
58-
}
59-
responseMap.put("assessments",list);
40+
throw new IllegalArgumentException("Input response does not contain required fields");
6041
}
6142

6243
return responseMap;

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

Lines changed: 229 additions & 46 deletions
Large diffs are not rendered by default.

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

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import io.sentrius.agent.analysis.agents.interpreters.ObjectListInterpreter;
1414
import io.sentrius.agent.analysis.agents.interpreters.TerminalListInterpreter;
1515
import io.sentrius.agent.analysis.agents.interpreters.TerminalOutputInterpreter;
16+
import io.sentrius.agent.analysis.model.AssessedTerminal;
1617
import io.sentrius.agent.analysis.model.Assessment;
1718
import io.sentrius.sso.core.dto.HostSystemDTO;
1819
import io.sentrius.sso.core.dto.ztat.AgentExecution;
@@ -110,12 +111,12 @@ public List<ObjectNode> fetchTerminalOutput(TokenDTO token, List<HostSystemDTO>
110111
" Requires sessionId, risk, and description in a json object.",
111112
requiresTokenManagement = true,
112113
outputInterpreter = TerminalOutputInterpreter.class, inputInterpreter = AsessmentListInterpreter.class)
113-
public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution execution, List<Assessment> dtos)
114+
public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution execution, List<AssessedTerminal> dtos)
114115
throws ZtatException, IOException {
115116
try {
116117
List<ObjectNode> responses = new ArrayList<>();
117118
log.info("Terminal list response: {}", dtos);
118-
for (Assessment dto : dtos) {
119+
for (AssessedTerminal dto : dtos) {
119120

120121
// submit the kill
121122
if (dto != null){
@@ -126,8 +127,8 @@ public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution
126127
}
127128

128129

129-
var risk =dto.getRisk();
130-
var description = dto.getDescription();
130+
var risk =dto.getAssessment().getRisk();
131+
var description = dto.getAssessment().getDescription();
131132
if (null != risk && null != description) {
132133
switch(risk) {
133134
case "low":
@@ -136,13 +137,13 @@ public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution
136137
case "medium":
137138
case "high":
138139
// kill the session
139-
log.info("Killing terminal session: {}", dto.getSessionId());
140+
log.info("Killing terminal session: {}", dto.getAssessment().getSessionId());
140141
break;
141142
default:
142143
throw new RuntimeException("Unknown risk level: " + risk);
143144
}
144145
try {
145-
var sessionId = URLEncoder.encode(dto.getSessionId(), StandardCharsets.UTF_8);
146+
var sessionId = URLEncoder.encode(dto.getAssessment().getSessionId(), StandardCharsets.UTF_8);
146147
var response = zeroTrustClientService.callPutOnApi(
147148
execution, "/ssh/terminal/kill",
148149
Maps.immutableEntry("sessionId", List.of(sessionId))
@@ -151,7 +152,7 @@ public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution
151152
// Successfully retrieved logs
152153
log.info("Terminal output response: {}", response);
153154
var obj = JsonUtil.MAPPER.createObjectNode();
154-
obj.put("id", dto.getSessionId());
155+
obj.put("id", dto.getAssessment().getSessionId());
155156
obj.put("terminalOutput", response);
156157
responses.add(obj);
157158
}
@@ -166,13 +167,19 @@ public List<ObjectNode> killTerminalSessionWithTerminalAssessment(AgentExecution
166167
.justification(description)
167168
.summary("Kill a Terminal session because it is high risk")
168169
.build();
169-
log.info("Obtaining approval. Justification: {}", description);
170+
log.info("Obtaining approval. Justification: {} {}", description, ztatRequestDTO);
170171
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser()
171172
,ztatRequestDTO);
172173

173174
ztatRequestDTO.setRequestId(request);
174175

175-
agentVerbs.justifyAgent(execution, ztatRequestDTO, dto.getDescription());
176+
var token = agentVerbs.justifyAgent(execution, ztatRequestDTO, dto);
177+
execution.setZtatToken(token);
178+
var sessionId = URLEncoder.encode(dto.getAssessment().getSessionId(), StandardCharsets.UTF_8);
179+
var response = zeroTrustClientService.callPutOnApi(
180+
execution, "/ssh/terminal/kill",
181+
Maps.immutableEntry("sessionId", List.of(sessionId))
182+
);
176183
}
177184
}
178185

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package io.sentrius.agent.analysis.model;
2+
3+
import java.util.List;
4+
import io.sentrius.sso.genai.Message;
5+
import lombok.AllArgsConstructor;
6+
import lombok.Builder;
7+
import lombok.Data;
8+
import lombok.Getter;
9+
import lombok.NoArgsConstructor;
10+
11+
@Data
12+
@Builder
13+
@Getter
14+
@NoArgsConstructor
15+
@AllArgsConstructor
16+
public class AssessedTerminal {
17+
Assessment assessment;
18+
List<Message> messages;
19+
20+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,6 @@
1414
public class ZtatAsessment {
1515
String requestId;
1616
boolean approved;
17-
String questionToUser;
17+
String questionToAgent;
1818

1919
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package io.sentrius.agent.analysis.model;
2+
3+
import lombok.AllArgsConstructor;
4+
import lombok.Builder;
5+
import lombok.Data;
6+
import lombok.Getter;
7+
import lombok.NoArgsConstructor;
8+
9+
@Data
10+
@Builder
11+
@Getter
12+
@NoArgsConstructor
13+
@AllArgsConstructor
14+
public class ZtatResponse {
15+
String requestId;
16+
String justificationToAgent;
17+
18+
}

‎ai-agent/src/main/resources/assess-ztat.json‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
[
2-
{
3-
"requestId": "<id>",
4-
"approved": "<true|false>",
5-
"questionToUser": "<question to the agent>"
6-
}
7-
]
1+
{
2+
"requestId": "<id>",
3+
"approved": "<true|false>",
4+
"questionToAgent": "<question to the agent>"
5+
}

‎ai-agent/src/main/resources/challenger-config.yaml‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
description: "Agent that challenges other agents on Access token requests."
22
context: |
3-
Access tokens are used to authenticate and authorize users to access resources and perform operations. Your job
4-
is to challenge other agents on their access token requests. You should ask them to provide a justification for the
5-
task they are doing and apply pushback if it doesn't make sense through the function calls that interact with
6-
the agents. Get the ztat requests from the api server and then create questions for the agents, if any.
7-
3+
Access tokens are used to authenticate and authorize users to access resources and perform operations.
4+
List the ztat requests from the api server and then create questions for the agents, if any. Another agent
5+
is responsible for listing open terminal sessions, you only get the ztat requests and assess their justification.
6+
You will be provided the messages from the agent that lists terminal output in the role of "user". You cannot
7+
communicate with the user, you can only ask questions of the agent requesting the action. Do not request to access
8+
terminal data directly as you don't have access.

‎ai-agent/src/main/resources/respond-ztat.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
[
2+
{
3+
"requestId": "<id>",
4+
"justificationToAgent": "<Justification for why your operation should be approved, responding to the question asked by the agent>",
5+
}
6+
]

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@ public ResponseEntity<String> handleAllExceptions(Throwable ex, RedirectAttribut
5050
return ResponseEntity.status(428).body(responseStatusException.getReason());
5151
}
5252
}
53-
53+
ex.printStackTrace();
54+
log.info("ahhasldigjudaslkgj {}", ex.getMessage());
55+
log.error("asldkjgadlskgj " + ex.getCause(), ex);
5456
String message = "Received Error Message: " + ex.getCause();
5557
ErrorOutput errorOutput = ErrorOutput.builder()
5658
.errorType(ex.getClass().getName())
@@ -64,8 +66,7 @@ public ResponseEntity<String> handleAllExceptions(Throwable ex, RedirectAttribut
6466
// Add messageId as a redirect attribute
6567
redirectAttributes.addAttribute("errorId", MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR));
6668

67-
ex.printStackTrace();
68-
log.info("ahhasldigjudaslkgj {}", ex.getMessage());
69+
6970
// Redirect to "/mydashboard" with the messageId parameter
7071
URI redirectUri = URI.create("/sso/v1/dashboard?errorId=" + MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR));
7172
return ResponseEntity.status(HttpStatus.FOUND).location(redirectUri).build();

0 commit comments

Comments
 (0)