Recaf 4X decompile script for command line use

decompile.java

/*
An example script showcasing how batch decompilation can be enriched with
multiple Recaf features such as:
- Inserting comments into classes & members
- Pre-processing classes to clean up some obfuscation

Usage: Add arguments for this task after the ';' character - see below for how this is handled and arg options
  java -jar recaf.jar -i <jar-to-decompile> -s decompile.java -h -q ; <args-for-this-script>
  java -jar recaf.jar -i <jar-to-decompile> -s decompile.java -h -q
*/
import jakarta.annotation.Nonnull;
import jakarta.annotation.Nullable;
import jakarta.enterprise.context.Dependent;
import jakarta.inject.Inject;
import picocli.CommandLine;
import picocli.CommandLine.Command;
import picocli.CommandLine.Option;
import software.coley.recaf.info.ClassInfo;
import software.coley.recaf.info.JvmClassInfo;
import software.coley.recaf.info.member.MethodMember;
import software.coley.recaf.launch.LaunchArguments;
import software.coley.recaf.launch.LaunchCommand;
import software.coley.recaf.path.ClassPathNode;
import software.coley.recaf.services.comment.ClassComments;
import software.coley.recaf.services.comment.CommentManager;
import software.coley.recaf.services.comment.WorkspaceComments;
import software.coley.recaf.services.decompile.DecompileResult;
import software.coley.recaf.services.decompile.DecompilerManager;
import software.coley.recaf.services.decompile.DecompilerManagerConfig;
import software.coley.recaf.services.decompile.JvmDecompiler;
import software.coley.recaf.services.mapping.MappingApplierService;
import software.coley.recaf.services.search.SearchService;
import software.coley.recaf.services.transform.TransformationApplier;
import software.coley.recaf.services.workspace.WorkspaceManager;
import software.coley.recaf.services.workspace.patch.PatchApplier;
import software.coley.recaf.util.ZipCreationUtils;
import software.coley.recaf.workspace.model.Workspace;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

@Dependent
public class DecompileTask implements Runnable {
	/** Used as a delimiter between base {@link LaunchCommand Recaf args} and {@link Args args for this task}. */
	private static final String argSplit = ";";
	// Injected services
	private final DecompilerManager decompileManager;
	private final DecompilerManagerConfig decompilerManagerConfig;
	private final CommentManager commentManager;
	private final WorkspaceManager workspaceManager;
	private final LaunchArguments arguments;

	@Inject
	public DecompileTask(@Nonnull DecompilerManager decompilerManager,
	                     @Nonnull DecompilerManagerConfig decompilerManagerConfig,
	                     @Nonnull CommentManager commentManager,
	                     @Nonnull WorkspaceManager workspaceManager,
	                     @Nonnull LaunchArguments arguments) {
		this.decompileManager = decompilerManager;
		this.decompilerManagerConfig = decompilerManagerConfig;
		this.commentManager = commentManager;
		this.workspaceManager = workspaceManager;
		this.arguments = arguments;
	}

	@Override
	public void run() {
		// Input validation
		if (!workspaceManager.hasCurrentWorkspace()) {
			System.err.println("No workspace provided");
			return;
		}

		// Parse args for this task
		String[] rawArgs = arguments.getArgs();
		int argSplitIndex = Arrays.asList(rawArgs).indexOf(argSplit);
		Args args = new Args();
		if (argSplitIndex > 0) {
			try {
				String[] subArgs = Arrays.copyOfRange(rawArgs, argSplitIndex + 1, rawArgs.length);
				CommandLine.populateCommand(args, subArgs);
			} catch (Throwable t) {
				System.err.println("Error parsing inputs: " + t);
				return;
			}
		}

		// Configure decompiler options
		configureDecompiler(args);

		// Decompile all classes
		ZipCreationUtils.ZipBuilder builder = new ZipCreationUtils.ZipBuilder();
		StringBuilder failures = new StringBuilder();
		Workspace workspace = workspaceManager.getCurrent();
		workspace.classesStream(false).forEach(cp -> {
			// Some pre-processing
			annotateClass(cp);

			// Decompile the class
			String name = cp.getValue().getName();
			String decompiled = decompile(workspace, cp);

			// Append the decompiled class to the output
			if (decompiled != null) {
				builder.add(name + ".java", decompiled.getBytes(StandardCharsets.UTF_8));
			} else {
				failures.append(name).append('\n');
			}
		});
		if (!failures.isEmpty())
			builder.add("failures.txt", failures.toString().getBytes(StandardCharsets.UTF_8));

		// Write to output
		try {
			Files.write(Paths.get("decompiled.zip"), builder.bytes());
		} catch (IOException e) {
			System.err.println("Failed to write decompiled output");
		}
	}

