From 6941aa7036a09514ccc75a75391347587c3de86f Mon Sep 17 00:00:00 2001 From: Marc Parisi Date: Tue, 24 Jun 2025 17:59:05 -0400 Subject: [PATCH 1/2] proxy Chat --- .local.env | 14 ++-- .local.env.bak | 14 ++-- .../launcher/service/PodLauncherService.java | 14 +++- .../controllers/api/AgentApiController.java | 20 +++++- .../controllers/api/EnclaveApiController.java | 4 +- .../controllers/api/UserApiController.java | 10 +++ .../sso/locator/KubernetesAgentLocator.java | 24 +++++++ .../startup/ConfigurationApplicationTask.java | 50 +++++++++++++- .../websocket/AgentHandshakeInterceptor.java | 33 ++++++++++ .../websocket/AgentWebSocketProxyHandler.java | 43 ++++++++++++ .../sso/websocket/WebSocketConfig.java | 1 + .../sso/websocket/WebSocketRouteConfig.java | 36 ++++++++++ .../db/migration/V1__Initial_schema.sql | 2 + api/src/main/resources/static/js/chat.js | 13 +++- .../sso/core/services/UserService.java | 4 +- .../configuration/InstallConfiguration.java | 9 +++ docker/keycloak/realms/sentrius-realm.json | 1 + ops-scripts/local/deploy-helm.sh | 1 + .../templates/configmap.yaml | 4 +- .../templates/llm-proxy-alias-service.yaml | 8 +++ sentrius-chart-launcher/values.yaml | 1 + sentrius-chart/templates/configmap.yaml | 65 +++++++++++++++++-- 22 files changed, 341 insertions(+), 30 deletions(-) create mode 100644 api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java create mode 100644 api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java create mode 100644 api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java create mode 100644 api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java create mode 100644 sentrius-chart-launcher/templates/llm-proxy-alias-service.yaml diff --git a/.local.env b/.local.env index 251e76a4..2810524e 100644 --- a/.local.env +++ b/.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 \ No newline at end of file +SENTRIUS_VERSION=1.1.92 +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.27 \ No newline at end of file diff --git a/.local.env.bak b/.local.env.bak index 91ad3e25..164106b3 100644 --- a/.local.env.bak +++ b/.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 \ No newline at end of file +SENTRIUS_VERSION=1.1.92 +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.26 \ No newline at end of file diff --git a/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java b/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java index 6778af12..92f0e8f4 100644 --- a/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java +++ b/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"), diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java index 545ced9c..bb1a0ef8 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java +++ b/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") diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java index 907bc794..8404f4f0 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java @@ -75,8 +75,8 @@ public ResponseEntity setAssignments(HttpServletRequest request, H List newUserList = new ArrayList<>(); for(var userId : (List) 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); diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/UserApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/UserApiController.java index cb6e484b..b7ac26eb 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/UserApiController.java +++ b/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); diff --git a/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java b/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java new file mode 100644 index 00000000..4095c202 --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java @@ -0,0 +1,24 @@ +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 { + + @Value("${sentrius.agent.namespace}") + private String agentNamespace; + + @Value("${sentrius.agent.port:8080}") + private int agentPort; + + public URI resolveWebSocketUri(String agentId) { + // DNS: sentrius-agent-[ID].[namespace].svc.cluster.local + String fqdn = String.format("ws://sentrius-agent-%s.%s.svc.cluster.local:%d/ws", + agentId, agentNamespace, agentPort); + return URI.create(fqdn); + } +} diff --git a/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java b/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java index 28c873a0..b70a2092 100644 --- a/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java +++ b/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java @@ -189,6 +189,8 @@ public List 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,7 +683,7 @@ protected List createUsers( } } if (action){ - user = userService.getUser(user.getId()); + user = userService.getUser(user.getId()).orElseThrow(); var definition = userDTO.getAtlpDefinition(); if (null != definition && !definition.isEmpty()) { Optional policy = policyList.stream() @@ -809,7 +811,7 @@ protected List createNPEs( } } if (action){ - user = userService.getUser(user.getId()); + user = userService.getUser(user.getId()).orElseThrow(); var definition = userDTO.getAtlpDefinition(); if (null != definition && !definition.isEmpty()) { Optional policy = policyList.stream() @@ -880,6 +882,50 @@ protected List createAdminUser(InstallConfiguration installConfigura return sideEffects; } + @Transactional + public List createSystemAdmin(InstallConfiguration installConfiguration, boolean action) throws NoSuchAlgorithmException { + + var user = installConfiguration.getSystemUser(); + + if (null == user) { + throw new IllegalStateException("Admin user not found in configuration"); + } + List 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 { diff --git a/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java b/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java new file mode 100644 index 00000000..6a083ad4 --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java @@ -0,0 +1,33 @@ +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; + +public class AgentHandshakeInterceptor implements HandshakeInterceptor { + + @Override + public boolean beforeHandshake(ServerHttpRequest request, + ServerHttpResponse response, + WebSocketHandler wsHandler, + Map attributes) { + + String path = request.getURI().getPath(); // e.g. /api/v1/agents/ws/agent-123 + String[] segments = path.split("/"); + String agentId = segments[segments.length - 1]; // assumes agentId is at the end + + attributes.put("agentId", agentId); + return true; + } + + @Override + public void afterHandshake(ServerHttpRequest request, + ServerHttpResponse response, + WebSocketHandler wsHandler, + Exception exception) { + // no-op + } +} diff --git a/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java b/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java new file mode 100644 index 00000000..3d0b6b21 --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java @@ -0,0 +1,43 @@ +package io.sentrius.sso.websocket; + +import java.net.URI; +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 reactor.core.publisher.Mono; + +@Component +@RequiredArgsConstructor +public class AgentWebSocketProxyHandler implements WebSocketHandler { + +private final KubernetesAgentLocator agentLocator; + +@Override +public Mono handle(WebSocketSession clientSession) { + String agentId = (String) clientSession.getAttributes().get("agentId"); + URI agentUri = agentLocator.resolveWebSocketUri(agentId); + + ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient(); + + return proxyClient.execute(agentUri, agentSession -> { + // Forward messages from client to agent + Mono clientToAgent = clientSession.receive() + .map(WebSocketMessage::getPayload) + .map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer)) + .as(agentSession::send); + + // Forward messages from agent to client + Mono 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(); + }); +} +} \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/websocket/WebSocketConfig.java b/api/src/main/java/io/sentrius/sso/websocket/WebSocketConfig.java index 031f4a89..961000d1 100644 --- a/api/src/main/java/io/sentrius/sso/websocket/WebSocketConfig.java +++ b/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") diff --git a/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java b/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java new file mode 100644 index 00000000..b9c32294 --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java @@ -0,0 +1,36 @@ +package io.sentrius.sso.websocket; + +import java.util.HashMap; +import java.util.Map; +import lombok.RequiredArgsConstructor; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.web.reactive.HandlerMapping; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.server.WebSocketService; +import org.springframework.web.reactive.socket.server.support.HandshakeWebSocketService; +import org.springframework.web.reactive.socket.server.upgrade.ReactorNettyRequestUpgradeStrategy; +import org.springframework.web.socket.server.support.WebSocketHandlerMapping; + +@Configuration +@RequiredArgsConstructor +public class WebSocketRouteConfig { + + private final AgentWebSocketProxyHandler agentWebSocketProxyHandler; + + @Bean + public WebSocketHandlerMapping webSocketMapping() { + Map map = new HashMap<>(); + map.put("/api/v1/agents/ws/{agentId}", agentWebSocketProxyHandler); + + WebSocketHandlerMapping mapping = new WebSocketHandlerMapping(); + mapping.setUrlMap(map); + mapping.setOrder(-1); // Ensure it's picked up early + return mapping; + } + + @Bean + public WebSocketService webSocketService() { + return new HandshakeWebSocketService(new ReactorNettyRequestUpgradeStrategy()); + } +} diff --git a/api/src/main/resources/db/migration/V1__Initial_schema.sql b/api/src/main/resources/db/migration/V1__Initial_schema.sql index c38bc2d3..7e4e0d85 100644 --- a/api/src/main/resources/db/migration/V1__Initial_schema.sql +++ b/api/src/main/resources/db/migration/V1__Initial_schema.sql @@ -366,6 +366,8 @@ INSERT INTO usertypes (id, user_type_name, automation_access, system_access, rul 'CAN_VIEW_ZTATS', 'CAN_MANAGE_APPLICATION'); + + INSERT INTO usertypes (id, user_type_name, automation_access, system_access, rule_access, user_access, ztat_access, application_access) VALUES (-4, 'Base User', 'CAN_RUN_AUTOMATION', 'CAN_MANAGE_SYSTEMS', 'CAN_VIEW_RULES', 'CAN_MANAGE_USERS', 'CAN_VIEW_ZTATS', diff --git a/api/src/main/resources/static/js/chat.js b/api/src/main/resources/static/js/chat.js index fc886539..e73f34cb 100644 --- a/api/src/main/resources/static/js/chat.js +++ b/api/src/main/resources/static/js/chat.js @@ -237,17 +237,26 @@ export function switchToAgent(agentName,agentId, sessionId, agentHost) { export function sendMessage(event) { console.log("Send message event:", event); - if (event.key !== "Enter") return; + if (event.key !== "Enter"){ + console.log("Key pressed is not Enter, ignoring."); + return; + } const input = document.getElementById("chat-input"); const messageText = input.value.trim(); - if (!messageText) return; + if (!messageText) { + console.log("Empty message, ignoring."); + return; + } const container = document.getElementById("chat-container"); const agentId = container.dataset.agentId; const session = chatSessions.get(agentId); if (session) { session.send(messageText); + } else { + console.error("No active chat session found for agent:", agentId); + return; } input.value = ""; diff --git a/dataplane/src/main/java/io/sentrius/sso/core/services/UserService.java b/dataplane/src/main/java/io/sentrius/sso/core/services/UserService.java index 31f83431..00fae487 100644 --- a/dataplane/src/main/java/io/sentrius/sso/core/services/UserService.java +++ b/dataplane/src/main/java/io/sentrius/sso/core/services/UserService.java @@ -301,8 +301,8 @@ public User addUscer(User user) { * @param userId The ID of the user. * @return The user object if found. */ - public User getUser(Long userId) { - return userRepository.getById(userId); + public Optional getUser(Long userId) { + return userRepository.findById(userId); } /** diff --git a/dataplane/src/main/java/io/sentrius/sso/install/configuration/InstallConfiguration.java b/dataplane/src/main/java/io/sentrius/sso/install/configuration/InstallConfiguration.java index 00957798..eb75b1a6 100644 --- a/dataplane/src/main/java/io/sentrius/sso/install/configuration/InstallConfiguration.java +++ b/dataplane/src/main/java/io/sentrius/sso/install/configuration/InstallConfiguration.java @@ -4,6 +4,7 @@ import java.io.InputStream; import java.util.ArrayList; import java.util.List; +import java.util.UUID; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import io.sentrius.sso.core.dto.HostSystemDTO; @@ -46,6 +47,14 @@ public class InstallConfiguration { .password("changeme") .build(); + @Builder.Default + private UserDTO systemUser = + UserDTO.builder() + .authorizationType(UserType.createSystemAdmin().toDTO()) + .username("SYSTEM") + .password(UUID.randomUUID().toString()) + .build(); + @Builder.Default private List userTypes = new ArrayList<>(); @Builder.Default private List managementGroups = new ArrayList<>(); diff --git a/docker/keycloak/realms/sentrius-realm.json b/docker/keycloak/realms/sentrius-realm.json index 41ba2002..f9c27ee6 100644 --- a/docker/keycloak/realms/sentrius-realm.json +++ b/docker/keycloak/realms/sentrius-realm.json @@ -9,6 +9,7 @@ "secret": "nGkEukexSWTvDzYjSkDmeUlM0FJ5Jhh0", "rootUrl": "${ROOT_URL}", "baseUrl": "${ROOT_URL}", + "serviceAccountsEnabled": true, "redirectUris": [ "${REDIRECT_URIS}/*" ], diff --git a/ops-scripts/local/deploy-helm.sh b/ops-scripts/local/deploy-helm.sh index bdb9eb9e..7491a971 100755 --- a/ops-scripts/local/deploy-helm.sh +++ b/ops-scripts/local/deploy-helm.sh @@ -61,6 +61,7 @@ helm upgrade --install sentrius-agents ./sentrius-chart-launcher --namespace ${T --set sentriusNamespace=${TENANT} \ --set keycloakFQDN=sentrius-keycloak.${TENANT}.svc.cluster.local \ --set sentriusFQDN=sentrius-sentrius.${TENANT}.svc.cluster.local \ + --set llmProxyFQDN=sentrius-llmproxy.${TENANT}.svc.cluster.local \ --set subdomain="sentrius-sentrius" \ --set keycloakSubdomain="sentrius-keycloak" \ --set keycloakHostname="sentrius-keycloak:8081" \ diff --git a/sentrius-chart-launcher/templates/configmap.yaml b/sentrius-chart-launcher/templates/configmap.yaml index 905411e3..216eb4ef 100644 --- a/sentrius-chart-launcher/templates/configmap.yaml +++ b/sentrius-chart-launcher/templates/configmap.yaml @@ -131,6 +131,6 @@ data: sentrius.agent.namespace={{ .Values.tenant }} agents.ai.chat.agent.enabled=true - agent.callback.url=http://localhost:8090 + sentrius.agent.callback.format.url=http://sentrius-agent-%s.%s.svc.cluster.local:8090 agent.api.url={{ .Values.sentriusDomain }} - agent.open.ai.endpoint=http://localhost:8084/ + agent.open.ai.endpoint=http://sentrius-llmproxy:8084/ \ No newline at end of file diff --git a/sentrius-chart-launcher/templates/llm-proxy-alias-service.yaml b/sentrius-chart-launcher/templates/llm-proxy-alias-service.yaml new file mode 100644 index 00000000..b7295b27 --- /dev/null +++ b/sentrius-chart-launcher/templates/llm-proxy-alias-service.yaml @@ -0,0 +1,8 @@ +apiVersion: v1 +kind: Service +metadata: + name: sentrius-llmproxy + namespace: dev-agents +spec: + type: ExternalName + externalName: {{ .Values.sentriusFQDN }} \ No newline at end of file diff --git a/sentrius-chart-launcher/values.yaml b/sentrius-chart-launcher/values.yaml index d7a93300..4a45e039 100644 --- a/sentrius-chart-launcher/values.yaml +++ b/sentrius-chart-launcher/values.yaml @@ -14,6 +14,7 @@ keycloakDomain: https://{{ .Values.keycloakSubdomain }} sentriusDomain: https://{{ .Values.subdomain }} keycloakFQDN: sentrius-keycloak.dev.svc.cluster.local sentriusFQDN: sentrius-sentrius.dev.svc.cluster.local +llmProxyFQDN: sentrius-llmproxy.dev.svc.cluster.local certificates: enabled: false # Disable certs for local; enable for cloud diff --git a/sentrius-chart/templates/configmap.yaml b/sentrius-chart/templates/configmap.yaml index 321fbd5a..b1ba18f4 100644 --- a/sentrius-chart/templates/configmap.yaml +++ b/sentrius-chart/templates/configmap.yaml @@ -288,11 +288,68 @@ data: spring.kafka.properties.retry.backoff.ms=1000 provenance.kafka.topic=sentrius-provenance spring.kafka.bootstrap-servers=sentrius-kafka:9092 - spring.kafka.producer.key-serializer=org.apache.kafka.common.serialization.StringSerializer - spring.kafka.producer.value-serializer=org.springframework.kafka.support.serializer.JsonSerializer - # Optional: trust package to avoid class cast issues with JSON - spring.kafka.producer.properties.spring.json.trusted.packages=io.sentrius.* sentrius.agent.launcher.service=http://sentrius-agents-launcherservice:8080/ + sentrius.agent.register.bootstrap.allow=true + sentrius.agent.bootstrap.policy=default-policy.yaml + default-policy.yaml: | + --- + version: "v0" + description: "Default Policy For Unregistered Agents ( if configured )" + match: + agent_tags: + - "env:prod" + - "classification:observer" + behavior: + minimum_positive_runs: 5 + max_incidents: 1 + incident_types: + denylist: + - "policy_violation" + actions: + on_success: "allow" + on_failure: "deny" + on_marginal: + action: "require_ztat" + ztat_provider: "ztat-service.internal" + capabilities: + primitives: + - id: "accessLLM" + description: "access llm" + endpoint: + - "/api/v1/chat/completions" + composed: + ztat: + provider: "ztat-service.internal" + ttl: "5m" + approved_issuers: + - "http://localhost:8080/" + key_binding: "RSA2048" + approval_required: true + policy_id: "f3326ce2-f46f-405d-94b6-bda2b26db423" + identity: + issuer: "https://keycloak.test.sentrius.cloud" + subject_prefix: "agent-" + mfa_required: true + certificate_authority: "Sentrius-CA" + provenance: + source: "https://test.sentrius.cloud" + signature_required: true + approved_committers: + - "alice@example.com" + runtime: + enclave_required: true + attestation_type: "aws-nitro" + verified_at_boot: true + allow_drift: true + trust_score: + minimum: 80 + marginalThreshold: 50 + weightings: + identity: 0.5 + provenance: 0.2 + runtime: 0.3 + behavior: 0.2 + dynamic.properties: | auditorClass=io.sentrius.sso.automation.auditing.AccessTokenAuditor twopartyapproval.option.LOCKING_SYSTEMS=true From 4d32ac175b0bba5d17080658398a51dcd0645172 Mon Sep 17 00:00:00 2001 From: Marc Parisi Date: Tue, 24 Jun 2025 20:08:38 -0400 Subject: [PATCH 2/2] fix bug --- .local.env | 4 +- .local.env.bak | 4 +- .../sso/locator/KubernetesAgentLocator.java | 12 ++-- .../startup/ConfigurationApplicationTask.java | 14 ++--- .../websocket/AgentHandshakeInterceptor.java | 14 +++-- .../websocket/AgentWebSocketProxyHandler.java | 60 ++++++++++++------- .../sso/websocket/WebSocketRouteConfig.java | 2 +- api/src/main/resources/static/js/chat.js | 5 +- .../core/services/agents/AgentService.java | 8 ++- 9 files changed, 76 insertions(+), 47 deletions(-) diff --git a/.local.env b/.local.env index 2810524e..9d613f13 100644 --- a/.local.env +++ b/.local.env @@ -1,7 +1,7 @@ -SENTRIUS_VERSION=1.1.92 +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.27 \ No newline at end of file +LAUNCHER_VERSION=1.0.29 \ No newline at end of file diff --git a/.local.env.bak b/.local.env.bak index 164106b3..75d14785 100644 --- a/.local.env.bak +++ b/.local.env.bak @@ -1,7 +1,7 @@ -SENTRIUS_VERSION=1.1.92 +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.26 \ No newline at end of file +LAUNCHER_VERSION=1.0.28 \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java b/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java index 4095c202..82c0cd03 100644 --- a/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java +++ b/api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java @@ -9,16 +9,12 @@ @Component public class KubernetesAgentLocator { - @Value("${sentrius.agent.namespace}") - private String agentNamespace; - @Value("${sentrius.agent.port:8080}") - private int agentPort; - - public URI resolveWebSocketUri(String agentId) { + public URI resolveWebSocketUri(String host, String sessionId, String chatGroupId, String ztat) { // DNS: sentrius-agent-[ID].[namespace].svc.cluster.local - String fqdn = String.format("ws://sentrius-agent-%s.%s.svc.cluster.local:%d/ws", - agentId, agentNamespace, agentPort); + ///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); } } diff --git a/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java b/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java index b70a2092..c9034615 100644 --- a/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java +++ b/api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java @@ -683,16 +683,16 @@ protected List createUsers( } } if (action){ - user = userService.getUser(user.getId()).orElseThrow(); + var newUser = userService.getUser(user.getId()); var definition = userDTO.getAtlpDefinition(); if (null != definition && !definition.isEmpty()) { Optional 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); } } @@ -811,16 +811,16 @@ protected List createNPEs( } } if (action){ - user = userService.getUser(user.getId()).orElseThrow(); + var newUser = userService.getUser(user.getId()); var definition = userDTO.getAtlpDefinition(); if (null != definition && !definition.isEmpty()) { Optional 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); } } diff --git a/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java b/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java index 6a083ad4..373e2e9b 100644 --- a/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java +++ b/api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java @@ -7,6 +7,8 @@ import java.util.Map; +import org.springframework.web.util.UriComponentsBuilder; + public class AgentHandshakeInterceptor implements HandshakeInterceptor { @Override @@ -15,12 +17,16 @@ public boolean beforeHandshake(ServerHttpRequest request, WebSocketHandler wsHandler, Map attributes) { - String path = request.getURI().getPath(); // e.g. /api/v1/agents/ws/agent-123 - String[] segments = path.split("/"); - String agentId = segments[segments.length - 1]; // assumes agentId is at the end + String query = request.getURI().getQuery(); + Map 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")); - attributes.put("agentId", agentId); return true; + } @Override diff --git a/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java b/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java index 3d0b6b21..6cefa16e 100644 --- a/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java +++ b/api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java @@ -1,43 +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 handle(WebSocketSession clientSession) { - String agentId = (String) clientSession.getAttributes().get("agentId"); - URI agentUri = agentLocator.resolveWebSocketUri(agentId); - - ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient(); - - return proxyClient.execute(agentUri, agentSession -> { - // Forward messages from client to agent - Mono clientToAgent = clientSession.receive() - .map(WebSocketMessage::getPayload) - .map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer)) - .as(agentSession::send); - - // Forward messages from agent to client - Mono 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(); - }); + 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 clientToAgent = clientSession.receive() + .map(WebSocketMessage::getPayload) + .map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer)) + .as(agentSession::send); + + // Forward messages from agent to client + Mono 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); + } } } \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java b/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java index b9c32294..62b554aa 100644 --- a/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java +++ b/api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java @@ -21,7 +21,7 @@ public class WebSocketRouteConfig { @Bean public WebSocketHandlerMapping webSocketMapping() { Map map = new HashMap<>(); - map.put("/api/v1/agents/ws/{agentId}", agentWebSocketProxyHandler); + map.put("/api/v1/agents/ws/{host}/{sessionId}/{chatGroupId}/{ztat}", agentWebSocketProxyHandler); WebSocketHandlerMapping mapping = new WebSocketHandlerMapping(); mapping.setUrlMap(map); diff --git a/api/src/main/resources/static/js/chat.js b/api/src/main/resources/static/js/chat.js index e73f34cb..2f57b3c2 100644 --- a/api/src/main/resources/static/js/chat.js +++ b/api/src/main/resources/static/js/chat.js @@ -73,7 +73,10 @@ class ChatSession { // Step 3: Open WebSocket with ZTAT token - const uri = `${phost}/api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)}`; + //const uri = `${phost}/api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)}`; + //const uri = `/api/v1/agents/ws/${encodeURIComponent(phost)}/${encodeURIComponent(this.sessionId)}/${encodeURIComponent(this.chatGroupId)}/${encodeURIComponent(jwt)}`; + const uri = `/api/v1/agents/ws?phost=${encodeURIComponent(phost)}&sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${encodeURIComponent(this.chatGroupId)}&jwt=${encodeURIComponent(jwt)}`; + console.log("Connecting to chat server with ZTAT at:", uri); this.connection = new WebSocket(uri); diff --git a/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java b/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java index 69d52cba..6f6c7cf8 100644 --- a/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java +++ b/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java @@ -129,7 +129,13 @@ public List getAllAgents(boolean encryptId, List filteredIds, dtoBuilder.agentName(heartbeat.getAgentName()); var callback = callbackUrls.get(heartbeat.getAgentId()); if (callback != null) { - dtoBuilder.agentCallback(callback); + try { + + var encryptedCallback = cryptoService.encrypt(callback); // Ensure callback is decrypted + dtoBuilder.agentCallback(encryptedCallback); + } catch (GeneralSecurityException e) { + throw new RuntimeException("Error encrypting callback URL", e); + } } } if (encryptId){