Update counts when connected to chat with agents

Author: phrocker

.local.env

@@ -1,8 +1,8 @@
-SENTRIUS_VERSION=1.1.171
+SENTRIUS_VERSION=1.1.178
 SENTRIUS_SSH_VERSION=1.1.33
 SENTRIUS_KEYCLOAK_VERSION=1.1.46
 SENTRIUS_AGENT_VERSION=1.1.33
 SENTRIUS_AI_AGENT_VERSION=1.1.60
 LLMPROXY_VERSION=1.0.43
 LAUNCHER_VERSION=1.0.50
-AGENTPROXY_VERSION=1.0.60
+AGENTPROXY_VERSION=1.0.63

.local.env.bak

@@ -1,8 +1,8 @@
-SENTRIUS_VERSION=1.1.171
+SENTRIUS_VERSION=1.1.178
 SENTRIUS_SSH_VERSION=1.1.33
 SENTRIUS_KEYCLOAK_VERSION=1.1.46
 SENTRIUS_AGENT_VERSION=1.1.33
 SENTRIUS_AI_AGENT_VERSION=1.1.60
 LLMPROXY_VERSION=1.0.43
 LAUNCHER_VERSION=1.0.50
-AGENTPROXY_VERSION=1.0.60
+AGENTPROXY_VERSION=1.0.63

agent-proxy/src/main/java/io/sentrius/sso/controller/SessionController.java

@@ -0,0 +1,30 @@
+package io.sentrius.sso.controller;
+
+import java.util.List;
+import io.sentrius.sso.core.dto.TerminalLogDTO;
+import io.sentrius.sso.service.ActiveWebSocketSessionManager;
+import org.springframework.http.HttpHeaders;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.reactive.function.client.WebClient;
+import reactor.core.publisher.Mono;
+
+@RestController
+@RequestMapping("/api/v1/sessions")
+public class SessionController {
+
+    private final WebClient webClient;
+    private final ActiveWebSocketSessionManager activeWebSocketSessionManager;
+
+    public SessionController(WebClient.Builder builder, ActiveWebSocketSessionManager activeWebSocketSessionManager) {
+        this.webClient = builder.baseUrl("http://sentrius-agent-proxy").build();
+        this.activeWebSocketSessionManager = activeWebSocketSessionManager;
+    }
+
+    @GetMapping("/list")
+    public List<TerminalLogDTO> listSessions() {
+        return activeWebSocketSessionManager.getActiveSessions();
+    }
+
+}

agent-proxy/src/main/java/io/sentrius/sso/service/ActiveWebSocketSessionManager.java

@@ -1,7 +1,12 @@
 package io.sentrius.sso.service;
 
+import java.sql.Timestamp;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+import io.sentrius.sso.core.dto.TerminalLogDTO;
 import org.springframework.stereotype.Component;
 import org.springframework.web.reactive.socket.WebSocketSession;
 
@@ -20,4 +25,16 @@ public void unregister(String sessionId) {
     public WebSocketSession get(String sessionId) {
         return sessions.get(sessionId);
     }
+
+    public List<TerminalLogDTO> getActiveSessions() {
+        return sessions.values().stream()
+            .map(session -> TerminalLogDTO.builder()
+                .sessionId(session.getId())
+                .user("Unknown")
+                .host(Objects.requireNonNull(session.getHandshakeInfo().getRemoteAddress()).toString())
+                .closed(!session.isOpen())
+                .sessionTime(new Timestamp(System.currentTimeMillis()))
+                .build())
+            .collect(Collectors.toList());
+    }
 }

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

@@ -137,7 +137,9 @@ public ResponseEntity<?> heartbeat(
             log.warn("Heartbeat status is empty");
             return ResponseEntity.status(HttpStatus.SC_BAD_REQUEST).body("Heartbeat status is empty");
         }
+
         agentService.recordHeartbeat(operatingUser.getUserId(),status.getName(), status);
+        agentService.setCallBack(operatingUser, status.getAgentUrl());
         log.info("Heartbeat status recorded for agent: {} {}", agentId, status);
         return ResponseEntity.ok(Map.of("status", "success"));
     }

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

@@ -1,9 +1,12 @@
 package io.sentrius.sso.controllers.api;
 
 import java.security.GeneralSecurityException;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import io.sentrius.sso.config.AppConfig;
 import io.sentrius.sso.core.annotations.LimitAccess;
 import io.sentrius.sso.core.config.SystemOptions;
 import io.sentrius.sso.core.controllers.BaseController;
@@ -17,12 +20,18 @@
 import io.sentrius.sso.core.services.UserService;
 import io.sentrius.sso.core.services.auditing.AuditService;
 import io.sentrius.sso.core.services.security.CryptoService;
+import io.sentrius.sso.core.services.security.KeycloakService;
 import io.sentrius.sso.core.services.terminal.SessionTrackingService;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.http.HttpEntity;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.*;
