Update wrappers 2/2

Author: bitterpanda63

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

@@ -6,11 +6,19 @@
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
 import java.io.File;
+import java.io.IOException;
 import java.lang.instrument.Instrumentation;
+import java.util.jar.JarFile;
 
 import static dev.aikido.agent.ByteBuddyInitializer.createAgentBuilder;
 import static dev.aikido.agent.DaemonStarter.startDaemon;
@@ -19,17 +27,31 @@
 
 public class Agent {
     private static final Logger logger = LogManager.getLogger(Agent.class);
+
     public static void premain(String agentArgs, Instrumentation inst) {
         // Check for 'AIKIDO_DISABLE' :
         if (new BooleanEnv("AIKIDO_DISABLE", /*default value*/ false).getValue()) {
             return; // AIKIDO_DISABLE is true, so we will not be wrapping anything.
         }
         logger.info("Zen by Aikido v%s starting.", Config.pkgVersion);
-        setAikidoSysProperties();
+        try {
+            setAikidoSysProperties(inst);
+        } catch (Exception e) {
+            logger.error(e.getMessage());
+        }
+
+        // Modify bootstrap class path (includes the core Java classes), so we can have some shared
+        // code for core java classes when wrapping.
+        try {
+            inst.appendToBootstrapClassLoaderSearch(new JarFile(getPathToAikidoDirectory() + "/agent_bootstrap.jar"));
+        } catch (IOException e) {
+            logger.error(e.getMessage());
+        }
 
         // Test loading of zen binaries :
         loadLibrary();
 
+
         ElementMatcher.Junction wrapperTypeDescriptors = ElementMatchers.none();
         for(Wrapper wrapper: WRAPPERS) {
             wrapperTypeDescriptors = wrapperTypeDescriptors.or(wrapper.getTypeMatcher());
@@ -56,11 +78,14 @@ public static AgentBuilder.Transformer get() {
             return adviceAgentBuilder;
         }
     }
-    private static void setAikidoSysProperties() {
+    private static String getPathToAikidoDirectory() {
         String pathToAgentJar = Agent.class.getProtectionDomain().getCodeSource().getLocation().getPath();
-        String pathToAikidoDirectory = new File(pathToAgentJar).getParent();
-        String jarPath = "file:" + pathToAikidoDirectory + "/agent_api.jar";
-        System.setProperty("AIK_agent_dir", pathToAikidoDirectory);
+        return new File(pathToAgentJar).getParent();
+    }
+    private static void setAikidoSysProperties(Instrumentation inst) throws IOException {
+
+        String jarPath = "file:" + getPathToAikidoDirectory() + "/agent_api.jar";
+        System.setProperty("AIK_agent_dir", getPathToAikidoDirectory());
         System.setProperty("AIK_agent_api_jar", jarPath);
     }
-}
+}

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

@@ -1,21 +1,12 @@
 package dev.aikido.agent.wrappers;
+import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
 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.io.IOException;
-import java.lang.reflect.Executable;
-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.nameContains;
 
 public class RuntimeExecWrapper implements Wrapper {
     public String getName() {
@@ -34,48 +25,11 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
     }
 
     public static class CommandExecAdvice {
-        // 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, optional = true) Object target,
-                @Advice.Origin Executable method,
                 @Advice.Argument(0) Object argument
         ) throws Throwable {
-            if (!(argument instanceof String)) {
-                return;
-            }
-            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, argument);
-                        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());
-            }
+            AikidoBootstrapClass.invoke(AikidoBootstrapClass.COMMAND_COLLECTOR_REPORT, argument);
         }
     }
 }