refactor: new ProbeRunner class (#195)

Author: fraxken

src/ProbeRunner.js

@@ -0,0 +1,141 @@
+// Import all the probes
+import isUnsafeCallee from "./probes/isUnsafeCallee.js";
+import isLiteral from "./probes/isLiteral.js";
+import isLiteralRegex from "./probes/isLiteralRegex.js";
+import isRegexObject from "./probes/isRegexObject.js";
+import isVariableDeclaration from "./probes/isVariableDeclaration.js";
+import isRequire from "./probes/isRequire.js";
+import isImportDeclaration from "./probes/isImportDeclaration.js";
+import isMemberExpression from "./probes/isMemberExpression.js";
+import isArrayExpression from "./probes/isArrayExpression.js";
+import isFunction from "./probes/isFunction.js";
+import isAssignmentExpression from "./probes/isAssignmentExpression.js";
+import isObjectExpression from "./probes/isObjectExpression.js";
+import isUnaryExpression from "./probes/isUnaryExpression.js";
+import isWeakCrypto from "./probes/isWeakCrypto.js";
+import isClassDeclaration from "./probes/isClassDeclaration.js";
+import isMethodDefinition from "./probes/isMethodDefinition.js";
+
+// Import Internal Dependencies
+import { SourceFile } from "./SourceFile.js";
+
+/**
+ * @typedef {Object} Probe
+ * @property {string} name
+ * @property {any} validateNode
+ * @property {(node: any, options: any) => any} main
+ * @property {(options: any) => void} teardown
+ * @property {boolean} [breakOnMatch=false]
+ * @property {string} [breakGroup]
+ */
+
+export const ProbeSignals = Object.freeze({
+  Break: Symbol.for("breakWalk"),
+  Skip: Symbol.for("skipWalk")
+});
+
+export class ProbeRunner {
+  /**
+   * Note:
+   * The order of the table has an importance/impact on the correct execution of the probes
+   *
+   * @type {Probe[]}
+   */
+  static Defaults = [
+    isUnsafeCallee,
+    isLiteral,
+    isLiteralRegex,
+    isRegexObject,
+    isVariableDeclaration,
+    isRequire,
+    isImportDeclaration,
+    isMemberExpression,
+    isAssignmentExpression,
+    isObjectExpression,
+    isArrayExpression,
+    isFunction,
+    isUnaryExpression,
+    isWeakCrypto,
+    isClassDeclaration,
+    isMethodDefinition
+  ];
+
+  /**
+   *
+   * @param {!SourceFile} sourceFile
+   * @param {Probe[]} [probes=ProbeRunner.Defaults]
+   */
+  constructor(
+    sourceFile,
+    probes = ProbeRunner.Defaults
+  ) {
+    this.sourceFile = sourceFile;
+
+    // TODO: Assert/validate all probes? (Good first issue)
+    this.probes = probes;
+  }
+
+  /**
+   * @param {!Probe} probe
+   * @param {!any} node
+   * @returns {null|void|Symbol}
+   */
+  #runProbe(probe, node) {
+    const validationFns = Array.isArray(probe.validateNode) ?
+      probe.validateNode : [probe.validateNode];
+    for (const validateNode of validationFns) {
+      const [isMatching, data = null] = validateNode(
+        node,
+        this.sourceFile
+      );
+
+      if (isMatching) {
+        return probe.main(node, {
+          analysis: this.sourceFile,
+          data
+        });
+      }
+    }
+
+    return null;
+  }
+
+  walk(node) {
+    /** @type {Set<string>} */
+    const breakGroups = new Set();
+
+    for (const probe of this.probes) {
+      if (breakGroups.has(probe.breakGroup)) {
+        continue;
+      }
+
+      try {
+        const result = this.#runProbe(probe, node);
+        if (result === null) {
+          continue;
+        }
+
+        if (result === ProbeSignals.Skip) {
+          return "skip";
+        }
+        if (result === ProbeSignals.Break || probe.breakOnMatch) {
+          const breakGroup = probe.breakGroup || null;
+
+          if (breakGroup === null) {
+            break;
+          }
+          else {
+            breakGroups.add(breakGroup);
+          }
+        }
+      }
+      finally {
+        if (probe.teardown) {
+          probe.teardown({ analysis: this.sourceFile });
+        }
+      }
+    }
+
+    return null;
+  }
+}

src/SourceFile.js

