proxy Chat

.local.env

@@ -1,7 +1,7 @@
-SENTRIUS_VERSION=1.1.81
-SENTRIUS_SSH_VERSION=1.1.17
-SENTRIUS_KEYCLOAK_VERSION=1.1.23
-SENTRIUS_AGENT_VERSION=1.1.17
-SENTRIUS_AI_AGENT_VERSION=1.1.32
-LLMPROXY_VERSION=1.0.17
-LAUNCHER_VERSION=1.0.22
+SENTRIUS_VERSION=1.1.95
+SENTRIUS_SSH_VERSION=1.1.18
+SENTRIUS_KEYCLOAK_VERSION=1.1.25
+SENTRIUS_AGENT_VERSION=1.1.18
+SENTRIUS_AI_AGENT_VERSION=1.1.33
+LLMPROXY_VERSION=1.0.18
+LAUNCHER_VERSION=1.0.29

.local.env.bak

@@ -1,7 +1,7 @@
-SENTRIUS_VERSION=1.1.81
-SENTRIUS_SSH_VERSION=1.1.17
-SENTRIUS_KEYCLOAK_VERSION=1.1.22
-SENTRIUS_AGENT_VERSION=1.1.17
-SENTRIUS_AI_AGENT_VERSION=1.1.32
-LLMPROXY_VERSION=1.0.17
-LAUNCHER_VERSION=1.0.22
+SENTRIUS_VERSION=1.1.95
+SENTRIUS_SSH_VERSION=1.1.18
+SENTRIUS_KEYCLOAK_VERSION=1.1.25
+SENTRIUS_AGENT_VERSION=1.1.18
+SENTRIUS_AI_AGENT_VERSION=1.1.33
+LLMPROXY_VERSION=1.0.18
+LAUNCHER_VERSION=1.0.28

agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java