+import org.springframework.web.client.RestTemplate;
+import org.springframework.web.reactive.function.client.WebClient;
 
 @Slf4j
 @RestController
@@ -31,18 +40,32 @@ public class AuditApiController extends BaseController {
     private final AuditService auditService;
     final CryptoService cryptoService;
     final SessionTrackingService sessionTrackingService;
+    final AppConfig appConfig;
+    final RestTemplate restTemplate = new RestTemplate();
+    final KeycloakService keycloakService;
+
+    private WebClient webClient;
 
     public AuditApiController(
         UserService userService,
         SystemOptions systemOptions,
         ErrorOutputService errorOutputService,
         AuditService auditService,
-        CryptoService cryptoService,SessionTrackingService sessionTrackingService
+        CryptoService cryptoService, SessionTrackingService sessionTrackingService, AppConfig appConfig,
+        KeycloakService keycloakService
     ) {
         super(userService, systemOptions, errorOutputService);
         this.auditService = auditService;
         this.cryptoService = cryptoService;
         this.sessionTrackingService = sessionTrackingService;
+        this.appConfig = appConfig;
+        this.keycloakService = keycloakService;
+        try {
+            this.webClient = WebClient.builder().baseUrl(appConfig.getAgentProxyExternalUrl()).build();
+        }
+        catch (Exception e) {
+            this.webClient = null;
+        }
     }
 
     public SessionLog createSession(@RequestParam String username, @RequestParam String ipAddress) {
@@ -57,14 +80,44 @@ public List<TerminalLogs> getLogs(@PathVariable Long sessionId) {
     @GetMapping("/list")
     @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_LOG_IN}, sshAccess = {SSHAccessEnum.CAN_VIEW_SYSTEMS})
     public List<TerminalLogDTO> listOpenSessions(HttpServletRequest request, HttpServletResponse response) {
-        return sessionTrackingService.getConnectedSession().stream().map(
+        List<TerminalLogDTO> dtos = new ArrayList<>();
+        sessionTrackingService.getConnectedSession().stream().map(
             x -> {
                 try {
                     return x.toDTO(cryptoService.encrypt(x.getSession().getId().toString()));
                 } catch (GeneralSecurityException e) {
                     throw new RuntimeException(e);
                 }
-            }).collect(Collectors.toList());
+            }).forEach(dtos::add);
+
+        var agentProxyUrl = appConfig.getAgentProxyExternalUrl();
+
+
+        try {
+
+            if (null != webClient) {
+                HttpHeaders headers = new HttpHeaders();
+                headers.setBearerAuth(keycloakService.getKeycloakToken()); // or however you're retrieving the token
+                HttpEntity<?> requestEntity = new HttpEntity<>(headers);
+
+                var agentDtos = webClient.get()
+                    .uri("/api/v1/sessions/list")
+                    .header(HttpHeaders.AUTHORIZATION, "Bearer " + keycloakService.getKeycloakToken())
+                    .retrieve()
+                    .bodyToFlux(TerminalLogDTO.class)
+                    .collectList()
+                    .block(); // blocks for compatibility with a synchronous controller
+
+                if (agentDtos != null) dtos.addAll(agentDtos);
+            } else {
+                log.warn("Agent Proxy URL is not configured or WebClient could not be initialized. Cannot retrieve agent proxy sessions.");
+            }
+        } catch (Exception e) {
+            log.warn("Failed to retrieve agent proxy sessions: {}", e.getMessage());
+        }
+
+
+        return dtos;
     }
 
     @GetMapping("/audit/list")

api/src/main/resources/templates/fragments/dashboard/dashboard_cards.html

@@ -4,7 +4,7 @@
         <i class="fas fa-cog card-icon"></i> <!-- Icon for the card -->
     </div>
     <div class="card-body">
-        <h5 class="card-title">Open SSH Connections</h5>
+        <h5 class="card-title">Open Connections</h5>
         <p class="card-text"><span id="open-cxn-count"></span></p>
     </div>
 </div>

dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java

@@ -109,11 +109,13 @@ public List<AgentDTO> getAllAgents(boolean encryptId, List<String> filteredIds,
             .filter(heartbeat -> {
                 var user = userService.getUserByUsername(heartbeat.getAgentName());
                 log.info("Agent {}: {}", heartbeat.getAgentId(), user);
-                log.info("Excluding {}: {}", heartbeat.getAgentId(), filteredIds);
+                log.info("Excluding {}: {} -- include/exclude {}", heartbeat.getAgentId(), filteredIds, include);
                 if (include){
+                    log.info("Including {}: {}", heartbeat.getAgentId(), filteredIds);
                     return user != null && filteredIds.contains(user.getUserId());
                 }
                 else {
+                    log.info("Excluding {}: {}", heartbeat.getAgentId(), filteredIds);
                     return user != null && !filteredIds.contains(user.getUserId());
                 }
             })