Rework lexer/parser to use Antlr4 rather than something custom.

Author: plucia-mitre

src/main/java/org/mitre/pickledcanary/PickledCanary.java

@@ -1,19 +1,21 @@
 
 // Copyright (C) 2023 The MITRE Corporation All Rights Reserved
 
-package org.mitre.pickledcanary;
+package ghidra.pickledcanary;
 
 import java.io.File;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 
-import org.mitre.pickledcanary.patterngenerator.frontend.PatternAssembler;
-import org.mitre.pickledcanary.patterngenerator.frontend.PatternAssembler.AssembleType;
-import org.mitre.pickledcanary.querylanguage.lexer.Lexer;
-import org.mitre.pickledcanary.querylanguage.lexer.ParseTree;
-import org.mitre.pickledcanary.querylanguage.tokenizer.Token;
-import org.mitre.pickledcanary.querylanguage.tokenizer.Tokenizer;
+import org.antlr.v4.runtime.BaseErrorListener;
+import org.antlr.v4.runtime.CharStreams;
+import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.RecognitionException;
+import org.antlr.v4.runtime.Recognizer;
+
+import org.mitre.pickledcanary.patterngenerator.PCVisitor;
+import org.mitre.pickledcanary.patterngenerator.generated.pc_grammar;
+import org.mitre.pickledcanary.patterngenerator.generated.pc_lexer;
 import org.mitre.pickledcanary.search.Pattern;
 import org.mitre.pickledcanary.search.SavedDataAddresses;
 import org.mitre.pickledcanary.search.VmSearch;
