Skip to content

Commit 497e51a

Browse files
phrockerphrocker
committed
Update SSH Proxy
1 parent 24a768a commit 497e51a

File tree

15 files changed

+632
-339
lines changed

15 files changed

+632
-339
lines changed

‎.local.env‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
SENTRIUS_VERSION=1.1.371
1+
SENTRIUS_VERSION=1.1.375
22
SENTRIUS_SSH_VERSION=1.1.41
33
SENTRIUS_KEYCLOAK_VERSION=1.1.53
44
SENTRIUS_AGENT_VERSION=1.1.42
55
SENTRIUS_AI_AGENT_VERSION=1.1.264
66
LLMPROXY_VERSION=1.0.78
77
LAUNCHER_VERSION=1.0.82
88
AGENTPROXY_VERSION=1.0.85
9-
SSHPROXY_VERSION=1.0.40
9+
SSHPROXY_VERSION=1.0.84

‎.local.env.bak‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
SENTRIUS_VERSION=1.1.371
1+
SENTRIUS_VERSION=1.1.375
22
SENTRIUS_SSH_VERSION=1.1.41
33
SENTRIUS_KEYCLOAK_VERSION=1.1.53
44
SENTRIUS_AGENT_VERSION=1.1.42
55
SENTRIUS_AI_AGENT_VERSION=1.1.264
66
LLMPROXY_VERSION=1.0.78
77
LAUNCHER_VERSION=1.0.82
88
AGENTPROXY_VERSION=1.0.85
9-
SSHPROXY_VERSION=1.0.40
9+
SSHPROXY_VERSION=1.0.84

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,9 @@ public String setOption(HttpServletRequest request, HttpServletResponse response
128128
case "java.lang.Float":
129129
results.add(systemOptions.setValue(option.getName(), Float.valueOf(entry.getValue()[0]), false));
130130
break;
131+
case "java.lang.Double":
132+
results.add(systemOptions.setValue(option.getName(), Double.valueOf(entry.getValue()[0]), false));
133+
break;
131134
default:
132135
log.error("Unsupported type: {}", option.getClosestType());
133136
}

‎dataplane/src/main/java/io/sentrius/sso/core/model/sessions/SessionOutput.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,11 +248,13 @@ public AuditOutput waitForOutput(Long time,
248248
}
249249

250250
if (!persistentMessage.isEmpty()){
251+
log.info("Persistent Message: {}", persistentMessage);
251252
var trigger = persistentMessage.pop();
252253
triggers.add( getTrigger(trigger));
253254
}
254255

255256
if (!prompt.isEmpty()){
257+
log.info("Prompt: {}", prompt);
256258
var trigger = prompt.pop();
257259
triggers.add( getTrigger(trigger));
258260
}

