Incremental commit

phrocker · phrocker · commit f8e19ffa8d66 · 2025-05-10T08:17:22.000-04:00
diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java
@@ -86,7 +86,7 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
             var encryptedSecret = agentRegistrationDTO.getClientSecret();
             var decryptedSecret = agentKeyService.
                 decryptWithPrivateKey(encryptedSecret, keyPair.getPrivate());
-            keycloakService.createKeycloakClient(agentRegistrationDTO.getAgentName(),
+            keycloakService.createKeycloakClient(agentConfigOptions.getName(),
                 decryptedSecret);
 
 
@@ -131,7 +131,7 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
 
         while(running) {
 
-                log.info("Registering v1.0.2 agent failed. Retrying in 10 seconds...");
+                log.info("Agent Registered...");
                 try {
 
                     Thread.sleep(5_000);
diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/api/AgentController.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/api/AgentController.java
@@ -58,6 +58,7 @@ public ResponseEntity<AgentRegistrationDTO> describeAgent() {
             .agentName(agentConfig.getName())
             .agentPublicKey(agentKeyService.getKeyPair().getPublic().toString())
             .agentPublicKeyAlgo(agentKeyService.getKeyPair().getPublic().getAlgorithm())
+            .agentType(agentConfig.getType())
             .build();
         return ResponseEntity.ok(dto);
     }
diff --git a/ai-agent/src/main/java/io/sentrius/agent/config/AgentConfigOptions.java b/ai-agent/src/main/java/io/sentrius/agent/config/AgentConfigOptions.java
@@ -15,4 +15,5 @@ public class AgentConfigOptions {
 
 
     private String name;
+    private String type;
 }
diff --git a/ai-agent/src/main/resources/chat-helper.properties b/ai-agent/src/main/resources/chat-helper.properties
@@ -78,4 +78,5 @@ otel.metrics.exporter=none
 otel.logs.exporter=none
 otel.resource.attributes.service.name=chat-helper
 otel.traces.sampler=always_on
-otel.exporter.otlp.timeout=10s
+otel.exporter.otlp.timeout=10s
+agent.type="chat-helper"
diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java b/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java
@@ -1,11 +1,16 @@
 package io.sentrius.sso.controllers.api;
 
+import java.io.IOException;
+import java.io.InputStream;
 import java.security.GeneralSecurityException;
 import java.sql.SQLException;
 import io.sentrius.sso.config.ApiPaths;
 import io.sentrius.sso.core.config.SystemOptions;
 import io.sentrius.sso.core.controllers.BaseController;
 import io.sentrius.sso.core.dto.AgentRegistrationDTO;
+import io.sentrius.sso.core.model.security.IdentityType;
+import io.sentrius.sso.core.model.security.UserType;
+import io.sentrius.sso.core.model.users.User;
 import io.sentrius.sso.core.services.ATPLPolicyService;
 import io.sentrius.sso.core.services.ErrorOutputService;
 import io.sentrius.sso.core.services.UserService;
@@ -19,6 +24,7 @@
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PostMapping;
@@ -39,6 +45,13 @@ public class AgentBootstrapController extends BaseController {
     final ZeroTrustRequestService ztrService;
     final AgentService agentService;
 
+
+    @Value("${sentrius.agent.register.bootstrap.allow:false}")
+    private boolean allowRegistration;
+
+    @Value("${sentrius.agent.bootstrap.policy:default-policy.yaml}")
+    private String defaultPolicyFile;
+
     public AgentBootstrapController(
         UserService userService,
         SystemOptions systemOptions,
@@ -62,22 +75,56 @@ public AgentBootstrapController(
 
     @PostMapping("/register")
     // no LimitAccess
-    public ResponseEntity<AgentRegistrationDTO> bootsrap(
-        @RequestBody AgentRegistrationDTO registrationDTO) throws GeneralSecurityException {
+    public ResponseEntity<AgentRegistrationDTO> bootstrap(
+        @RequestBody AgentRegistrationDTO registrationDTO) throws GeneralSecurityException, IOException {
         log.info("Registering agent {}", registrationDTO);
-        var secret = keycloakService.registerAgentClient(registrationDTO.getAgentName());
+        // need a pre-shared secret to register the agent or ztat approval
+        var unencryptedRegistration = keycloakService.registerAgentClient(registrationDTO);
 
-        var secretKey = CryptoService.encryptWithPublicKey(secret,
+        var secretKey = CryptoService.encryptWithPublicKey(unencryptedRegistration.getClientSecret(),
             CryptoService.decodePublicKey(registrationDTO.getAgentPublicKey(),
             registrationDTO.getAgentPublicKeyAlgo()));
 
         var newDTO = AgentRegistrationDTO.builder()
-            .agentName(registrationDTO.getAgentName())
+            .agentName(unencryptedRegistration.getAgentName())
             .agentPublicKey(registrationDTO.getAgentPublicKey())
             .agentPublicKeyAlgo(registrationDTO.getAgentPublicKeyAlgo())
             .clientSecret(secretKey)
+            .clientId(unencryptedRegistration.getClientId())
             .build();
 
+        if (allowRegistration) {
+            log.info("Registering {}", registrationDTO.getAgentName());
+            User user = userService.getUserByUsername(newDTO.getAgentName());
+            if (user == null) {
+                var type = userService.getUserType(
+                    UserType.createUnknownUser());
+
+                user = User.builder()
+                    .username(newDTO.getAgentName())
+                    .name(newDTO.getAgentName())
+                    .emailAddress(newDTO.getAgentName())
+                    .userId(unencryptedRegistration.getClientId())
+                    .authorizationType(type.get())
+                    .identityType(IdentityType.NON_PERSON_ENTITY)
+                    .build();
+                log.info("Creating new user: {}", user);
+                userService.save(user);
+            }
+            try(InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream(defaultPolicyFile)) {
+                if (terminalHelperStream == null) {
+                    throw new RuntimeException(defaultPolicyFile + "not found on classpath");
+
+                }
+                String defaultYaml = new String(terminalHelperStream.readAllBytes());
+                log.info("Default policy file: {}", defaultPolicyFile);
+                var policy = atplPolicyService.createPolicy(user, defaultYaml);
+            }
+
+        }
+        else {
+            log.info("Not Registering {}", registrationDTO.getAgentName());
+        }
         // bootstrap with a default policy
         return ResponseEntity.ok(newDTO);
     }
diff --git a/api/src/main/resources/application.properties b/api/src/main/resources/application.properties
@@ -88,3 +88,6 @@ otel.logs.exporter=none
 otel.resource.attributes.service.name=sentrius-api
 otel.traces.sampler=always_on
 otel.exporter.otlp.timeout=10s
+
+sentrius.agent.register.bootstrap.allow=true
+sentrius.agent.bootstrap.policy=default-policy.yaml
diff --git a/api/src/main/resources/default-policy.yaml b/api/src/main/resources/default-policy.yaml
@@ -0,0 +1,53 @@
+---
+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:
+  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
diff --git a/api/src/test/resources/default-policy.yaml b/api/src/test/resources/default-policy.yaml
@@ -1,6 +1,6 @@
 ---
 version: "v0"
-description: "Default Policy For Chat Agents"
+description: "Default Policy For Unregistered Agents ( if configured )"
 match:
   agent_tags:
     - "env:prod"
@@ -19,18 +19,7 @@ actions:
     ztat_provider: "ztat-service.internal"
 capabilities:
   primitives:
-    - id: "accessLLM"
-      description: "access llm"
-      endpoint:
-        - "/api/v1/chat/completions"
-      tags: null
   composed:
-    - id: "observe"
-      description: null
-      endpoint: []
-      tags: null
-      includes:
-        - "readLogs"
 ztat:
   provider: "ztat-service.internal"
   ttl: "5m"
diff --git a/core/src/main/java/io/sentrius/sso/config/KeycloakManager.java b/core/src/main/java/io/sentrius/sso/config/KeycloakManager.java
@@ -6,13 +6,15 @@
 import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
+import lombok.extern.slf4j.Slf4j;
 import org.keycloak.OAuth2Constants;
 import org.keycloak.admin.client.Keycloak;
 import org.keycloak.admin.client.KeycloakBuilder;
 
 @Builder
 @Getter
 @Setter
+@Slf4j
 public class KeycloakManager {
     private Keycloak keycloak;
 
@@ -29,6 +31,8 @@ public void addPublicKey(String kid, PublicKey publicKey) {
 
 
     public Keycloak createKeycloakClient(String serverUrl, String realm, String clientId, String clientSecret) {
+        log.info("Creating Keycloak client for server: {}, realm: {}, clientId: {} shhh {}", serverUrl, realm,
+            clientId, clientSecret);
         return KeycloakBuilder.builder()
             .serverUrl(serverUrl)
             .realm(realm)
diff --git a/core/src/main/java/io/sentrius/sso/core/dto/AgentRegistrationDTO.java b/core/src/main/java/io/sentrius/sso/core/dto/AgentRegistrationDTO.java
@@ -17,4 +17,6 @@ public class AgentRegistrationDTO {
     private final String agentPublicKey;
     private final String agentPublicKeyAlgo;
     private final String clientSecret;
+    private final String clientId;
+    private final String agentType;
 }
diff --git a/core/src/main/java/io/sentrius/sso/core/services/security/KeycloakService.java b/core/src/main/java/io/sentrius/sso/core/services/security/KeycloakService.java
@@ -2,10 +2,12 @@
 import io.jsonwebtoken.Jwts;
 import io.sentrius.sso.config.KeycloakConfig;
 import io.sentrius.sso.config.KeycloakManager;
+import io.sentrius.sso.core.dto.AgentRegistrationDTO;
 import jakarta.ws.rs.core.Response;
 import lombok.extern.slf4j.Slf4j;
 import org.keycloak.admin.client.resource.ClientResource;
 import org.keycloak.admin.client.resource.ClientsResource;
+import org.keycloak.admin.client.resource.RoleMappingResource;
 import org.keycloak.admin.client.resource.UsersResource;
 import org.keycloak.representations.idm.ClientRepresentation;
 import org.keycloak.representations.idm.CredentialRepresentation;
@@ -102,12 +104,12 @@ public String extractUsername(String token) {
         return claims.get("preferred_username", String.class); // Extracts agent identity
     }
 
-    public String registerAgentClient(String agentName) {
+    public AgentRegistrationDTO registerAgentClient(AgentRegistrationDTO agent) {
         ClientsResource clients = keycloak.getKeycloak().realm(realm).clients();
 
         // Step 1: Build client representation
         ClientRepresentation client = new ClientRepresentation();
-        client.setClientId(agentName);
+        client.setClientId(agent.getAgentName());
         client.setEnabled(true);
         client.setProtocol("openid-connect");
         client.setServiceAccountsEnabled(true);
@@ -129,9 +131,34 @@ public String registerAgentClient(String agentName) {
             String clientId = response.getLocation().getPath().replaceAll(".*/([^/]+)$", "$1");
             ClientResource createdClient = clients.get(clientId);
 
+            String serviceAccountUserId = createdClient.getServiceAccountUser().getId();
+
+            // Step 5: Assign realm-management roles
+            RoleMappingResource roleMapping = keycloak.getKeycloak().realm(realm)
+                .users().get(serviceAccountUserId).roles();
+
+            var clientRoles = keycloak.getKeycloak().realm(realm)
+                .clients()
+                .findByClientId("realm-management")
+                .stream()
+                .findFirst()
+                .map(cr -> keycloak.getKeycloak().realm(realm).clients().get(cr.getId()).roles().list())
+                .orElseThrow(() -> new RuntimeException("realm-management client not found"));
+
+            var toAssign = clientRoles.stream()
+                .filter(role -> List.of("view-clients", "query-clients").contains(role.getName()))
+                .toList();
+
+            roleMapping.clientLevel(clientRoles.get(0).getContainerId()).add(toAssign);
+
             // Step 4: Get generated secret
             CredentialRepresentation secret = createdClient.getSecret();
-            return secret.getValue(); // Return this to the agent
+            return AgentRegistrationDTO.builder().
+                agentName(createdClient.getServiceAccountUser().getUsername())
+                .agentPublicKey(agent.getAgentPublicKey())
+                .clientSecret(secret.getValue())
+                .clientId(clientId)
+                .build();
         }
     }
 
diff --git a/dataplane/src/main/java/io/sentrius/sso/core/services/ATPLPolicyService.java b/dataplane/src/main/java/io/sentrius/sso/core/services/ATPLPolicyService.java
@@ -76,6 +76,7 @@ public Optional<ATPLPolicy> getPolicy(User operatingUser) {
         List<AgentPolicyAssignment> assignments = agentPolicyAssignmentRepository.findByUserUsernameOrderByAssignedAtDesc(operatingUser.getUsername());
 
         if (assignments.isEmpty()) {
+            log.info("No policy assignments found for user {}", operatingUser.getUsername());
             return Optional.empty();
         }
 

Original file line number	Diff line number	Diff line change
`@@ -58,6 +58,7 @@ public ResponseEntity<AgentRegistrationDTO> describeAgent() {`
`58`	`58`	`.agentName(agentConfig.getName())`
`59`	`59`	`.agentPublicKey(agentKeyService.getKeyPair().getPublic().toString())`
`60`	`60`	`.agentPublicKeyAlgo(agentKeyService.getKeyPair().getPublic().getAlgorithm())`
	`61`	`+ .agentType(agentConfig.getType())`
`61`	`62`	`.build();`
`62`	`63`	`return ResponseEntity.ok(dto);`
`63`	`64`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,4 +15,5 @@ public class AgentConfigOptions {`
`15`	`15`
`16`	`16`
`17`	`17`	`private String name;`
	`18`	`+ private String type;`
`18`	`19`	`}`
Original file line number	Diff line number	Diff line change
`@@ -17,4 +17,6 @@ public class AgentRegistrationDTO {`
`17`	`17`	`private final String agentPublicKey;`
`18`	`18`	`private final String agentPublicKeyAlgo;`
`19`	`19`	`private final String clientSecret;`
	`20`	`+ private final String clientId;`
	`21`	`+ private final String agentType;`
`20`	`22`	`}`