@@ -6,7 +6,7 @@ import { VariableTracer } from "@nodesecure/estree-ast-utils";
 import { rootLocation, toArrayLocation } from "./utils.js";
 import { generateWarning } from "./warnings.js";
 import { isObfuscatedCode, hasTrojanSource } from "./obfuscators/index.js";
-import { runOnProbes } from "./probes/index.js";
+import { ProbeRunner } from "./ProbeRunner.js";
 
 // CONSTANTS
 const kDictionaryStrParts = [
@@ -45,6 +45,7 @@ export class SourceFile {
     this.encodedLiterals = new Map();
     this.warnings = [];
     this.literalScores = [];
+    this.probesRunner = new ProbeRunner(this);
 
     if (hasTrojanSource(sourceCodeString)) {
       this.addWarning("obfuscated-code", "trojan-source");
@@ -161,7 +162,7 @@ export class SourceFile {
       this.inTryStatement = false;
     }
 
-    return runOnProbes(node, this);
+    return this.probesRunner.walk(node);
   }
 }
 

src/probes/index.js

@@ -26,5 +26,7 @@ function main(node, { analysis }) {
 
 export default {
   name: "isArrayExpression",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isArrayExpression.js

@@ -25,5 +25,7 @@ function main(node, options) {
 
 export default {
   name: "isAssignmentExpression",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isAssignmentExpression.js

@@ -49,5 +49,7 @@ function walkBinaryExpression(expr, level = 1) {
 
 export default {
   name: "isBinaryExpression",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isBinaryExpression.js

@@ -21,5 +21,7 @@ function main(node, options) {
 
 export default {
   name: "isClassDeclaration",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isClassDeclaration.js

@@ -34,5 +34,7 @@ function main(node, options) {
 
 export default {
   name: "isFunctionDeclaration",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isFunction.js

@@ -26,5 +26,8 @@ function main(node, options) {
 
 export default {
   name: "isImportDeclaration",
-  validateNode, main, breakOnMatch: true, breakGroup: "import"
+  validateNode,
+  main,
+  breakOnMatch: true,
+  breakGroup: "import"
 };

src/probes/isImportDeclaration.js

@@ -57,5 +57,7 @@ function main(node, options) {
 
 export default {
   name: "isLiteral",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isLiteral.js

@@ -25,5 +25,7 @@ function main(node, options) {
 
 export default {
   name: "isLiteralRegex",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isLiteralRegex.js

@@ -12,5 +12,8 @@ function main(node, options) {
 
 export default {
   name: "isMemberExpression",
-  validateNode, main, breakOnMatch: true, breakGroup: "import"
+  validateNode,
+  main,
+  breakOnMatch: true,
+  breakGroup: "import"
 };

src/probes/isMemberExpression.js

@@ -21,5 +21,7 @@ function main(node, options) {
 
 export default {
   name: "isMethodDefinition",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isMethodDefinition.js

@@ -25,5 +25,7 @@ function main(node, options) {
 
 export default {
   name: "isObjectExpression",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isObjectExpression.js

@@ -43,5 +43,7 @@ function isRegexConstructor(node) {
 
 export default {
   name: "isRegexObject",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };

src/probes/isRegexObject.js

@@ -11,6 +11,9 @@ import {
   getCallExpressionArguments
 } from "@nodesecure/estree-ast-utils";
 
+// Import Internal Dependencies
+import { ProbeSignals } from "../ProbeRunner.js";
+
 function validateNode(node, { tracer }) {
   const id = getCallExpressionIdentifier(node);
   if (id === null) {
@@ -96,7 +99,7 @@ function main(node, options) {
       analysis.addWarning("unsafe-import", null, node.loc);
 
       // We skip walking the tree to avoid anymore warnings...
-      return Symbol.for("skipWalk");
+      return ProbeSignals.Skip;
     }
 
     default:
@@ -160,5 +163,8 @@ function walkRequireCallExpression(nodeToWalk, tracer) {
 
 export default {
   name: "isRequire",
-  validateNode, main, breakOnMatch: true, breakGroup: "import"
+  validateNode,
+  main,
+  breakOnMatch: true,
+  breakGroup: "import"
 };

src/probes/isRequire.js

@@ -22,5 +22,7 @@ function main(node, options) {
 
 export default {
   name: "isUnaryExpression",
-  validateNode, main, breakOnMatch: false
+  validateNode,
+  main,
+  breakOnMatch: false
 };