draft: Add protection to ProcessBuilder

agent/src/main/java/dev/aikido/agent/Wrappers.java

@@ -23,7 +23,11 @@ private Wrappers() {}
             new SpringControllerWrapper(),
             new FileConstructorSingleArgumentWrapper(),
             new FileConstructorMultiArgumentWrapper(),
+
+            // Shell wrappers
             new RuntimeExecWrapper(),
+            new ProcessBuilderWrapper(),
+
             new MysqlCJWrapper(),
             new MSSQLWrapper(),
             new MariaDBWrapper(),

agent/src/main/java/dev/aikido/agent/wrappers/ProcessBuilderWrapper.java

@@ -0,0 +1,74 @@
+package dev.aikido.agent.wrappers;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.is;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+public class ProcessBuilderWrapper implements Wrapper {
+    public String getName() {
+        // Wrap ProcessBuilder start().
+        // https://docs.oracle.com/en/java/javase/25/docs/api/java.base/java/lang/ProcessBuilder.html#start()
+        return ProcessBuilderAdvice.class.getName();
+    }
+    public ElementMatcher<? super MethodDescription> getMatcher() {
+        return ElementMatchers.isDeclaredBy(ProcessBuilder.class)
+                .and(named("start"));
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return is(ProcessBuilder.class);
+    }
+
+    public static class ProcessBuilderAdvice {
+        // Since we have to wrap a native Java Class stuff gets more complicated
+        // The classpath is not the same anymore, and we can't import our modules directly.
+        // To bypass this issue we load collectors from a .jar file.
+        @Advice.OnMethodEnter
+        public static void before(
+            @Advice.This(typing = DYNAMIC) ProcessBuilder target
+        ) throws Throwable {
+            String jarFilePath = System.getProperty("AIK_agent_api_jar");
+            URLClassLoader classLoader = null;
+            try {
+                URL[] urls = { new URL(jarFilePath) };
+                classLoader = new URLClassLoader(urls);
+            } catch (MalformedURLException ignored) {}
+            if (classLoader == null) {
+                return;
+            }
+
+            try {
+                // Load the class from the JAR
+                Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.CommandCollector");
+
+                // Run report with "argument"
+                for (Method method2: clazz.getMethods()) {
+                    if(method2.getName().equals("report")) {
+                        method2.invoke(null, target.command());
+                        break;
+                    }
+                }
+                classLoader.close(); // Close the class loader
+            } catch (InvocationTargetException invocationTargetException) {
+                if(invocationTargetException.getCause().toString().startsWith("dev.aikido.agent_api.vulnerabilities")) {
+                    throw invocationTargetException.getCause();
+                }
+                // Ignore non-aikido throwables.
+            } catch(Throwable e) {
+                System.out.println("AIKIDO: " + e.getMessage());
+            }
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/RuntimeExecWrapper.java

@@ -14,8 +14,7 @@
 import java.net.URLClassLoader;
 
 import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
-import static net.bytebuddy.matcher.ElementMatchers.is;
-import static net.bytebuddy.matcher.ElementMatchers.nameContains;
+import static net.bytebuddy.matcher.ElementMatchers.*;
 
 public class RuntimeExecWrapper implements Wrapper {
     public String getName() {
@@ -24,8 +23,11 @@ public String getName() {
         return CommandExecAdvice.class.getName();
     }
     public ElementMatcher<? super MethodDescription> getMatcher() {
+        // We only monkey-patch Runtime.exec(string), technically all Runtime.exec calls end up at the ProcessBuilder
+        // but at that point they are already split into arguments, so scanning a single command requires us to add
+        // a wrapper here, only for single strings.
         return ElementMatchers.isDeclaredBy(Runtime.class)
-                .and(ElementMatchers.nameContainsIgnoreCase("exec"));
+                .and(ElementMatchers.nameContainsIgnoreCase("exec")).and(takesArgument(0,  String.class));
     }
 
     @Override
@@ -39,13 +41,8 @@ public static class CommandExecAdvice {
         // To bypass this issue we load collectors from a .jar file.
         @Advice.OnMethodEnter
         public static void before(
-                @Advice.This(typing = DYNAMIC, optional = true) Object target,
-                @Advice.Origin Executable method,
-                @Advice.Argument(0) Object argument
+            @Advice.Argument(value = 0, typing = DYNAMIC) String command
         ) throws Throwable {
-            if (!(argument instanceof String)) {
-                return;
-            }
             String jarFilePath = System.getProperty("AIK_agent_api_jar");
             URLClassLoader classLoader = null;
             try {
@@ -63,7 +60,7 @@ public static void before(
                 // Run report with "argument"
                 for (Method method2: clazz.getMethods()) {
                     if(method2.getName().equals("report")) {
-                        method2.invoke(null, argument);
+                        method2.invoke(null, command);
                         break;
                     }
                 }

agent_api/src/main/java/dev/aikido/agent_api/collectors/CommandCollector.java

@@ -7,23 +7,34 @@
 import dev.aikido.agent_api.vulnerabilities.Scanner;
 import dev.aikido.agent_api.vulnerabilities.Vulnerabilities;
 
+import java.util.Arrays;
+import java.util.List;
+
 public final class CommandCollector {
     private CommandCollector() {}
     private static final Logger logger = LogManager.getLogger(CommandCollector.class);
-    public static void report(Object command) {
-        if (command instanceof String commandStr) {
-            if (commandStr.isEmpty()) {
-                return; // Empty command, don't scan.
-            }
+    public static void report(String command) {
+        if (command.isEmpty()) {
+            return; // Empty command, don't scan.
+        }
 
-            logger.trace("Scanning command: %s", commandStr);
+        logger.trace("Scanning command: %s", command);
 
-            // report stats
-            StatisticsStore.registerCall("runtime.Exec", OperationKind.EXEC_OP);
+        // report stats
+        StatisticsStore.registerCall("runtime.Exec", OperationKind.EXEC_OP);
 
-            // scan
-            Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.ShellInjectionVulnerability();
-            Scanner.scanForGivenVulnerability(vulnerability, "runtime.Exec", new String[]{commandStr});
-        }
+        // scan
+        Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.ShellInjectionVulnerability();
+        Scanner.scanForGivenVulnerability(vulnerability, "runtime.Exec", new String[]{command});
+    }
+
+    public static void report(List<String> commandArgs) {
+        // This happens when Runtime.exec()'s are being called with multiple arguments -> gets forwarded.
+        // or when new ProcessBuilder() is called. While we don't protect for argument injections, we do protect
+        // against cases like ["sh", "-c", "<command>"]
+        logger.trace("Scanning command arguments: %s", commandArgs);
+        StatisticsStore.registerCall("ProcessBuilder.start", OperationKind.EXEC_OP);
+        Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.ShellInjectionVulnerability();
+        //Scanner.scanForGivenVulnerability(vulnerability, "runtime.Exec", commandArgs);
     }
 }

agent_api/src/test/java/wrappers/ProcessBuilderTest.java

@@ -0,0 +1,185 @@
+package wrappers;
+
+import dev.aikido.agent_api.context.Context;
+import dev.aikido.agent_api.storage.ServiceConfigStore;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import utils.EmptySampleContextObject;
+
+import java.io.IOException;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class ProcessBuilderTest {
+    @AfterEach
+    void cleanup() {
+        Context.set(null);
+    }
+    @BeforeEach
+    void beforeEach() {
+        cleanup();
+        ServiceConfigStore.updateBlocking(true);
+    }
+    private void setContextAndLifecycle(String url) {
+        Context.set(new EmptySampleContextObject(url));
+    }
+
+    @Test
+    public void testShellInjection() {
+        setContextAndLifecycle(" -la");
+        Exception exception1 = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("yjytjyjty", "-c", "ls -la").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception1.getMessage());
+
+        cleanup();
+        setContextAndLifecycle("whoami");
+        Exception exception2 = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("bash", "-c", "whoami").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception2.getMessage());
+
+        cleanup();
+        assertDoesNotThrow(() -> {
+            Runtime.getRuntime().exec("whoami && ls -la");
+        });
+        assertThrows(IllegalArgumentException.class, () -> {
+            Runtime.getRuntime().exec("");
+        });
+    }
+
+    @Test
+    public void testOnlyScansStrings() {
+        setContextAndLifecycle("whoami");
+        assertDoesNotThrow(() -> {
+            Runtime.getRuntime().exec(new String[]{"whoami"});
+        });
+        assertDoesNotThrow(() -> {
+            Runtime.getRuntime().exec(new String[]{"whoami"}, new String[]{"MyEnvironmentVar=1"});
+        });
+
+        Exception exception1 = assertThrows(RuntimeException.class, () -> {
+            Runtime.getRuntime().exec("whoami", new String[]{"MyEnvironmentVar=1"});
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception1.getMessage());
+    }
+
+    // --- NEW TEST CASES ---
+
+    @Test
+    public void testProcessBuilderCommandModification() {
+        setContextAndLifecycle("whoami");
+        ProcessBuilder builder = new ProcessBuilder();
+        assertDoesNotThrow(() -> {
+            builder.command("whoami");
+            builder.start();
+        });
+
+        Exception exception = assertThrows(RuntimeException.class, () -> {
+            builder.command("sh", "-c", "whoami");
+            builder.start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception.getMessage());
+    }
+
+    @Test
+    public void testProcessBuilderWithDifferentShells() {
+        setContextAndLifecycle("whoami");
+        Exception shException = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("sh", "-c", "whoami").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", shException.getMessage());
+
+        Exception bashException = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("bash", "-c", "whoami").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", bashException.getMessage());
+
+        Exception zshException = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("zsh", "-c", "whoami").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", zshException.getMessage());
+    }
+
+    @Test
+    public void testProcessBuilderWithDirectCommand() {
+        setContextAndLifecycle("whoami");
+        assertDoesNotThrow(() -> {
+            new ProcessBuilder("whoami").start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithArguments() {
+        setContextAndLifecycle("whoami");
+        assertDoesNotThrow(() -> {
+            new ProcessBuilder("ls", "-l", "/tmp").start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithEnvironment() {
+        setContextAndLifecycle("whoami");
+        ProcessBuilder builder = new ProcessBuilder("whoami");
+        builder.environment().put("MY_VAR", "1");
+        assertDoesNotThrow(() -> {
+            builder.start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithShellInjectionInCommand() {
+        setContextAndLifecycle("whoami; ls");
+        Exception exception = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("sh", "-c", "whoami; ls").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception.getMessage());
+    }
+
+    @Test
+    public void testProcessBuilderWithComplexShellCommand() {
+        setContextAndLifecycle("whoami && ls -la");
+        Exception exception = assertThrows(RuntimeException.class, () -> {
+            new ProcessBuilder("bash", "-c", "whoami && ls -la").start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception.getMessage());
+    }
+
+    @Test
+    public void testProcessBuilderWithSafeCommand() {
+        setContextAndLifecycle("whoami");
+        assertDoesNotThrow(() -> {
+            new ProcessBuilder("whoami").start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithEmptyCommand() {
+        assertThrows(IndexOutOfBoundsException.class, () -> {
+            new ProcessBuilder().start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithNullCommand() {
+        assertThrows(NullPointerException.class, () -> {
+            new ProcessBuilder((String[]) null).start();
+        });
+    }
+
+    @Test
+    public void testProcessBuilderWithCommandModificationAfterStart() {
+        setContextAndLifecycle("whoami");
+        ProcessBuilder builder = new ProcessBuilder("whoami");
+        assertDoesNotThrow(() -> {
+            builder.start();
+        });
+        // Modifying command after start should not affect previous process
+        builder.command("sh", "-c", "whoami");
+        Exception exception = assertThrows(RuntimeException.class, () -> {
+            builder.start();
+        });
+        assertEquals("Aikido Zen has blocked Shell Injection", exception.getMessage());
+    }
+}