‎dataplane/src/main/java/io/sentrius/sso/core/services/SshListenerService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@ public void sendToTerminalSession(String terminalSessionId, ConnectedSystem conn
169169

170170
public void processTerminalMessage(
171171
ConnectedSystem terminalSessionId, Session.TerminalMessage terminalMessage) {
172+
log.info("process terminal messsage");
172173
if (!terminalSessionId.getSession().getClosed() && terminalMessage.getType() != Session.MessageType.HEARTBEAT) {
173174

174175
try {

‎llm-dataplane/src/main/java/io/sentrius/sso/automation/auditing/rules/TwoPartyAIMonitor.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ public Optional<Trigger> trigger(String cmd) {
120120
}
121121

122122
if (llmResponse.get() != null) {
123+
log.info("OpenAI analysis completed. Malicious: {}, response: {}, question: {}", flaggedAsMalicious, llmResponse.get(), llmQuestion.get());
123124
Trigger trg = llmQuestion.get() != null ? new Trigger(TriggerAction.PROMPT_ACTION, llmResponse.get(),
124125
llmQuestion.get()) :
125126
new Trigger(TriggerAction.PERSISTENT_MESSAGE, llmResponse.get());
@@ -183,6 +184,7 @@ public Optional<Trigger> onMessage(Session.TerminalMessage text) {
183184
analysis.get();
184185

185186
if (llmResponse.get() != null && llmQuestion.get() != null) {
187+
log.info("OpenAI analysis completed. Malicious: {}, response: {}, question: {}", flaggedAsMalicious, llmResponse.get(), llmQuestion.get());
186188
Trigger trg = llmQuestion.get() != null && enableLLMQuestions ? new Trigger(TriggerAction.PROMPT_ACTION,
187189
llmResponse.get(),
188190
llmQuestion.get()) :
@@ -197,11 +199,14 @@ public Optional<Trigger> onMessage(Session.TerminalMessage text) {
197199
}
198200
}
199201
if ((connectedSystem.getWebsocketListenerSessionId() == null || connectedSystem.getWebsocketListenerSessionId().isEmpty() ) && flaggedAsMalicious) {
202+
log.info("Flagged as malicious but no websocket session ID available. Returning JIT action.");
200203
if (llmQuestion.get()!= null){
204+
log.info("Flagged as malicious but no websocket session ID available. Returning prompt action.");
201205
Trigger trg = new Trigger(TriggerAction.PROMPT_ACTION, DESCRIPTION);
202206
return Optional.of(trg);
203207
}
204208
else {
209+
log.info("Flagged as malicious but no websocket session ID available. Returning JIT action.");
205210
Trigger trg = new Trigger(TriggerAction.JIT_ACTION, DESCRIPTION);
206211
return Optional.of(trg);
207212
}

‎ssh-proxy/pom.xml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@
3434
<artifactId>sentrius-dataplane</artifactId>
3535
<version>1.0.0-SNAPSHOT</version>
3636
</dependency>
37-
37+
<dependency>
38+
<groupId>io.sentrius</groupId>
39+
<artifactId>llm-dataplane</artifactId>
40+
<version>1.0.0-SNAPSHOT</version>
41+
</dependency>
3842
<!-- Kubernetes client -->
3943
<dependency>
4044
<groupId>io.kubernetes</groupId>

‎ssh-proxy/src/main/java/io/sentrius/sso/sshproxy/config/TaskConfig.java‎

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package io.sentrius.sso.sshproxy.config;
22

33
import java.util.concurrent.Executor;
4+
import java.util.concurrent.ExecutorService;
5+
import java.util.concurrent.ThreadPoolExecutor;
46
import io.sentrius.sso.core.services.TerminalService;
57
import jakarta.annotation.PreDestroy;
8+
import lombok.RequiredArgsConstructor;
69
import lombok.extern.slf4j.Slf4j;
710
import org.springframework.beans.factory.annotation.Autowired;
811
import org.springframework.context.annotation.Bean;
@@ -13,31 +16,37 @@
1316
@Slf4j
1417
@Configuration
1518
@EnableAsync
19+
@RequiredArgsConstructor
1620
public class TaskConfig {
1721

22+
private final TerminalService terminalService;
23+
24+
// Keep a reference so we can shut it down explicitly on destroy, if desired.
1825
private ThreadPoolTaskExecutor executor;
1926

2027
@Bean(name = "taskExecutor")
21-
public Executor taskExecutor() {
22-
ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
23-
executor.setCorePoolSize(15);
24-
executor.setMaxPoolSize(20);
25-
executor.setQueueCapacity(100);
26-
executor.setThreadNamePrefix("SentriusTask-");
27-
executor.initialize();
28-
return executor;
28+
public ThreadPoolTaskExecutor taskExecutor() {
29+
ThreadPoolTaskExecutor exec = new ThreadPoolTaskExecutor();
30+
exec.setCorePoolSize(15);
31+
exec.setMaxPoolSize(20);
32+
exec.setQueueCapacity(100);
33+
exec.setThreadNamePrefix("ProxySession-");
34+
exec.setWaitForTasksToCompleteOnShutdown(true);
35+
exec.setAwaitTerminationSeconds(30);
36+
exec.initialize();
37+
38+
this.executor = exec; // assign the field, not a shadowed local
39+
return exec; // expose as Executor for @Async
2940
}
3041

3142
@PreDestroy
3243
public void shutdownExecutor() {
3344
if (executor != null) {
45+
log.info("Shutting down task executor");
3446
executor.shutdown();
3547
}
36-
log.info("Shutting down executor");
37-
// Call shutdown on SshListenerService to close streams
48+
// If you truly want this on application shutdown:
49+
log.info("Shutting down TerminalService");
3850
terminalService.shutdown();
3951
}
40-
41-
@Autowired
42-
private TerminalService terminalService;
43-
}
52+
}

‎ssh-proxy/src/main/java/io/sentrius/sso/sshproxy/handler/ResponseServiceSession.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public class ResponseServiceSession implements DataSession {
2222
private final InputStream in;
2323
private final OutputStream out;
2424
private final BaseAccessTokenAuditor auditor;
25+
private String persistentMessage = "";
2526
private ConnectedSystem connectedSystem;
2627

2728

@@ -83,7 +84,13 @@ public void sendMessage(WebSocketMessage<?> message) throws IOException {
8384
break;
8485

8586
case PERSISTENT_MESSAGE:
86-
msg = formatPersistentMessage(trigger, auditLog);
87+
if (!persistentMessage.equals(trigger.getDescription())) {
88+
log.info(ANSI_BOLD + "Persistent message: " + ANSI_RESET + trigger.getDescription());
89+
msg = formatPersistentMessage(trigger, auditLog);
90+
}
91+
else {
92+
log.info(ANSI_BOLD + "Persistent message: samesies" + ANSI_RESET);
93+
}
8794
break;
8895
case APPROVE_ACTION:
8996
msg = formatApproveMessage(trigger, auditLog);
@@ -106,6 +113,7 @@ public void sendMessage(WebSocketMessage<?> message) throws IOException {
106113

107114

108115

116+
109117
}
110118
}
111119

@@ -163,11 +171,15 @@ private String formatRecordMessage(Session.Trigger trigger, Session.TerminalMess
163171
}
164172

165173
private String formatPersistentMessage(Session.Trigger trigger, Session.TerminalMessage auditLog) {
174+
if (trigger.getDescription() == null || trigger.getDescription().isEmpty()) {
175+
return "";
176+
}
166177
StringBuilder sb = new StringBuilder();
167178
sb.append("\r\n");
168-
sb.append(ANSI_BLUE).append(ANSI_BOLD).append("💬 MESSAGE").append(ANSI_RESET).append("\r\n");
179+
sb.append(ANSI_BLUE).append(ANSI_BOLD).append("💬 AI Monitor").append(ANSI_RESET).append("\r\n");
169180
sb.append(ANSI_BLUE).append(trigger.getDescription()).append(ANSI_RESET).append("\r\n");
170181
sb.append("\r\n");
182+
persistentMessage = trigger.getDescription();
171183
return sb.toString();
172184
}
173185

0 commit comments

Comments
 (0)