<commit_before>package me.coley.recaf.agent;

import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.objectweb.asm.tree.ClassNode;

import me.coley.recaf.LaunchParams;
import me.coley.recaf.Recaf;
import me.coley.recaf.asm.Asm;
import me.coley.recaf.util.Streams;

/**
 * Agent entry-point for Recaf.
 * 
 * <pre>
java -javaagent:/path/to/agent.jar -cp Test.jar TestMain
 * </pre>
 * 
 * @author Matt
 */
public class Agent {
	/**
	 * Map of loaded classes.
	 */
	public static final Map<String, byte[]> classes = new HashMap<>();
	/**
	 * Instrumentation instance.
	 */
	public static Instrumentation instrument;
	/**
	 * Exclude core java classes from being loaded.
	 */
	private static boolean excludeInternals = true;

	/**
	 * Called when attached through javaagent jvm arg.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	public static void premain(String agentArgs, Instrumentation inst) {
		agent(agentArgs, inst);
	}

	/**
	 * Called when attached to a jvm externally.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	public static void agentmain(String agentArgs, Instrumentation inst) {
		agent(agentArgs, inst);
	}

	/**
	 * Populate {@link #instrument} and invoke Recaf.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	private static void agent(String agentArgs, Instrumentation inst) {
		instrument = inst;
		if (agentArgs == null) {
			agentArgs = "";
		} else {
			// TODO: Lets say for some reason the user wishes to edit some
			// core-classes. They would want to set this boolean to false so
			// core classes are loaded... But if this boolean is set to false,
			// some wacky things happen.
			// Screenshot of issue: https://my.mixtape.moe/tptxwl.jpg
			// I have no idea why the instrumentation API is being such a child,
			// but I don't know why this is happening.
			if (agentArgs.equals("keep")) {
				System.out.println(Recaf.class);
				excludeInternals = false;
			}
			if (agentArgs.equals("keep2")) {
				excludeInternals = false;
			}
		}
		try {
			// Add transformer and invoke retransform so already loaded classes
			// are transformed.
			ClassPopulator transformer = new ClassPopulator();
			instrument.addTransformer(transformer, true);
			instrument.retransformClasses(getModifable());
		} catch (Throwable e) {
			e.printStackTrace();
		}
		// Manually set params agent value.
		// I keep getting mis-matched args via picoli when attaching to other
		// processes.
		LaunchParams params = new LaunchParams();
		params.isAgent = true;
		Recaf.start(new String[0], params);
	}

	/**
	 * Creates an array of modifable loaded classes.
	 * 
	 * @return
	 */
	@SuppressWarnings("rawtypes")
	private static Class[] getModifable() {
		Set<Class> set = new HashSet<>();
		for (Class clazz : instrument.getAllLoadedClasses()) {
			if (clazz.isArray() || clazz.isPrimitive() || clazz.isSynthetic()) {
				continue;
			}
			if (valid(clazz.getName()) && instrument.isModifiableClass(clazz)) {
				set.add(clazz);
			}
		}
		return set.toArray(new Class[set.size()]);
	}

	/**
	 * Update edits made to classes in the current VM.
	 */
	public static void apply() {
		try {
			Recaf.INSTANCE.logging.info("Applying class redefinintions:");
			ClassDefinition[] definitions = Marker.getDefinitions();
			for (ClassDefinition define : definitions) {
				Recaf.INSTANCE.logging.info(define.getDefinitionClass().getName(), 1);
			}
			instrument.redefineClasses(definitions);
			Recaf.INSTANCE.logging.info("Done applying redefinitions.");

		} catch (Exception e) {
			Recaf.INSTANCE.logging.error(e);
		}
	}

	/**
	 * Retrieve the map of ClassNodes using {@link #instrument}.
	 * 
	 * @return Map of JVM classes.
	 * @throws IOException
	 *             Thrown if classes could not be read.
	 */
	@SuppressWarnings("rawtypes")
	public static Map<String, ClassNode> getNodesViaInst() {
		Map<String, ClassNode> map = new HashMap<>();
		for (Class clazz : instrument.getAllLoadedClasses()) {
			try {
				String name = clazz.getName().replace(".", "/");
				if (!valid(name)) {
					continue;
				}
				ClassNode node = getNode(clazz);
				if (node != null) {
					map.put(node.name, node);
					Recaf.INSTANCE.logging.fine("Loaded runtime class: " + node.name, 1);
				}
			} catch (Exception e) {
				Recaf.INSTANCE.logging.fine("Failed loading runtime class: " + clazz.getName().replace(".", "/"));
				Recaf.INSTANCE.logging.error(e, false);
			}
		}
		return map;
	}