	/**
	 * Example of some pre-processing.
	 * <br>
	 * In this simple case, we just add some comments to different parts of the class.
	 * But you could also do more such as:
	 * <ul>
	 *     <li>Apply patches via {@link PatchApplier}</li>
	 *     <li>Deobfuscation via {@link TransformationApplier}</li>
	 *     <li>Mapping via {@link MappingApplierService}</li>
	 *     <li>Commenting on members with found patterns via {@link SearchService}</li>
	 * </ul>
	 */
	private void annotateClass(@Nonnull ClassPathNode classPath) {
		ClassInfo classInfo = classPath.getValue();
		String className = classInfo.getName();

		WorkspaceComments comments = Objects.requireNonNull(commentManager.getCurrentWorkspaceComments());
		ClassComments classComments = comments.getOrCreateClassComments(classPath);
		if (className.startsWith("software/coley/recaf")) {
			classComments.setClassComment("This is a class that belongs to Recaf");
		}

		for (MethodMember method : classInfo.getMethods()) {
			if (method.getName().equals("main")) {
				classComments.setMethodComment(method, "This is the main method");
			}
		}
	}

	/**
	 * Configure pre-processing at the decompilation step.
	 * These do not affect the content in the workspace and are applied on a per-decompile basis.
	 */
	private void configureDecompiler(@Nonnull Args args) {
		if (args.filterSynthetics) decompilerManagerConfig.getFilterSynthetics().setValue(true);
		if (args.filterGenerics) decompilerManagerConfig.getFilterSignatures().setValue(true);
		if (args.filterLong > 0) {
			decompilerManagerConfig.getFilterLongAnnotations().setValue(true);
			decompilerManagerConfig.getFilterLongAnnotationsLength().setValue(args.filterLong);
			decompilerManagerConfig.getFilterLongExceptions().setValue(true);
			decompilerManagerConfig.getFilterLongExceptionsLength().setValue(args.filterLong);
		}
		if (args.decompiler != null) decompilerManagerConfig.getPreferredJvmDecompiler().setValue(args.decompiler);
	}

	/**
	 * Decompiles the specified class.
	 */
	@Nullable
	private String decompile(@Nonnull Workspace workspace, @Nonnull ClassPathNode classPath) {
		JvmDecompiler decompiler = decompileManager.getTargetJvmDecompiler();
		try {
			JvmClassInfo classInfo = classPath.getValue().asJvmClass();
			DecompileResult result = decompileManager.decompile(decompiler, workspace, classInfo).get(5, TimeUnit.MINUTES);
			if (result.getException() != null) {
				System.err.println("Error occurred when decompiling: " + result.getException());
				return null;
			}

			return result.getText();
		} catch (TimeoutException e) {
			System.err.println("Decompilation took too long");
		} catch (Throwable e) {
			System.err.println("Unhandled error occurred when decompiling: " + e);
		}
		return null;
	}

	@Command(name = "", version = "", description = "")
	public static class Args implements Callable<Void> {
		@Option(names = {"-fs"}, description = "Remove synthetic flags from members.")
		private boolean filterSynthetics;

		@Option(names = {"-fg"}, description = "Remove invalid generics from members.")
		private boolean filterGenerics;

		@Option(names = {"-fl"}, description = "Max length of long type names (exceptions, annotations, etc).")
		private int filterLong;

		@Option(names = {"-dec"}, description = "Decompiler to use.")
		private String decompiler;


		@Override
		public Void call() throws Exception {
			return null;
		}
	}
}