@@ -29,178 +31,111 @@
 import ghidra.util.task.TaskMonitor;
 
 /**
- * This Class holds high-level static methods that are useful for parsing and/or
- * searching pickled canary patterns.
+ * This Class holds high-level static methods that are useful for parsing and/or searching pickled
+ * canary patterns.
  */
 public class PickledCanary {
 
+	static class MyErrorListener extends BaseErrorListener {
+		@Override
+		public void syntaxError(Recognizer<?, ?> recognizer, Object offendingSymbol, int line,
+				int charPositionInLine,
+				String msg, RecognitionException e) {
+			throw new RuntimeException("Pattern lexer encountered error when processing line " + line + ":" + charPositionInLine + " " + msg);
+		}
+	}
+	
 	/**
-	 * You probably want to use
-	 * {@link #parseAndAssemble(TaskMonitor, Program, Address, String, Boolean)
-	 * parseAndAssemble} or
-	 * {@link #parseAndRunAll(TaskMonitor, Program, Address, String) parseAndRunAll}
-	 * instead.
+	 * Creates and runs the PC lexer and visitor across a given string pattern. The result can be
+	 * used to generate a JSON or {@link Pattern} output.
+	 * <p>
+	 * You probably want to use {@link #compile(TaskMonitor, String, Program, Address, boolean)
+	 * compile} or {@link #compile(TaskMonitor, String, Program, Address) compile} instead which
+	 * handle these later steps for you as well. This is probably only useful if you're looking to
+	 * capture both types of output later.
+	 * 
+	 * @param monitor
+	 * @param pattern
+	 *            The pattern to lex and visit.
+	 * @param currentProgram
+	 *            The program to use when lexing and visiting the given pattern.
+	 * @param currentAddress
+	 *            The address to use when lexing and visiting.
+	 * @return PCVisitor instance which has already visited all nodes of the given pattern.
 	 */
-	public static ParseTree parsePattern(TaskMonitor monitor, String query) {
+	public static PCVisitor createAndRunVisitor(TaskMonitor monitor, String pattern,
+			final Program currentProgram, final Address currentAddress) {
 
 		monitor.setIndeterminate(true);
+		
+		MyErrorListener errorListener = new MyErrorListener();
 
-		monitor.setMessage("Tokenizing query.");
-		final Tokenizer tokenizer = new Tokenizer(query);
-
-		final LinkedList<Token> tokens = tokenizer.tokenize(true);
+		var chars = CharStreams.fromString(pattern);
+		var lexer = new pc_lexer(chars);
+		lexer.addErrorListener(errorListener);
+		var commonTokenStream = new CommonTokenStream(lexer);
+		var parser = new pc_grammar(commonTokenStream);
+		parser.addErrorListener(errorListener);
 
-		monitor.setMessage("Lexing query.");
-		final Lexer lexer = new Lexer(tokens);
+		var progContext = parser.prog();
 
-		ParseTree out = lexer.lex();
+		var visitor = new PCVisitor(currentProgram, currentAddress, monitor);
+		visitor.visit(progContext);
 
 		monitor.setIndeterminate(false);
-		return out;
-	}
-
-	private static Object assembleInternal(AssembleType type, TaskMonitor monitor, Program program,
-			Address currentAddress, ParseTree parseTree, Boolean removeDebugInfo) {
-		// start a transaction so that we can undo any overwrites we do to the binary
-		int transactionID = program.startTransaction("Pickled Canary Pattern Assemble");
-
-		monitor.setMessage("Creating pattern assembler");
-		final PatternAssembler patternAssembler = new PatternAssembler();
-
-		Object pattern;
-
-		try {
-
-			monitor.setMessage("Assembling pattern");
-
-			pattern = patternAssembler.assemble(type, program, parseTree, currentAddress, program.getLanguage(),
-					monitor, removeDebugInfo);
-
-			// end transaction - discard all overwrites to the binary
-			program.endTransaction(transactionID, false);
-
-		} catch (Exception e) {
-			// end transaction - discard all overwrites to the binary
-			program.endTransaction(transactionID, false);
-			throw e;
-		}
-
-		return pattern;
+		return visitor;
 	}
+	
 
 	/**
 	 * Returns a JSON string of the compiled pattern.
-	 * <p>
-	 * You probably want to use
-	 * {@link #parseAndAssemble(TaskMonitor, Program, Address, String, Boolean)
-	 * parseAndAssemble} instead.
 	 */
-	public static String assemble(TaskMonitor monitor, Program program, Address currentAddress, ParseTree parseTree) {
-		return (String) assembleInternal(AssembleType.JSON, monitor, program, currentAddress, parseTree, false);
+	public static String compile(TaskMonitor monitor, String pattern, Program program,
+			Address address, boolean removeDebugInfo) {
+		return createAndRunVisitor(monitor, pattern, program, address).getJSON(removeDebugInfo);
 	}
 
-	/**
-	 * Returns a JSON string of the compiled pattern without the compile_info key
-	 * information.
-	 * <p>
-	 * You probably want to use
-	 * {@link #parseAndAssemble(TaskMonitor, Program, Address, String, Boolean)
-	 * parseAndAssemble} instead.
-	 */
-	public static String assemble(TaskMonitor monitor, Program program, Address currentAddress, ParseTree parseTree,
-			boolean removeDebugFlag) {
-		return (String) assembleInternal(AssembleType.JSON, monitor, program, currentAddress, parseTree,
-				removeDebugFlag);
+	public static Pattern compile(TaskMonitor monitor, String pattern, Program program,
+			Address address) {
+		return createAndRunVisitor(monitor, pattern, program, address).getPattern();
 	}
 
 	/**
-	 * Consider using
-	 * {@link #assemblePatternWrapped(TaskMonitor, Program, Address, ParseTree)
-	 * assemblePatternWrapped} which does the same thing as this function, but adds
-	 * a starting .* and match instructions
+	 * Runs the given pattern, returning all results in the given program.
 	 */
-	public static Pattern assemblePattern(TaskMonitor monitor, Program program, Address currentAddress,
-			ParseTree parseTree) {
-		return (Pattern) assembleInternal(AssembleType.PATTERN, monitor, program, currentAddress, parseTree, false);
-	}
+	public static List<SavedDataAddresses> parseAndRunAll(TaskMonitor monitor,
+			Program program, Address address, String pattern) {
+		Pattern patternCompiled = compileWrapped(monitor, pattern, program, address);
 
-	/**
-	 * This is the same as running {@link #parsePattern(TaskMonitor, String)
-	 * parsePattern} followed by
-	 * {@link #assemble(TaskMonitor, Program, Address, ParseTree) assemble).
-	 * 
-	 * Returns a compiled JSON pattern.
-	 */
-	public static String parseAndAssemble(TaskMonitor monitor, Program program, Address currentAddress, String query,
-			Boolean removeDebugInfo) {
-
-		final ParseTree parseTree = PickledCanary.parsePattern(monitor, query);
-		return (String) assembleInternal(AssembleType.JSON, monitor, program, currentAddress, parseTree,
-				removeDebugInfo);
-	}
-
-	/**
-	 * Runs the given pattern (query), returning all results in the given program.
-	 */
-	public static List<SavedDataAddresses> parseAndRunAll(TaskMonitor monitor, Program program, Address currentAddress,
-			String query) {
-		final ParseTree parseTree = PickledCanary.parsePattern(monitor, query);
-
-		Pattern pattern = assemblePatternWrapped(monitor, program, currentAddress, parseTree);
-
-		return runAll(monitor, program, pattern);
-	}
-
-	/**
-	 * Runs the given pattern (query), returning all results in the given program.
-	 */
-	public static void parseAndRunAll(TaskMonitor monitor, Program program, Address currentAddress, String query,
-			Accumulator<SavedDataAddresses> accumulator) {
-		final ParseTree parseTree = PickledCanary.parsePattern(monitor, query);
-
-		Pattern pattern = assemblePatternWrapped(monitor, program, currentAddress, parseTree);
-
-		runAll(monitor, program, pattern, accumulator);
+		return runAll(monitor, program, patternCompiled);
 	}
 
 	/**
-	 * Runs the given pattern (query), returning all results in the given program.
+	 * Runs the given pattern, returning all results in the given program.
 	 */
-	public static String parseAssembleAndRunAll(TaskMonitor monitor, Program program, Address currentAddress,
-			String query, Accumulator<SavedDataAddresses> accumulator) {
-		final ParseTree parseTree = PickledCanary.parsePattern(monitor, query);
-
-		Pattern pattern = assemblePatternWrapped(monitor, program, currentAddress, parseTree);
+	public static void parseAndRunAll(TaskMonitor monitor, Program program,
+			Address address,
+			Accumulator<SavedDataAddresses> accumulator, String pattern) {
+		Pattern patternCompiled = compileWrapped(monitor, pattern, program, address);
 
-		runAll(monitor, program, pattern, accumulator);
-
-		return assemble(monitor, program, currentAddress, parseTree);
+		runAll(monitor, program, patternCompiled, accumulator);
 	}
 
 	/**
-	 * Similar to assemblePattern, but adds a .* to the start of the pattern and
-	 * adds instructions to record the start of the match and when the pattern has
-	 * matched.
+	 * Similar to compile, but adds a .* to the start of the pattern and adds instructions to record
+	 * the start of the match and when the pattern has matched.
 	 */
-	public static Pattern assemblePatternWrapped(TaskMonitor monitor, Program program, Address currentAddress,
-			ParseTree parseTree) {
-
-		final Pattern pattern = assemblePattern(monitor, program, currentAddress, parseTree);
-
-		monitor.setMessage("Preparing pattern");
-		Pattern start = Pattern.getDotStar();
-		start.append(Pattern.getSaveStart());
-		pattern.prepend(start);
-		pattern.append(Pattern.getMatch());
-		return pattern;
-
+	public static Pattern compileWrapped(TaskMonitor monitor, String pattern, Program program,
+			Address address) {
+		return PickledCanary.createAndRunVisitor(monitor, pattern, program, address).getPatternWrapped();
 	}
 
 	/**
 	 * Runs the given pattern on the given program. You may prefer to use
-	 * {@link #parseAndRunAll(TaskMonitor, Program, Address, String) parseAndRunAll}
+	 * {@link #parseAndRunAll(TaskMonitor, String, Program, Address) parseAndRunAll}
 	 */
-	public static List<SavedDataAddresses> runAll(TaskMonitor monitor, Program program, Pattern pattern) {
+	public static List<SavedDataAddresses> runAll(TaskMonitor monitor, Program program,
+			Pattern pattern) {
 		monitor.setMessage("Searching");
 		VmSearch vm = new VmSearch(pattern, program.getMemory());
 
@@ -209,7 +144,7 @@ public static List<SavedDataAddresses> runAll(TaskMonitor monitor, Program progr
 
 	/**
 	 * Runs the given pattern on the given program. You may prefer to use
-	 * {@link #parseAndRunAll(TaskMonitor, Program, Address, String) parseAndRunAll}
+	 * {@link #parseAndRunAll(TaskMonitor, String, Program, Address) parseAndRunAll}
 	 */
 	public static void runAll(TaskMonitor monitor, Program program, Pattern pattern,
 			Accumulator<SavedDataAddresses> accumulator) {
@@ -224,13 +159,13 @@ public enum AskFileType {
 	}
 
 	/**
-	 * Custom version of Ghidra's {@code askFile} method. Allows user to select a
-	 * file.
+	 * Custom version of Ghidra's {@code askFile} method. Allows user to select a file.
 	 * 
-	 * @param isSave true if asking user to save JSON file; false if asking user to
-	 *               choose ptn file
+	 * @param isSave
+	 *            true if asking user to save JSON file; false if asking user to choose ptn file
 	 * @return file object to read or write to
-	 * @throws CancelledException for if user clicks cancel button
+	 * @throws CancelledException
+	 *             for if user clicks cancel button
 	 */
 	public static File pcAskFile(final boolean isSave, final AskFileType type, File previousFile)
 			throws CancelledException {
@@ -246,7 +181,8 @@ public static File pcAskFile(final boolean isSave, final AskFileType type, File
 					String jsonFileName;
 					if ((type == AskFileType.JSON) && ptnFileName.endsWith(".ptn")) {
 						jsonFileName = ptnFileName.replace(".ptn", ".json");
-					} else {
+					}
+					else {
 						jsonFileName = ptnFileName;
 					}
 					selectedFile = new File(jsonFileName);