Fixup issue with regexes

Author: phrocker

api/src/main/resources/db/migration/V14__command_categorizer_gin.sql

@@ -0,0 +1,3 @@
+CREATE EXTENSION IF NOT EXISTS pg_trgm;
+
+CREATE INDEX idx_command_pattern_trgm ON command_categories USING gin (pattern gin_trgm_ops);

core/src/main/java/io/sentrius/sso/core/repository/CommandCategoryRepository.java

@@ -4,6 +4,7 @@
 import io.sentrius.sso.core.model.categorization.CommandCategory;
 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.data.jpa.repository.Query;
+import org.springframework.data.repository.query.Param;
 import org.springframework.stereotype.Repository;
 
 @Repository
@@ -12,4 +13,8 @@ public interface CommandCategoryRepository extends JpaRepository<CommandCategory
     List<CommandCategory> findAllOrderedByPriority();
 
     List<CommandCategory> findByPattern(String pattern);
+
+    @Query(value = "SELECT * FROM command_categories WHERE :command ~ pattern", nativeQuery = true)
+    List<CommandCategory> findMatchingCategories(@Param("command") String command);
+
 }

core/src/main/java/io/sentrius/sso/core/services/openai/categorization/CommandCategorizer.java

@@ -1,8 +1,10 @@
 package io.sentrius.sso.core.services.openai.categorization;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
@@ -33,42 +35,44 @@ public class CommandCategorizer {
     private final CommandCategoryRepository commandCategoryRepository;
 
 
-    private final CommandTrie commandTrie = new CommandTrie();
-
     private final Cache<String, CommandCategory> commandCache = Caffeine.newBuilder()
-        .maximumSize(10000)
-        .expireAfterWrite(24, TimeUnit.HOURS)
+        .maximumSize(1000)
+        .expireAfterWrite(1, TimeUnit.HOURS)
         .build();
 
-    @PostConstruct
-    public void initializeTrie() {
-        List<CommandCategory> categories = commandCategoryRepository.findAll();
-        for (CommandCategory category : categories) {
-            log.info("Adding command category {} to trie", category);
-            commandTrie.insert(category.getPattern(), category);
-        }
+
+
+    private CommandCategory fetchFromDatabase(String command) {
+        List<CommandCategory> matchingCategories = commandCategoryRepository.findMatchingCategories(command);
+        return matchingCategories.stream()
+            .min(Comparator.comparingInt(CommandCategory::getPriority))
+            .orElse(null);
     }
 
+
     @Transactional
     public CommandCategory categorizeCommand(String command) {
         return commandCache.get(command, this::categorizeWithRulesOrML);
     }
 
-
     protected List<CommandCategory> getDBCommandCategory(String command){
         return commandCategoryRepository.findByPattern(command);
     }
 
-    @Transactional
-    protected void addCommandCategory(String command, CommandCategory category) {
-        commandTrie.insert(command, category);
-        commandCategoryRepository.save(category);
+
+    public boolean isValidRegex(String regex) {
+        try {
+            Pattern.compile(regex);
+            return true; // Valid regex
+        } catch (PatternSyntaxException e) {
+            return false; // Invalid regex
+        }
     }
 
 
     @Transactional
     protected CommandCategory categorizeWithRulesOrML(String command) {
-        CommandCategory category = commandTrie.searchByPrefix(command);
+        CommandCategory category = fetchFromDatabase(command);
         if (category != null) {
             log.info("Found command category {} for {} ", category, command);
             return category;
@@ -93,8 +97,9 @@ protected CommandCategory categorizeWithRulesOrML(String command) {
             try {
                 category = commandCategorizer.generate(command);
 
-                addCommandCategory(category.getPattern(), category);
-                
+                if (isValidRegex(category.getPattern())) {
+                    addCommandCategory(category.getPattern(), category);
+                }
                 log.info("Categorized command: {}", category);
                 return category;
             } catch (Exception e) {
@@ -110,4 +115,9 @@ protected CommandCategory categorizeWithRulesOrML(String command) {
 
         return CommandCategory.builder().build();
     }
+
+    private void addCommandCategory(String pattern, CommandCategory category) {
+        commandCategoryRepository.save(category);
+        commandCache.put(pattern, category);
+    }
 }

core/src/main/java/io/sentrius/sso/genai/LLMCommandCategorizer.java

@@ -58,28 +58,28 @@ public CommandCategory generate(String on) throws HttpException, JsonProcessingE
     public String generateInput(String on) {
         return """
         Categorize the following command with a generalized pattern, defined as a **regex**, that captures the intent and considers risk factors. Include specific arguments or paths in the regex only if they significantly impact the risk level. If no regex is predefined for the command, generate one that appropriately generalizes the command's behavior while retaining any risk-relevant components.
-                
-                For example:
-                - 'cat /etc/passwd' is risky due to sensitive user data, so it should be included as-is with the regex '^cat /etc/passwd$'.
-                - 'cat /etc/hosts' has low risk, so it should be generalized to '^cat /etc/.*$' to cover all files in the `/etc` directory.
-                - 'sudo rm -rf /important_dir' should include '/important_dir' if it's sensitive, but otherwise generalized to '^sudo rm -rf .*'.
-                
-                Command: "%s"
-                
-                **Categories to choose from:**
-                - PRIVILEGED: Commands that require elevated permissions or pose a risk if misused.
-                - DESTRUCTIVE: Commands that can delete or alter critical files or system configurations.
-                - INFORMATIONAL: Commands that retrieve information without altering the system state.
-                - GENERAL: Commands that do not fit into the above categories.
-                
-                Respond in **JSON format** as follows:
-                {
-                    "category": "<Category Name>",
-                    "priority": <Numerical Priority>,
-                    "pattern": "<Generalized regex Pattern>",
-                    "rationale": "<Explain why this category and regex were chosen>"
-                }
-                
+         
+         For example:
+         - 'cat /etc/passwd' is risky due to sensitive user data, so it should be included as-is with the regex '^cat /etc/passwd$'.
+         - 'cat /etc/hosts' has low risk, so it should be generalized to '^cat /etc/.*$' to cover all files in the `/etc` directory.
+         - 'sudo rm -rf /important_dir' should include '/important_dir' if it's sensitive, but otherwise generalized to '^sudo rm -rf .*'.
+         
+         Command: "%s"
+         
+         **Categories to choose from:**
+         - PRIVILEGED: Commands that require elevated permissions or pose a risk if misused.
+         - DESTRUCTIVE: Commands that can delete or alter critical files or system configurations.
+         - INFORMATIONAL: Commands that retrieve information without altering the system state.
+         - GENERAL: Commands that do not fit into the above categories.
+         
+         Respond in **JSON format** as follows:
+         {
+             "category": "<Category Name>",
+             "priority": <Numerical Priority>,
+             "pattern": "<Generalized regex Pattern>",
+             "rationale": "<Explain why this category and regex were chosen>"
+         }
+            
     """.formatted(on);
     }