diff --git a/agent/src/main/java/dev/aikido/agent/Wrappers.java b/agent/src/main/java/dev/aikido/agent/Wrappers.java index a71c13b2..87b3e0bc 100644 --- a/agent/src/main/java/dev/aikido/agent/Wrappers.java +++ b/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(), diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/ProcessBuilderWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/ProcessBuilderWrapper.java new file mode 100644 index 00000000..7a9aa73f --- /dev/null +++ b/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 getMatcher() { + return ElementMatchers.isDeclaredBy(ProcessBuilder.class) + .and(named("start")); + } + + @Override + public ElementMatcher 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()); + } + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/RuntimeExecWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/RuntimeExecWrapper.java index 6032ae32..4fd12ee0 100644 --- a/agent/src/main/java/dev/aikido/agent/wrappers/RuntimeExecWrapper.java +++ b/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 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; } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/collectors/CommandCollector.java b/agent_api/src/main/java/dev/aikido/agent_api/collectors/CommandCollector.java index 719deb57..35cbdf06 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/collectors/CommandCollector.java +++ b/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 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", ""] + 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); } } diff --git a/agent_api/src/test/java/wrappers/ProcessBuilderTest.java b/agent_api/src/test/java/wrappers/ProcessBuilderTest.java new file mode 100644 index 00000000..2a325192 --- /dev/null +++ b/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()); + } +}