@@ -28,11 +28,19 @@ public class PodLauncherService {
     @Value("${sentrius.agent.registry.version}")
     private String agentVersion;
 
+    @Value("${sentrius.agent.callback.format.url:http://sentrius-agent-%s.%s.svc.cluster.local:8090}")
+    private String callbackFormatUrl;
+
     public PodLauncherService() throws IOException {
         ApiClient client = Config.defaultClient(); // in-cluster or kubeconfig
         this.coreV1Api = new CoreV1Api(client);
     }
 
+    private String buildAgentCallbackUrl(String agentId) {
+        return String.format(callbackFormatUrl, agentId, agentNamespace);
+    }
+
+
     public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception {
         if (agentRegistry != null ) {
             if ("local".equalsIgnoreCase(agentRegistry)) {
@@ -42,6 +50,8 @@ public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception
             }
         }
 
+        var constructedCallbackUrl = buildAgentCallbackUrl(agentId);
+
         String image = String.format("%ssentrius-launchable-agent:%s", agentRegistry, agentVersion);
 
         log.info("Launching agent pod with ID: {}, Image: {}, Callback URL: {}", agentId, image, callbackUrl);
@@ -56,7 +66,9 @@ public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception
                     .imagePullPolicy("IfNotPresent")
 
                     .args(List.of("--spring.config.location=file:/config/agent.properties",
-                        "--agent.namePrefix=" + agentId, "--agent.ai.config=/config/chat-helper.yaml", "--agent.listen.websocket=true"))
+                        "--agent.namePrefix=" + agentId, "--agent.ai.config=/config/chat-helper.yaml", "--agent.listen.websocket=true",
+                        "--agent.callback.url=" + constructedCallbackUrl
+                        ))
                     .resources(new V1ResourceRequirements()
                         .limits(Map.of(
                             "cpu", Quantity.fromString("500m"),

api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java

@@ -3,11 +3,13 @@
 import java.net.URLDecoder;
 import java.nio.charset.StandardCharsets;
 import java.security.GeneralSecurityException;
+import java.security.NoSuchAlgorithmException;
 import java.sql.SQLException;
 import java.time.LocalDateTime;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
@@ -40,6 +42,7 @@
 import io.sentrius.sso.provenance.kafka.ProvenanceKafkaProducer;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
+import jakarta.transaction.Transactional;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.http.HttpStatus;
 import org.jetbrains.annotations.NotNull;
@@ -172,7 +175,7 @@ public ResponseEntity<?> requestRegistration(
 
         // Approve the request if the agent has an active policy ( and it is known and allowed ).
         if (atplPolicyService.getPolicy(operatingUser).isPresent()) {
-            var admin = userService.getUser(UserType.createSystemAdmin().getId());
+            var admin = createOrGetSystemAdmin();
             var approval = ztatService.approveOpsAccessToken(ztatRequest, admin);
 
             return ResponseEntity.ok(Map.of("ztat_token", approval.getToken().toString(), "communication_id",communicationId ));
@@ -186,6 +189,21 @@ public ResponseEntity<?> requestRegistration(
 
 
 
+    }
+
+    @Transactional
+    protected synchronized User createOrGetSystemAdmin() throws NoSuchAlgorithmException {
+        var admin = userService.getUserByUsername("SYSTEM");
+        if (null == admin){
+                var systemAdmin = User.builder()
+                    .username("SYSTEM")
+                    .name("System Admin")
+                    .userId("SYSTEM")
+                    .emailAddress("email").password( userService.encodePassword(UUID.randomUUID().toString())).authorizationType(UserType.createSystemAdmin()).identityType(IdentityType.NON_PERSON_ENTITY);
+                return userService.save(systemAdmin.build());
+        }
+        return admin;
+
     }
 
     @PostMapping("/provenance/submit")

api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java

@@ -75,8 +75,8 @@ public ResponseEntity<HostGroupDTO> setAssignments(HttpServletRequest request, H
         List<User> newUserList = new ArrayList<>();
         for(var userId : (List<String>) payload.get("userIds")) {
             var u = userService.getUser(Long.valueOf(userId));
-            if (null != u) {
-                newUserList.add(u);
+            if (u.isPresent()) {
+                newUserList.add(u.get());
             }
         }
         hg.setUsers(newUserList);

api/src/main/java/io/sentrius/sso/controllers/api/UserApiController.java

@@ -142,9 +142,19 @@ public String deleteUser(@RequestParam("userId") String userId, @RequestParam(re
             log.info("Deleting non-person entity user with id: {}", userId);
             String userIdStr = cryptoService.decrypt(userId);
             var usr = userService.getUserByUserid(userIdStr);
+            if (usr.getId() < 0) {
+                log.info("User with id {} is a system user and cannot be deleted", usr.getId());
+                return "redirect:/sso/v1/users/list?message=" + MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR);
+
+            }
             userService.deleteUser(usr.getId());
         } else {
             Long id = Long.parseLong(cryptoService.decrypt(userId));
+            if (id < 0) {
+                log.info("User with id {} is a system user and cannot be deleted", id);
+                return "redirect:/sso/v1/users/list?message=" +
+                    MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR);
+            }
             userService.deleteUser(id);
         }
         return "redirect:/sso/v1/users/list?message=" + MessagingUtil.getMessageId(MessagingUtil.USER_DELETE_SUCCESS);

api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java

@@ -0,0 +1,20 @@
+package io.sentrius.sso.locator;
+
+
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.stereotype.Component;
+
+import java.net.URI;
+
+@Component
+public class KubernetesAgentLocator {
+
+
+    public URI resolveWebSocketUri(String host, String sessionId, String chatGroupId, String ztat) {
+        // DNS: sentrius-agent-[ID].[namespace].svc.cluster.local
+        ///api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)
+        String fqdn = String.format("%s/api/v1/chat/attach/subscribe?sessionId=%s&chatGroupId=%s&ztat=%s",
+            host,  sessionId, chatGroupId, ztat);
+        return URI.create(fqdn);
+    }
+}

api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java

@@ -189,6 +189,8 @@ public List<SideEffect> initialize(InstallConfiguration installConfiguration, bo
 
         // first we create the admin user, then the user types followed by all users
         sideEffects.addAll(createAdminUser(installConfiguration, action));
+        sideEffects.addAll(createSystemAdmin(installConfiguration, action));
+
 
         createSystemUser(installConfiguration);
 
@@ -681,16 +683,16 @@ protected List<User> createUsers(
                     }
                 }
                 if (action){
-                    user = userService.getUser(user.getId());
+                    var newUser = userService.getUser(user.getId());
                     var definition = userDTO.getAtlpDefinition();
                     if (null != definition && !definition.isEmpty()) {
                     Optional<ATPLPolicyEntity> policy = policyList.stream()
                         .filter(p -> p.getPolicyId().equals(definition))
                         .findFirst();
-                    if (policy.isPresent()) {
-                        atplPolicyService.assignPolicyToUser(user, policy.get());
+                    if (policy.isPresent() & newUser.isPresent()) {
+                        atplPolicyService.assignPolicyToUser(newUser.get(), policy.get());
                     } else {
-                        log.warn("No ATPL policy found for user {} with policy id {}", user.getUsername(),
+                        log.warn("No ATPL policy found for user {} with policy id {}", newUser.get().getUsername(),
                             definition);
                     }
                 }
@@ -809,16 +811,16 @@ protected List<User> createNPEs(
                     }
                 }
                 if (action){
-                    user = userService.getUser(user.getId());
+                    var newUser = userService.getUser(user.getId());
                     var definition = userDTO.getAtlpDefinition();
                     if (null != definition && !definition.isEmpty()) {
                         Optional<ATPLPolicyEntity> policy = policyList.stream()
                             .filter(p -> p.getPolicyId().equals(definition))
                             .findFirst();
                         if (policy.isPresent()) {
-                            atplPolicyService.assignPolicyToUser(user, policy.get());
+                            atplPolicyService.assignPolicyToUser(newUser.get(), policy.get());
                         } else {
-                            log.warn("No ATPL policy found for user {} with policy id {}", user.getUsername(),
+                            log.warn("No ATPL policy found for user {} with policy id {}", newUser.get().getUsername(),
                                 definition);
                         }
                     }
@@ -880,6 +882,50 @@ protected List<SideEffect> createAdminUser(InstallConfiguration installConfigura
         return sideEffects;
     }
 
+    @Transactional
+    public  List<SideEffect> createSystemAdmin(InstallConfiguration installConfiguration, boolean action) throws NoSuchAlgorithmException {
+
+        var user = installConfiguration.getSystemUser();
+
+        if (null == user) {
+            throw new IllegalStateException("Admin user not found in configuration");
+        }
+        List<SideEffect> sideEffects = new ArrayList<>();
+        userService.findByUsername("SYSTEM").ifPresentOrElse(
+            user1 -> {
+                // ignore
+            },
+            () -> {
+                sideEffects.add(SideEffect.builder().sideEffectDescription("Creating admin user " + user.getUsername()).type(
+                    SideEffectType.UPDATE_DATABASE).asset("Users").build());
+                if (action) {
+                    try {
+                        user.setUserId("SYSTEM");
+                        user.setPassword(userService.encodePassword(UUID.randomUUID().toString()));
+                        user.setAuthorizationType(UserType.createSystemAdmin().toDTO());
+                        user.setIdentityType(IdentityType.NON_PERSON_ENTITY.toString());
+
+                        var type =
+                            userService.getUserType(UserType.createSystemAdmin());
+                        if (type.isEmpty()){
+                            type = Optional.of( userService.saveUserType(UserType.createSystemAdmin()) );
+                        }
+
+                        userService.addUscer(User.from(user, type.get()));
+                    } catch (NoSuchAlgorithmException e) {
+                        throw new RuntimeException(e);
+                    }
+
+                    // insert default admin user
+
+                }
+            }
+        );
+
+
+        return sideEffects;
+    }
+
     @Transactional
     protected void createSystemUser(InstallConfiguration connection) throws NoSuchAlgorithmException {
 

api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java

@@ -0,0 +1,39 @@
+package io.sentrius.sso.websocket;
+
+import org.springframework.http.server.ServerHttpRequest;
+import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+
+import java.util.Map;
+
+import org.springframework.web.util.UriComponentsBuilder;
+
+public class AgentHandshakeInterceptor implements HandshakeInterceptor {
+
+    @Override
+    public boolean beforeHandshake(ServerHttpRequest request,
+                                   ServerHttpResponse response,
+                                   WebSocketHandler wsHandler,
+                                   Map<String, Object> attributes) {
+
+        String query = request.getURI().getQuery();
+        Map<String, String> queryParams = UriComponentsBuilder.fromUri(request.getURI()).build().getQueryParams().toSingleValueMap();
+
+        attributes.put("host", queryParams.get("phost"));
+        attributes.put("sessionId", queryParams.get("sessionId"));
+        attributes.put("chatGroupId", queryParams.get("chatGroupId"));
+        attributes.put("ztat", queryParams.get("ztat"));
+
+        return true;
+
+    }
+
+    @Override
+    public void afterHandshake(ServerHttpRequest request,
+                               ServerHttpResponse response,
+                               WebSocketHandler wsHandler,
+                               Exception exception) {
+        // no-op
+    }
+}

api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java

@@ -0,0 +1,61 @@
+package io.sentrius.sso.websocket;
+
+import java.net.URI;
+import java.security.GeneralSecurityException;
+
+import io.sentrius.sso.locator.KubernetesAgentLocator;
+import lombok.RequiredArgsConstructor;
+
+import org.springframework.stereotype.Component;
+import org.springframework.web.reactive.socket.WebSocketHandler;
+import org.springframework.web.reactive.socket.WebSocketMessage;
+import org.springframework.web.reactive.socket.WebSocketSession;
+import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient;
+
+import io.sentrius.sso.core.services.security.CryptoService;
+import lombok.extern.slf4j.Slf4j;
+import reactor.core.publisher.Mono;
+
+@Component
+@Slf4j
+@RequiredArgsConstructor
+public class AgentWebSocketProxyHandler implements WebSocketHandler {
+
+private final KubernetesAgentLocator agentLocator;
+private final CryptoService cryptoService;
+
+@Override
+public Mono<Void> handle(WebSocketSession clientSession) {
+    try {
+        String host = (String) clientSession.getAttributes().get("host");
+        var decryptedHost = cryptoService.decrypt(host); // Ensure host is decrypted if necessary
+        String sessionId = (String) clientSession.getAttributes().get("sessionId");
+        String chatGroupId = (String) clientSession.getAttributes().get("chatGroupId");
+        String ztat = (String) clientSession.getAttributes().get("ztat");
+        log.info("Handling WebSocket connection for host: {}, sessionId: {}, chatGroupId: {}, ztat: {}",
+                decryptedHost, sessionId, chatGroupId, ztat);
+        URI agentUri = agentLocator.resolveWebSocketUri(decryptedHost, sessionId, chatGroupId, ztat);
+
+        ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient();
+
+        return proxyClient.execute(agentUri, agentSession -> {
+            // Forward messages from client to agent
+            Mono<Void> clientToAgent = clientSession.receive()
+                    .map(WebSocketMessage::getPayload)
+                    .map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer))
+                    .as(agentSession::send);
+
+            // Forward messages from agent to client
+            Mono<Void> agentToClient = agentSession.receive()
+                    .map(WebSocketMessage::getPayload)
+                    .map(dataBuffer -> clientSession.binaryMessage(factory -> dataBuffer))
+                    .as(clientSession::send);
+
+            // Run both directions in parallel, complete when both are done
+            return Mono.zip(clientToAgent, agentToClient).then();
+        });
+    } catch (GeneralSecurityException ex) {
+        throw new RuntimeException("Failed to decrypt host", ex);
+    }
+}
+}

api/src/main/java/io/sentrius/sso/websocket/WebSocketConfig.java

@@ -14,6 +14,7 @@ public class WebSocketConfig implements WebSocketConfigurer {
     private final TerminalWSHandler customWebSocketHandler;
     private final AuditSocketHandler auditSocketHandler;
     private final ChatWSHandler chatWSHandler;
+    private final AgentWebSocketProxyHandler agentProxyHandler;
     @Override
     public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
         registry.addHandler(customWebSocketHandler, "/api/v1/ssh/terminal/subscribe")