backend/src/main/java/ai/decompile/engine/service/AbstractIdaEngineService.java

383 lines
12 KiB
Java

package ai.decompile.engine.service;
import ai.decompile.docker.service.ContainerService;
import ai.decompile.engine.config.EngineProperties;
import ai.decompile.engine.model.AnalysisResult;
import ai.decompile.engine.model.FunctionResult;
import ai.decompile.engine.model.LabelEntry;
import ai.decompile.engine.model.SegmentEntry;
import ai.decompile.engine.model.XrefEntry;
import com.github.dockerjava.api.model.Bind;
import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.CodingErrorAction;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Objects;
import lombok.extern.log4j.Log4j2;
@Log4j2
public abstract class AbstractIdaEngineService implements EngineService {
private static final String KEY_ENTRY_POINT = "\"entry_point\": \"";
private static final String KEY_ADDRESS = "\",\n \"address\": \"";
private static final String KEY_SIZE = "\",\n \"size\": ";
private static final String KEY_FLAGS = ",\n \"flags\": \"";
private static final String KEY_XREFS = "\",\n \"xrefs\": \"";
private static final String KEY_ASSEMBLY = "\",\n \"assembly\": \"";
private static final String KEY_LABELS = "\",\n \"labels\": \"";
private static final String KEY_ENTRY_END = "\"\n }";
private static final String KEY_SEGMENTS_START = "\"segments\": [";
private static final String KEY_SEGMENT_NAME = "\"name\": \"";
private static final String KEY_SEGMENT_START = "\", \"start\": \"";
private static final String KEY_SEGMENT_END = "\", \"end\": \"";
protected final ContainerService containers;
protected final EngineProperties engineProperties;
protected AbstractIdaEngineService(
ContainerService containers, EngineProperties engineProperties) {
this.containers = containers;
this.engineProperties = engineProperties;
}
protected abstract String getDockerfilePath();
protected abstract String getConfigKey();
protected abstract String getTempDirPrefix();
@Override
public AnalysisResult analyze(String filePath) {
var configKey = getConfigKey();
var config = engineProperties.getEngineConfig(configKey);
var absolutePath = Path.of(filePath).toAbsolutePath().normalize();
var filename = absolutePath.getFileName().toString();
containers.ensureImage(config.getImage(), new File(getDockerfilePath()));
var outputDir = createTempDir();
var containerId = createContainer(config.getImage(), filename, absolutePath, outputDir);
try {
var exitCode = containers.waitForExit(containerId, config.getTimeoutSeconds());
var logs = containers.captureLogs(containerId, 2);
validateExitCode(configKey, exitCode, logs);
var outputFile = validateOutputFile(configKey, outputDir, logs);
return parseOutput(readStringTolerant(outputFile));
} catch (Exception e) {
log.error("{} analysis failed for file {}: {}", configKey, filePath, e.getMessage());
throw new RuntimeException(configKey + " analysis failed: " + e.getMessage(), e);
} finally {
containers.stopAndRemove(containerId);
cleanupTempDir(outputDir);
}
}
private Path createTempDir() {
try {
return Files.createTempDirectory(getTempDirPrefix());
} catch (IOException e) {
throw new RuntimeException("Failed to create temporary output directory", e);
}
}
private String createContainer(String image, String filename, Path absolutePath, Path outputDir) {
var inputBind = Bind.parse(absolutePath.getParent() + ":/input:ro");
var outputBind = Bind.parse(outputDir.toAbsolutePath() + ":/output:rw");
var containerId =
containers.createContainer(image, new String[] {filename}, inputBind, outputBind);
containers.startContainer(containerId);
return containerId;
}
private static void validateExitCode(String configKey, int exitCode, String logs) {
if (exitCode != 0) {
log.error("{} container exited with code {}: {}", configKey, exitCode, logs);
throw new RuntimeException(
configKey + " analysis failed with exit code " + exitCode + ". Output: " + logs);
}
}
private static Path validateOutputFile(String configKey, Path outputDir, String logs) {
var outputFile = outputDir.resolve("output.json");
if (!Files.exists(outputFile)) {
throw new RuntimeException(
configKey + " analysis did not produce output.json. Container logs: " + logs);
}
return outputFile;
}
private static void cleanupTempDir(Path dir) {
try (var walk = Files.walk(dir)) {
walk.sorted(Comparator.reverseOrder()).forEach(AbstractIdaEngineService::deleteSilently);
} catch (IOException ignored) {
}
}
private static void deleteSilently(Path path) {
try {
Files.deleteIfExists(path);
} catch (IOException ignored) {
}
}
private static String readStringTolerant(Path path) throws IOException {
var decoder =
StandardCharsets.UTF_8
.newDecoder()
.onMalformedInput(CodingErrorAction.REPLACE)
.onUnmappableCharacter(CodingErrorAction.REPLACE);
return decoder.decode(ByteBuffer.wrap(Files.readAllBytes(path))).toString();
}
private static AnalysisResult parseOutput(String raw) {
var text = raw.replace("\r\n", "\n");
var entryPoint = parseEntryPoint(text);
var segments = parseSegments(text);
var functions = new ArrayList<FunctionResult>();
var sections = text.split("\"function\"\\s*:\\s*\"");
for (int i = 1; i < sections.length; i++) {
var func = parseFunctionEntry(sections[i]);
if (Objects.nonNull(func)) {
functions.add(func);
}
}
if (functions.isEmpty()) {
log.warn(
"IDA produced no parsable function entries. Raw output (first 500 chars): {}",
raw.length() > 500 ? raw.substring(0, 500) : raw);
}
return new AnalysisResult(entryPoint, segments, functions);
}
private static String parseEntryPoint(String text) {
var entryRemainder = after(text, KEY_ENTRY_POINT);
if (Objects.isNull(entryRemainder)) {
return "unknown";
}
var entryVal = extractField(entryRemainder, "\"");
if (Objects.isNull(entryVal)) {
return "unknown";
}
return entryVal;
}
private static FunctionResult parseFunctionEntry(String section) {
var name = extractField(section, KEY_ADDRESS);
if (Objects.isNull(name)) {
return null;
}
var remainder = after(section, KEY_ADDRESS);
if (Objects.isNull(remainder)) {
return null;
}
var address = extractField(remainder, KEY_SIZE);
if (Objects.isNull(address)) {
return null;
}
remainder = after(remainder, KEY_SIZE);
if (Objects.isNull(remainder)) {
return null;
}
var sizeStr = extractField(remainder, KEY_FLAGS);
if (Objects.isNull(sizeStr)) {
return null;
}
var size = parseSize(sizeStr);
remainder = after(remainder, KEY_FLAGS);
if (Objects.isNull(remainder)) {
return null;
}
var flags = extractField(remainder, KEY_XREFS);
if (Objects.isNull(flags)) {
return null;
}
remainder = after(remainder, KEY_XREFS);
if (Objects.isNull(remainder)) {
return null;
}
var xrefsRaw = extractField(remainder, KEY_ASSEMBLY);
if (Objects.isNull(xrefsRaw)) {
return null;
}
remainder = after(remainder, KEY_ASSEMBLY);
if (Objects.isNull(remainder)) {
return null;
}
var assembly = extractField(remainder, KEY_LABELS);
if (Objects.isNull(assembly)) {
return null;
}
remainder = after(remainder, KEY_LABELS);
if (Objects.isNull(remainder)) {
return null;
}
var labelsRaw = extractField(remainder, KEY_ENTRY_END);
if (Objects.isNull(labelsRaw)) {
return null;
}
var xrefs = parseXrefs(xrefsRaw.trim());
var labels = parseLabels(labelsRaw.trim());
return new FunctionResult(
name,
address,
size,
flags.endsWith("|") ? flags.substring(0, flags.length() - 1) : flags,
assembly,
null,
null,
xrefs,
labels);
}
private static int parseSize(String sizeStr) {
try {
return Integer.parseInt(sizeStr.trim());
} catch (NumberFormatException e) {
return 0;
}
}
private static String extractField(String text, String endDelim) {
var end = text.indexOf(endDelim);
return end < 0 ? null : text.substring(0, end);
}
private static String after(String text, String delim) {
var idx = text.indexOf(delim);
return idx < 0 ? null : text.substring(idx + delim.length());
}
private static List<SegmentEntry> parseSegments(String text) {
var segments = new ArrayList<SegmentEntry>();
var remainder = after(text, KEY_SEGMENTS_START);
if (Objects.isNull(remainder)) {
return segments;
}
var segmentsBlock = extractField(remainder, "\n ],\n");
if (Objects.isNull(segmentsBlock)) {
return segments;
}
var parts = segmentsBlock.split("},\\s*\\{");
for (var part : parts) {
var segment = parseSegmentEntry(part);
if (Objects.nonNull(segment)) {
segments.add(segment);
}
}
return segments;
}
private static SegmentEntry parseSegmentEntry(String part) {
var clean = part.replace("{", "").replace("}", "").trim();
if (clean.isBlank()) {
return null;
}
var nameRemainder = after(clean, KEY_SEGMENT_NAME);
if (Objects.isNull(nameRemainder)) {
return null;
}
var name = extractField(nameRemainder, KEY_SEGMENT_START);
if (Objects.isNull(name)) {
return null;
}
var startRemainder = after(nameRemainder, KEY_SEGMENT_START);
if (Objects.isNull(startRemainder)) {
return null;
}
var startAddress = extractField(startRemainder, KEY_SEGMENT_END);
if (Objects.isNull(startAddress)) {
return null;
}
var endRemainder = after(startRemainder, KEY_SEGMENT_END);
if (Objects.isNull(endRemainder)) {
return null;
}
var endAddress = extractField(endRemainder, "\"");
if (Objects.isNull(endAddress)) {
return null;
}
return new SegmentEntry(name, startAddress, endAddress);
}
private static List<XrefEntry> parseXrefs(String raw) {
if (Objects.isNull(raw) || raw.isBlank()) {
return List.of();
}
return Arrays.stream(raw.trim().split("\\s+"))
.map(AbstractIdaEngineService::parseXrefEntry)
.toList();
}
private static XrefEntry parseXrefEntry(String token) {
var colon = token.lastIndexOf(':');
if (colon < 0) {
return new XrefEntry(token, "CALL");
}
return new XrefEntry(token.substring(0, colon), token.substring(colon + 1));
}
private static List<LabelEntry> parseLabels(String raw) {
if (Objects.isNull(raw) || raw.isBlank()) {
return List.of();
}
return Arrays.stream(raw.trim().split("\\s+"))
.distinct()
.map(AbstractIdaEngineService::parseLabelEntry)
.toList();
}
private static LabelEntry parseLabelEntry(String token) {
var lastColon = token.lastIndexOf(':');
if (lastColon < 0) {
return new LabelEntry(token, "", "CALL");
}
var type = token.substring(lastColon + 1);
var nameAddr = token.substring(0, lastColon);
var addrColon = nameAddr.lastIndexOf(':');
if (addrColon < 0) {
return new LabelEntry(nameAddr, "", type);
}
return new LabelEntry(
nameAddr.substring(0, addrColon), nameAddr.substring(addrColon + 1), type);
}
}