	/**
	 * Get node via its internal class name.
	 * 
	 * @param name
	 * @return
	 * @throws IOException
	 */
	private static ClassNode getNode(String name) throws IOException {
		byte[] bytes = classes.get(name);
		if (bytes == null) {
			return null;
		}
		return Asm.getNode(bytes);
	}

	/**
	 * Get node via its class.
	 * 
	 * @param clazz
	 * @return
	 * @throws IOException
	 */
	public static ClassNode getNode(Class<?> clazz) throws IOException {
		String name = clazz.getName().replace(".", "/");
		byte[] bytes = classes.get(name);
		if (bytes == null) {
			InputStream is = clazz.getResourceAsStream("/" + clazz.getName().replace('.', '/') + ".class");
			if (is == null) {
				return null;
			}
			bytes = Streams.from(is);
		}
		return Asm.getNode(bytes);
	}

	/**
	 * Add class to map and current class tree if recaf's UI is already open.
	 * 
	 * @param name
	 * @param bytes
	 */
	public static void addClass(String name, byte[] bytes) {
		classes.put(name, bytes);
		Recaf rc = Recaf.INSTANCE;
		try {
			ClassNode cn = getNode(name.replace("/", ".").replace("$", "."));
			if (cn == null) {
				return;
			}
			if (rc.jarData == null) {
				// Edge case, happens if class loaded is occuring while recaf is
				// setting itself up
				rc.logging.info("Load while jar pop: " + name);
			} else {
				rc.jarData.classes.put(name, cn);
			}

		} catch (Exception e) {
			e.printStackTrace();
		}
		if (rc.ui != null && valid(name) && rc.configs.agent.autoRefresh) {
			rc.ui.refreshTree();
		}<commit_msg>Remove debug message<commit_after>package me.coley.recaf.agent;

import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.objectweb.asm.tree.ClassNode;

import me.coley.recaf.LaunchParams;
import me.coley.recaf.Recaf;
import me.coley.recaf.asm.Asm;
import me.coley.recaf.util.Streams;

/**
 * Agent entry-point for Recaf.
 * 
 * <pre>
java -javaagent:/path/to/agent.jar -cp Test.jar TestMain
 * </pre>
 * 
 * @author Matt
 */
public class Agent {
	/**
	 * Map of loaded classes.
	 */
	public static final Map<String, byte[]> classes = new HashMap<>();
	/**
	 * Instrumentation instance.
	 */
	public static Instrumentation instrument;
	/**
	 * Exclude core java classes from being loaded.
	 */
	private static boolean excludeInternals = true;

	/**
	 * Called when attached through javaagent jvm arg.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	public static void premain(String agentArgs, Instrumentation inst) {
		agent(agentArgs, inst);
	}

	/**
	 * Called when attached to a jvm externally.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	public static void agentmain(String agentArgs, Instrumentation inst) {
		agent(agentArgs, inst);
	}

	/**
	 * Populate {@link #instrument} and invoke Recaf.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	private static void agent(String agentArgs, Instrumentation inst) {
		instrument = inst;
		if (agentArgs == null) {
			agentArgs = "";
		} else {
			// TODO: Lets say for some reason the user wishes to edit some
			// core-classes. They would want to set this boolean to false so
			// core classes are loaded... But if this boolean is set to false,
			// some wacky things happen.
			// Screenshot of issue: https://my.mixtape.moe/tptxwl.jpg
			// I have no idea why the instrumentation API is being such a child,
			// but I don't know why this is happening.
			if (agentArgs.equals("keep")) {
				System.out.println(Recaf.class);
				excludeInternals = false;
			}
			if (agentArgs.equals("keep2")) {
				excludeInternals = false;
			}
		}
		try {
			// Add transformer and invoke retransform so already loaded classes
			// are transformed.
			ClassPopulator transformer = new ClassPopulator();
			instrument.addTransformer(transformer, true);
			instrument.retransformClasses(getModifable());
		} catch (Throwable e) {
			e.printStackTrace();
		}
		// Manually set params agent value.
		// I keep getting mis-matched args via picoli when attaching to other
		// processes.
		LaunchParams params = new LaunchParams();
		params.isAgent = true;
		Recaf.start(new String[0], params);
	}

	/**
	 * Creates an array of modifable loaded classes.
	 * 
	 * @return
	 */
	@SuppressWarnings("rawtypes")
	private static Class[] getModifable() {
		Set<Class> set = new HashSet<>();
		for (Class clazz : instrument.getAllLoadedClasses()) {
			if (clazz.isArray() || clazz.isPrimitive() || clazz.isSynthetic()) {
				continue;
			}
			if (valid(clazz.getName()) && instrument.isModifiableClass(clazz)) {
				set.add(clazz);
			}
		}
		return set.toArray(new Class[set.size()]);
	}

	/**
	 * Update edits made to classes in the current VM.
	 */
	public static void apply() {
		try {
			Recaf.INSTANCE.logging.info("Applying class redefinintions:");
			ClassDefinition[] definitions = Marker.getDefinitions();
			for (ClassDefinition define : definitions) {
				Recaf.INSTANCE.logging.info(define.getDefinitionClass().getName(), 1);
			}
			instrument.redefineClasses(definitions);
			Recaf.INSTANCE.logging.info("Done applying redefinitions.");

		} catch (Exception e) {
			Recaf.INSTANCE.logging.error(e);
		}
	}

	/**
	 * Retrieve the map of ClassNodes using {@link #instrument}.
	 * 
	 * @return Map of JVM classes.
	 * @throws IOException
	 *             Thrown if classes could not be read.
	 */
	@SuppressWarnings("rawtypes")
	public static Map<String, ClassNode> getNodesViaInst() {
		Map<String, ClassNode> map = new HashMap<>();
		for (Class clazz : instrument.getAllLoadedClasses()) {
			try {
				String name = clazz.getName().replace(".", "/");
				if (!valid(name)) {
					continue;
				}
				ClassNode node = getNode(clazz);
				if (node != null) {
					map.put(node.name, node);
					Recaf.INSTANCE.logging.fine("Loaded runtime class: " + node.name, 1);
				}
			} catch (Exception e) {
				Recaf.INSTANCE.logging.fine("Failed loading runtime class: " + clazz.getName().replace(".", "/"));
				Recaf.INSTANCE.logging.error(e, false);
			}
		}
		return map;
	}

	/**
	 * Get node via its internal class name.
	 * 
	 * @param name
	 * @return
	 * @throws IOException
	 */
	private static ClassNode getNode(String name) throws IOException {
		byte[] bytes = classes.get(name);
		if (bytes == null) {
			return null;
		}
		return Asm.getNode(bytes);
	}

	/**
	 * Get node via its class.
	 * 
	 * @param clazz
	 * @return
	 * @throws IOException
	 */
	public static ClassNode getNode(Class<?> clazz) throws IOException {
		String name = clazz.getName().replace(".", "/");
		byte[] bytes = classes.get(name);
		if (bytes == null) {
			InputStream is = clazz.getResourceAsStream("/" + clazz.getName().replace('.', '/') + ".class");
			if (is == null) {
				return null;
			}
			bytes = Streams.from(is);
		}
		return Asm.getNode(bytes);
	}

	/**
	 * Add class to map and current class tree if recaf's UI is already open.
	 * 
	 * @param name
	 * @param bytes
	 */
	public static void addClass(String name, byte[] bytes) {
		classes.put(name, bytes);
		Recaf rc = Recaf.INSTANCE;
		try {
			ClassNode cn = getNode(name.replace("/", ".").replace("$", "."));
			if (cn == null) {
				return;
			}
			// Its fine if it's not put on immediately. It'll make it's way to the gui.
			if (rc.jarData != null) {
				rc.jarData.classes.put(name, cn);
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		if (rc.ui != null && valid(name) && rc.configs.agent.autoRefresh) {
			rc.ui.refreshTree();
		}
{'fraction_non_alphanumeric': 0.13808701373468213, 'fraction_numerical': 0.0006579488444773419, 'mean_word_length': 3.1902136457615438, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3086802', 'n_tokens_mistral': 4377, 'n_tokens_neox': 3820, 'n_words': 1075}