Skip to content

Commit 3150e3e

Browse files
phrockerphrocker
committed
UPdate
1 parent f8e19ff commit 3150e3e

File tree

90 files changed

+3072
-99
lines changed

Some content is hidden

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

90 files changed

+3072
-99
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,8 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
135135
try {
136136

137137
Thread.sleep(5_000);
138-
} catch (InterruptedException ex) {
138+
agentClientService.heartbeat(execution, execution.getUser().getUsername());
139+
} catch (InterruptedException | ZtatException ex) {
139140
throw new RuntimeException(ex);
140141
}
141142

ai-agent/src/main/java/io/sentrius/agent/analysis/api/AgentController.java renamed to ai-agent/src/main/java/io/sentrius/agent/analysis/api/BotController.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
@Slf4j
1818
@RestController
1919
@RequestMapping("/api/v1/agent")
20-
public class AgentController {
20+
public class BotController {
2121

2222
KeycloakService keycloakService;
2323
AgentKeyService agentKeyService;

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import java.security.GeneralSecurityException;
77
import java.util.Base64;
88
import java.util.Map;
9-
import java.util.concurrent.ConcurrentHashMap;
109
import java.util.stream.Collectors;
1110
import java.util.stream.Stream;
1211
import io.sentrius.agent.analysis.api.UserCommunicationService;

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

Lines changed: 30 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import io.sentrius.sso.core.services.ATPLPolicyService;
1515
import io.sentrius.sso.core.services.ErrorOutputService;
1616
import io.sentrius.sso.core.services.UserService;
17+
import io.sentrius.sso.core.services.agents.AgentClientService;
1718
import io.sentrius.sso.core.services.agents.AgentService;
1819
import io.sentrius.sso.core.services.auditing.AuditService;
1920
import io.sentrius.sso.core.services.security.CryptoService;
@@ -44,6 +45,7 @@ public class AgentBootstrapController extends BaseController {
4445
final ZeroTrustAccessTokenService ztatService;
4546
final ZeroTrustRequestService ztrService;
4647
final AgentService agentService;
48+
private final AgentClientService agentClientService;
4749

4850

4951
@Value("${sentrius.agent.register.bootstrap.allow:false}")
@@ -59,7 +61,8 @@ public AgentBootstrapController(
5961
AuditService auditService,
6062
CryptoService cryptoService, SessionTrackingService sessionTrackingService, KeycloakService keycloakService,
6163
ATPLPolicyService atplPolicyService,
62-
ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService
64+
ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService,
65+
AgentClientService agentClientService
6366
) {
6467
super(userService, systemOptions, errorOutputService);
6568
this.auditService = auditService;
@@ -70,6 +73,7 @@ public AgentBootstrapController(
7073
this.ztatService = ztatService;
7174
this.ztrService = ztrService;
7275
this.agentService = agentService;
76+
this.agentClientService = agentClientService;
7377
}
7478

7579

@@ -97,30 +101,34 @@ public ResponseEntity<AgentRegistrationDTO> bootstrap(
97101
log.info("Registering {}", registrationDTO.getAgentName());
98102
User user = userService.getUserByUsername(newDTO.getAgentName());
99103
if (user == null) {
100-
var type = userService.getUserType(
101-
UserType.createUnknownUser());
102-
103-
user = User.builder()
104-
.username(newDTO.getAgentName())
105-
.name(newDTO.getAgentName())
106-
.emailAddress(newDTO.getAgentName())
107-
.userId(unencryptedRegistration.getClientId())
108-
.authorizationType(type.get())
109-
.identityType(IdentityType.NON_PERSON_ENTITY)
110-
.build();
111-
log.info("Creating new user: {}", user);
112-
userService.save(user);
113-
}
114-
try(InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream(defaultPolicyFile)) {
115-
if (terminalHelperStream == null) {
116-
throw new RuntimeException(defaultPolicyFile + "not found on classpath");
117-
104+
var type = userService.getUserType(
105+
UserType.createUnknownUser());
106+
107+
user = User.builder()
108+
.username(newDTO.getAgentName())
109+
.name(newDTO.getAgentName())
110+
.emailAddress(newDTO.getAgentName())
111+
.userId(unencryptedRegistration.getClientId())
112+
.authorizationType(type.get())
113+
.identityType(IdentityType.NON_PERSON_ENTITY)
114+
.build();
115+
log.info("Creating new user: {}", user);
116+
userService.save(user);
117+
118+
try(InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream(defaultPolicyFile)) {
119+
if (terminalHelperStream == null) {
120+
throw new RuntimeException(defaultPolicyFile + "not found on classpath");
121+
122+
}
123+
String defaultYaml = new String(terminalHelperStream.readAllBytes());
124+
log.info("Default policy file: {}", defaultPolicyFile);
125+
var policy = atplPolicyService.createPolicy(user, defaultYaml);
118126
}
119-
String defaultYaml = new String(terminalHelperStream.readAllBytes());
120-
log.info("Default policy file: {}", defaultPolicyFile);
121-
var policy = atplPolicyService.createPolicy(user, defaultYaml);
127+
122128
}
123129

130+
agentService.setCallBack(user, registrationDTO.getAgentCallbackUrl());
131+
124132
}
125133
else {
126134
log.info("Not Registering {}", registrationDTO.getAgentName());
@@ -129,14 +137,6 @@ public ResponseEntity<AgentRegistrationDTO> bootstrap(
129137
return ResponseEntity.ok(newDTO);
130138
}
131139

132-
@GetMapping("/register")
133-
// no LimitAccess
134-
public ResponseEntity<?> ohhia(
135-
136-
HttpServletRequest request, HttpServletResponse response) throws SQLException, GeneralSecurityException {
137-
138-
return ResponseEntity.ok("ohhiai");
139-
}
140140

141141

142142
}

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,17 @@
22

33
import java.security.GeneralSecurityException;
44
import java.time.ZoneOffset;
5+
import java.util.ArrayList;
56
import java.util.List;
67
import java.util.stream.Collectors;
78
import io.sentrius.sso.core.config.SystemOptions;
89
import io.sentrius.sso.core.controllers.BaseController;
10+
import io.sentrius.sso.core.dto.AgentDTO;
911
import io.sentrius.sso.core.services.UserService;
12+
import io.sentrius.sso.core.services.agents.AgentService;
1013
import io.sentrius.sso.core.services.security.CryptoService;
1114
import io.sentrius.sso.core.utils.AccessUtil;
15+
import io.sentrius.sso.genai.Response;
1216
import io.sentrius.sso.protobuf.Session.ChatMessage;
1317
import io.sentrius.sso.core.model.security.enums.SSHAccessEnum;
1418
import io.sentrius.sso.core.model.sessions.SessionLog;
@@ -36,19 +40,22 @@ public class ChatApiController extends BaseController {
3640
final CryptoService cryptoService;
3741
final SessionTrackingService sessionTrackingService;
3842
final ChatLogRepository chatLogRepository;
43+
final AgentService agentService;
3944

4045
public ChatApiController(
4146
UserService userService,
4247
SystemOptions systemOptions,
4348
ErrorOutputService errorOutputService,
4449
AuditService auditService,
45-
CryptoService cryptoService, SessionTrackingService sessionTrackingService, ChatLogRepository chatLogRepository
50+
CryptoService cryptoService, SessionTrackingService sessionTrackingService, ChatLogRepository chatLogRepository,
51+
AgentService agentService
4652
) {
4753
super(userService, systemOptions, errorOutputService);
4854
this.auditService = auditService;
4955
this.cryptoService = cryptoService;
5056
this.sessionTrackingService = sessionTrackingService;
5157
this.chatLogRepository = chatLogRepository;
58+
this.agentService = agentService;
5259
}
5360

5461
public SessionLog createSession(@RequestParam String username, @RequestParam String ipAddress) {
@@ -92,6 +99,21 @@ public ResponseEntity<List<ChatMessage>> getChatHistory(
9299
return ResponseEntity.ok(messages);
93100
}
94101

102+
@GetMapping("/agent/list")
103+
public ResponseEntity<List<AgentDTO>> listAvailableAgents(
104+
HttpServletRequest request,
105+
HttpServletResponse response)
106+
throws GeneralSecurityException {
107+
108+
List<AgentDTO> availableAgents = agentService.getAvailableAgents();
109+
110+
111+
112+
// get a list of registered agents.
113+
114+
return ResponseEntity.ok(availableAgents);
115+
}
116+
95117

96118

97119

‎api/src/main/resources/static/js/chat.js‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
// =========================
2+
// Section: Globals & Setup
3+
// =========================
4+
const chatSessions = new Map(); // key: agentId, value: ChatSession
5+
6+
window.addEventListener("beforeunload", persistChatSessions);
7+
8+
// Restore on page load
9+
(function restoreSessions() {
10+
const saved = localStorage.getItem("openChats");
11+
if (!saved) return;
12+
13+
const chatData = JSON.parse(saved);
14+
for (const [agentId, data] of Object.entries(chatData)) {
15+
const session = new ChatSession(data.agentId, data.sessionId, data.agentHost, data.messages);
16+
chatSessions.set(agentId, session);
17+
}
18+
})();
19+
20+
// =========================
21+
// Section: ChatSession Class
22+
// =========================
23+
class ChatSession {
24+
constructor(agentId, sessionId, agentHost, preloadMessages = []) {
25+
this.agentId = agentId;
26+
this.sessionId = sessionId;
27+
this.agentHost = agentHost;
28+
this.chatGroupId = `${agentId}-${sessionId}`;
29+
this.messages = preloadMessages || [];
30+
this.connection = null;
31+
32+
this.connect();
33+
}
34+
35+
connect() {
36+
const protocol = location.protocol === "https:" ? "https" : "http";
37+
const uri = `${protocol}://${this.agentHost}/api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}`;
38+
this.connection = new SockJS(uri);
39+
this.connection.onmessage = (e) => this.handleMessage(e);
40+
this.connection.onopen = () => this.heartbeat();
41+
}
42+
43+
handleMessage(e) {
44+
try {
45+
const binary = Uint8Array.from(atob(e.data), c => c.charCodeAt(0));
46+
const msg = proto.io.sentrius.protobuf.ChatMessage.deserializeBinary(binary);
47+
const sender = msg.getSender();
48+
const message = msg.getMessage();
49+
this.messages.push({ sender, message });
50+
51+
const activeAgentId = document.getElementById("chat-container").dataset.agentId;
52+
if (activeAgentId === this.agentId) {
53+
appendToChatWindow(sender, message);
54+
}
55+
56+
persistChatSessions();
57+
} catch (err) {
58+
console.error("Failed to handle chat message", err);
59+
}
60+
}
61+
62+
send(text) {
63+
const msg = new proto.io.sentrius.protobuf.ChatMessage();
64+
msg.setSender("user");
65+
msg.setMessage(text);
66+
67+
this.connection.send(btoa(String.fromCharCode(...msg.serializeBinary())));
68+
this.messages.push({ sender: "You", message: text });
69+
70+
const activeAgentId = document.getElementById("chat-container").dataset.agentId;
71+
if (activeAgentId === this.agentId) {
72+
appendToChatWindow("You", text);
73+
}
74+
75+
persistChatSessions();
76+
}
77+
78+
heartbeat() {
79+
if (!this.connection || this.connection.readyState !== 1) return;
80+
81+
const msg = new proto.io.sentrius.protobuf.ChatMessage();
82+
msg.setSender("system");
83+
msg.setMessage("heartbeat");
84+
85+
this.connection.send(btoa(String.fromCharCode(...msg.serializeBinary())));
86+
setTimeout(() => this.heartbeat(), 5000);
87+
}
88+
89+
toJSON() {
90+
return {
91+
agentId: this.agentId,
92+
sessionId: this.sessionId,
93+
agentHost: this.agentHost,
94+
messages: this.messages
95+
};
96+
}
97+
}
98+
99+
// =========================
100+
// Section: UI Interaction
101+
// =========================
102+
103+
function switchToAgent(agentId, sessionId, agentHost) {
104+
let session = chatSessions.get(agentId);
105+
if (!session) {
106+
session = new ChatSession(agentId, sessionId, agentHost);
107+
chatSessions.set(agentId, session);
108+
}
109+
110+
const container = document.getElementById("chat-container");
111+
container.dataset.agentId = agentId;
112+
container.dataset.agentHost = agentHost;
113+
document.getElementById("chat-agent-name").textContent = agentId;
114+
115+
const messagesDiv = document.getElementById("chat-messages");
116+
messagesDiv.innerHTML = "";
117+
session.messages.forEach(msg => {
118+
appendToChatWindow(msg.sender, msg.message);
119+
});
120+
121+
container.style.display = "block";
122+
}
123+
124+
function sendMessage(event) {
125+
if (event.key !== "Enter") return;
126+
127+
const input = document.getElementById("chat-input");
128+
const messageText = input.value.trim();
129+
if (!messageText) return;
130+
131+
const container = document.getElementById("chat-container");
132+
const agentId = container.dataset.agentId;
133+
const session = chatSessions.get(agentId);
134+
if (session) {
135+
session.send(messageText);
136+
}
137+
138+
input.value = "";
139+
}
140+
141+
function appendToChatWindow(sender, message) {
142+
const chatBox = document.getElementById("chat-messages");
143+
const div = document.createElement("div");
144+
div.classList.add("chat-message");
145+
div.innerHTML = `<strong>${sender}:</strong> ${message}`;
146+
chatBox.appendChild(div);
147+
chatBox.scrollTop = chatBox.scrollHeight;
148+
}
149+
150+
function toggleChat() {
151+
const container = document.getElementById("chat-container");
152+
container.classList.toggle("hidden");
153+
}
154+
155+
// =========================
156+
// Section: Persistence
157+
// =========================
158+
function persistChatSessions() {
159+
const obj = {};
160+
chatSessions.forEach((session, agentId) => {
161+
obj[agentId] = session.toJSON();
162+
});
163+
localStorage.setItem("openChats", JSON.stringify(obj));
164+
}

0 commit comments

Comments
 (0)