diff --git a/src/main/java/me/f1nal/trinity/execution/pattern/PatternReplaceSession.java b/src/main/java/me/f1nal/trinity/execution/pattern/PatternReplaceSession.java new file mode 100644 index 0000000..1ca9773 --- /dev/null +++ b/src/main/java/me/f1nal/trinity/execution/pattern/PatternReplaceSession.java @@ -0,0 +1,135 @@ +package me.f1nal.trinity.execution.pattern; + +import me.f1nal.trinity.Trinity; +import me.f1nal.trinity.events.EventMemberModified; +import me.f1nal.trinity.execution.MethodInput; +import me.f1nal.trinity.gui.windows.impl.assembler.AssemblerClipboardCodec; +import me.f1nal.trinity.gui.windows.impl.assembler.AssemblerDocument; +import me.f1nal.trinity.gui.windows.impl.assembler.AssemblerValidationResult; +import me.f1nal.trinity.gui.windows.impl.assembler.AssemblerValidator; +import org.objectweb.asm.tree.AbstractInsnNode; +import org.objectweb.asm.tree.InsnList; +import org.objectweb.asm.tree.LabelNode; +import org.objectweb.asm.tree.MethodNode; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +public final class PatternReplaceSession { + + public enum State { RUNNING, FINISHED } + + public record ReplaceResult(MethodInput method, int replacements) {} + + private final Trinity trinity; + private final InstructionPattern searchPattern; + private final String replacementText; + private final List methods; + private final List results = new ArrayList<>(); + private final List errors = new ArrayList<>(); + private int methodIndex; + private int totalReplacements; + private State state = State.RUNNING; + + public PatternReplaceSession(Trinity trinity, InstructionPattern searchPattern, String replacementText) { + this.trinity = trinity; + this.searchPattern = searchPattern; + this.replacementText = replacementText; + this.methods = trinity.getExecution().getClassList().stream() + .flatMap(c -> c.getMethodMap().values().stream()) + .filter(m -> m.getOwningClass().getDeclaredMethod(m.getName(), m.getDescriptor()) == m) + .toList(); + } + + public void advance(long budgetNanos) { + if (state != State.RUNNING) return; + long deadline = System.nanoTime() + Math.max(100_000L, budgetNanos); + do { + MethodInput method = methods.get(methodIndex++); + try { + int replaced = replaceInMethod(method); + if (replaced > 0) { + results.add(new ReplaceResult(method, replaced)); + totalReplacements += replaced; + } + } catch (Throwable t) { + String msg = t.getMessage() != null ? t.getMessage() : t.getClass().getSimpleName(); + errors.add(method.getOwningClass().getRealName() + "#" + + method.getName() + method.getDescriptor() + ": " + msg); + } + } while (methodIndex < methods.size() && System.nanoTime() < deadline); + + if (methodIndex >= methods.size()) { + state = State.FINISHED; + } + } + + private int replaceInMethod(MethodInput method) { + List matches = InstructionPatternMatcher.findAll(method, searchPattern); + if (matches.isEmpty()) return 0; + + AssemblerDocument doc = new AssemblerDocument(method); + MethodNode workingMethod = doc.getMethod(); + InsnList insnList = workingMethod.instructions; + + Set matchedSet = Collections.newSetFromMap(new IdentityHashMap<>()); + Set matchHeads = Collections.newSetFromMap(new IdentityHashMap<>()); + for (InstructionPatternMatch match : matches) { + for (AbstractInsnNode insn : match.instructions()) { + matchedSet.add(insn); + } + if (!match.instructions().isEmpty()) { + matchHeads.add(match.instructions().get(0)); + } + } + + List newOrder = new ArrayList<>(); + for (AbstractInsnNode insn : insnList) { + if (!matchedSet.contains(insn)) { + newOrder.add(insn); + } else if (matchHeads.contains(insn)) { + for (AbstractInsnNode rep : parseReplacement(workingMethod)) { + newOrder.add(rep.clone(new IdentityHashMap<>())); + } + } + } + + MethodNode candidate = doc.buildCandidate(newOrder); + AssemblerValidationResult validation = AssemblerValidator.validate(method.getOwningClass(), candidate); + if (!validation.isValid()) { + throw new IllegalStateException(String.join("; ", validation.getErrors())); + } + + doc.commit(candidate); + trinity.getExecution().getXrefMap().refreshMethod(method); + trinity.getEventManager().postEvent(new EventMemberModified(method)); + return matches.size(); + } + + private List parseReplacement(MethodNode context) { + Map existingLabels = new LinkedHashMap<>(); + for (AbstractInsnNode insn : context.instructions) { + if (insn instanceof LabelNode label) { + existingLabels.put("L" + existingLabels.size(), label); + } + } + AssemblerClipboardCodec.ParsedInstructions parsed = AssemblerClipboardCodec.parse( + replacementText, + name -> existingLabels.getOrDefault(name, new LabelNode())); + return parsed.instructions(); + } + + public boolean isFinished() { return state == State.FINISHED; } + public float progress() { return methods.isEmpty() ? 1f : (float) methodIndex / methods.size(); } + public int methodsProcessed() { return methodIndex; } + public int methodCount() { return methods.size(); } + public int totalReplacements() { return totalReplacements; } + public int methodsModified() { return results.size(); } + public List results() { return List.copyOf(results); } + public List errors() { return List.copyOf(errors); } +} diff --git a/src/main/java/me/f1nal/trinity/gui/actions/ApplicationActionRegistry.java b/src/main/java/me/f1nal/trinity/gui/actions/ApplicationActionRegistry.java index c53fe18..7ebd440 100644 --- a/src/main/java/me/f1nal/trinity/gui/actions/ApplicationActionRegistry.java +++ b/src/main/java/me/f1nal/trinity/gui/actions/ApplicationActionRegistry.java @@ -14,6 +14,7 @@ import me.f1nal.trinity.gui.windows.impl.cp.ProjectBrowserFrame; import me.f1nal.trinity.gui.windows.impl.navigation.NavigationHistoryWindow; import me.f1nal.trinity.gui.windows.impl.pattern.PatternSearchFrame; +import me.f1nal.trinity.gui.windows.impl.pattern.PatternReplaceFrame; import me.f1nal.trinity.gui.windows.impl.project.create.NewProjectFrame; import me.f1nal.trinity.gui.windows.impl.project.settings.ProjectSettingsWindow; import me.f1nal.trinity.gui.windows.impl.refactor.GlobalRenameWindow; @@ -42,6 +43,7 @@ public final class ApplicationActionRegistry { public static final String XREF_SEARCH = "inspect.xref_search"; public static final String CONSTANT_SEARCH = "inspect.constant_search"; public static final String PATTERN_SEARCH = "inspect.pattern_search"; + public static final String PATTERN_REPLACE = "inspect.pattern_replace"; public static final String VIEW_ALL_STRINGS = "inspect.all_strings"; public static final String GLOBAL_RENAME = "refactor.global_rename"; public static final String NAVIGATE_BACK = "navigation.back"; @@ -104,6 +106,9 @@ private void registerActions() { this.register(PATTERN_SEARCH, "Pattern Search", "Find bytecode instruction patterns", "Inspect", FontAwesomeIcons.Code, List.of("instruction search", "bytecode pattern"), this::hasProject, () -> this.openStatic(PatternSearchFrame.class)); + this.register(PATTERN_REPLACE, "Search & Replace", "Find and replace bytecode instruction patterns", "Inspect", + FontAwesomeIcons.ExchangeAlt, List.of("bytecode replace", "pattern replace", "instruction replace"), this::hasProject, + () -> this.openStatic(PatternReplaceFrame.class)); this.register(VIEW_ALL_STRINGS, "View All Strings", "List every string constant in the project", "Inspect", FontAwesomeIcons.Font, List.of("strings", "string constants"), this::hasProject, this::viewAllStrings); diff --git a/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternReplaceFrame.java b/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternReplaceFrame.java new file mode 100644 index 0000000..6319893 --- /dev/null +++ b/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternReplaceFrame.java @@ -0,0 +1,200 @@ +package me.f1nal.trinity.gui.windows.impl.pattern; + +import com.google.common.eventbus.Subscribe; +import imgui.ImGui; +import imgui.flag.ImGuiWindowFlags; +import me.f1nal.trinity.Main; +import me.f1nal.trinity.Trinity; +import me.f1nal.trinity.decompiler.output.colors.ColoredStringBuilder; +import me.f1nal.trinity.events.EventClassModified; +import me.f1nal.trinity.events.EventClassesLoaded; +import me.f1nal.trinity.events.EventMemberModified; +import me.f1nal.trinity.events.api.IEventListener; +import me.f1nal.trinity.execution.pattern.InstructionPatternCompiler; +import me.f1nal.trinity.execution.pattern.PatternDiagnostic; +import me.f1nal.trinity.execution.pattern.PatternReplaceSession; +import me.f1nal.trinity.gui.components.MemorableCheckboxComponent; +import me.f1nal.trinity.gui.viewport.notifications.Notification; +import me.f1nal.trinity.gui.viewport.notifications.NotificationType; +import me.f1nal.trinity.gui.viewport.notifications.SimpleCaption; +import me.f1nal.trinity.gui.windows.api.StaticWindow; +import me.f1nal.trinity.logging.Logging; +import me.f1nal.trinity.theme.CodeColorScheme; + +public final class PatternReplaceFrame extends StaticWindow implements IEventListener { + + private static final long FRAME_BUDGET_NANOS = 4_000_000L; + + private final AssemblerPatternEditor searchEditor; + private final AssemblerPatternEditor replaceEditor; + private final MemorableCheckboxComponent includeMetadata = new MemorableCheckboxComponent( + "patternReplaceIncludeMetadata", "Include Metadata", false); + + private InstructionPatternCompiler.Compilation searchCompilation; + private String currentSearchText = ""; + private String currentReplaceText = ""; + private PatternReplaceSession session; + private String staleMessage; + private String lastError; + + public PatternReplaceFrame(Trinity trinity) { + super("Bytecode Search & Replace", 700, 530, trinity); + this.setDialog(true); + this.windowFlags = ImGuiWindowFlags.NoCollapse; + this.searchEditor = new AssemblerPatternEditor(trinity); + this.replaceEditor = new AssemblerPatternEditor(trinity); + this.searchCompilation = InstructionPatternCompiler.compile("", false); + trinity.getEventManager().registerListener(this); + } + + @Override + protected void onOpen() { + searchEditor.requestFocus(); + } + + @Override + protected void renderFrame() { + ImGui.textColored(CodeColorScheme.DISABLED, + "Find an instruction pattern across all methods and replace every match."); + + float availableHeight = ImGui.getContentRegionAvailY(); + boolean sessionActive = session != null; + float controlsHeight = ImGui.getFrameHeightWithSpacing() * (sessionActive ? 3.5f : 2.5f); + float splitHeight = Math.max(80f, (availableHeight - controlsHeight) * 0.5f - 4f); + + drawSectionHeader("Search Pattern", CodeColorScheme.KEYWORD_DATA); + AssemblerPatternEditor.EditorState searchState = searchEditor.draw(splitHeight); + currentSearchText = searchState.text(); + if (searchState.changed()) { + searchCompilation = InstructionPatternCompiler.compile(currentSearchText, includeMetadata.isChecked()); + invalidateSession("Pattern changed"); + } + + ImGui.spacing(); + drawSectionHeader("Replacement", CodeColorScheme.NOTIFY_SUCCESS); + AssemblerPatternEditor.EditorState replaceState = replaceEditor.draw(splitHeight); + currentReplaceText = replaceState.text(); + + ImGui.spacing(); + drawDiagnostic(); + ImGui.spacing(); + + boolean metaBefore = includeMetadata.isChecked(); + includeMetadata.draw(); + if (metaBefore != includeMetadata.isChecked()) { + searchCompilation = InstructionPatternCompiler.compile(currentSearchText, includeMetadata.isChecked()); + invalidateSession("Metadata setting changed"); + } + + ImGui.sameLine(); + + if (sessionActive) { + session.advance(FRAME_BUDGET_NANOS); + float barWidth = Math.max(1f, ImGui.getContentRegionAvailX() - 80f); + ImGui.progressBar(session.progress(), barWidth, 0f, + session.methodsProcessed() + " / " + session.methodCount() + " methods"); + ImGui.sameLine(); + if (ImGui.button("Cancel")) { + cancelSession(); + } + if (session != null && session.isFinished()) { + onSessionFinished(); + } + } else { + boolean canReplace = searchCompilation.valid(); + if (!canReplace) ImGui.beginDisabled(); + boolean clicked = ImGui.button("Replace All"); + if (!canReplace) ImGui.endDisabled(); + if (clicked && canReplace) startSession(); + + if (lastError != null) { + ImGui.sameLine(); + ImGui.textColored(CodeColorScheme.NOTIFY_ERROR, lastError); + } + } + } + + private void drawSectionHeader(String label, int color) { + float x = ImGui.getCursorScreenPosX(); + float y = ImGui.getCursorScreenPosY(); + float w = ImGui.getContentRegionAvailX(); + float h = ImGui.getTextLineHeight() + 4f; + ImGui.getWindowDrawList().addRectFilled(x, y, x + w, y + h, + CodeColorScheme.setAlpha(color, 35), 2f); + ImGui.setCursorPosX(ImGui.getCursorPosX() + 4f); + ImGui.textColored(color, label); + ImGui.spacing(); + } + + private void drawDiagnostic() { + if (staleMessage != null) { + ImGui.textColored(CodeColorScheme.NOTIFY_WARN, staleMessage); + return; + } + PatternDiagnostic diagnostic = searchCompilation.primaryDiagnostic(); + if (diagnostic == null) { + ImGui.textColored(CodeColorScheme.NOTIFY_SUCCESS, "Search pattern valid \u2013 ready to replace"); + return; + } + int color = diagnostic.severity() == PatternDiagnostic.Severity.ERROR + ? CodeColorScheme.NOTIFY_ERROR : CodeColorScheme.DISABLED; + ImGui.textColored(color, "Line " + diagnostic.line() + ", col " + + diagnostic.column() + ": " + diagnostic.message()); + } + + private void startSession() { + lastError = null; + staleMessage = null; + session = new PatternReplaceSession(trinity, searchCompilation.pattern(), currentReplaceText); + } + + private void onSessionFinished() { + PatternReplaceSession done = session; + session = null; + staleMessage = null; + + if (!done.errors().isEmpty()) { + lastError = done.errors().size() + " method(s) failed \u2013 see log"; + done.errors().forEach(Logging::warn); + } + + int replaced = done.totalReplacements(); + int methods = done.methodsModified(); + NotificationType type = replaced > 0 ? NotificationType.SUCCESS : NotificationType.INFO; + String msg = replaced > 0 + ? "Replaced " + replaced + " match(es) across " + methods + " method(s)." + : "No matches found for the given pattern."; + + Main.getDisplayManager().addNotification(new Notification( + type, + new SimpleCaption("Search & Replace"), + ColoredStringBuilder.create().fmt(msg).get() + )); + } + + private void cancelSession() { + session = null; + staleMessage = "Replace cancelled"; + } + + private void invalidateSession(String reason) { + session = null; + staleMessage = reason + " \u2013 run Replace All again"; + lastError = null; + } + + @Subscribe + public void onClassesLoaded(EventClassesLoaded event) { + invalidateSession("Project changed"); + } + + @Subscribe + public void onClassModified(EventClassModified event) { + invalidateSession("Project changed"); + } + + @Subscribe + public void onMemberModified(EventMemberModified event) { + invalidateSession("Project changed"); + } +} diff --git a/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternSearchFrame.java b/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternSearchFrame.java index e675a21..fd92b2e 100644 --- a/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternSearchFrame.java +++ b/src/main/java/me/f1nal/trinity/gui/windows/impl/pattern/PatternSearchFrame.java @@ -14,6 +14,7 @@ import me.f1nal.trinity.execution.pattern.PatternSearchSession; import me.f1nal.trinity.gui.components.MemorableCheckboxComponent; import me.f1nal.trinity.gui.windows.api.StaticWindow; +import me.f1nal.trinity.gui.windows.impl.pattern.PatternReplaceFrame; import me.f1nal.trinity.theme.CodeColorScheme; public final class PatternSearchFrame extends StaticWindow implements IEventListener { @@ -66,6 +67,10 @@ protected void renderFrame() { cancel = ImGui.button("Cancel"); } + ImGui.sameLine(); + if (ImGui.button("Search & Replace")) { + Main.getWindowManager().addStaticWindow(PatternReplaceFrame.class); + } ImGui.sameLine(); boolean metadataBefore = includeMetadata.isChecked(); includeMetadata.draw();