feat(customProbes): inject custom probes as param for AstAnalyser

index.js

@@ -9,10 +9,16 @@ function runASTAnalysis(
 ) {
   const {
     customParser = new JsSourceParser(),
+    customProbes = [],
+    skipDefaultProbes = false,
     ...opts
   } = options;
 
-  const analyser = new AstAnalyser(customParser);
+  const analyser = new AstAnalyser({
+    customParser,
+    customProbes,
+    skipDefaultProbes
+  });
 
   return analyser.analyse(str, opts);
 }
@@ -23,10 +29,16 @@ async function runASTAnalysisOnFile(
 ) {
   const {
     customParser = new JsSourceParser(),
+    customProbes = [],
+    skipDefaultProbes = false,
     ...opts
   } = options;
 
-  const analyser = new AstAnalyser(customParser);
+  const analyser = new AstAnalyser({
+    customParser,
+    customProbes,
+    skipDefaultProbes
+  });
 
   return analyser.analyseFile(pathToFile, opts);
 }

src/AstAnalyser.js

@@ -14,10 +14,17 @@ import { JsSourceParser } from "./JsSourceParser.js";
 export class AstAnalyser {
   /**
    * @constructor
-   * @param { SourceParser } [parser]
+   * @param {object} [options={}]
+   * @param {SourceParser} [options.customParser]
+   * @param {Array<object>} [options.customProbes]
+   * @param {boolean} [options.skipDefaultProbes=false]
    */
-  constructor(parser = new JsSourceParser()) {
-    this.parser = parser;
+  constructor(options = {}) {
+    this.parser = options.customParser ?? new JsSourceParser();
+    this.probesOptions = {
+      customProbes: options.customProbes ?? [],
+      skipDefaultProbes: options.skipDefaultProbes ?? false
+    };
   }
 
   analyse(str, options = Object.create(null)) {
@@ -31,7 +38,7 @@ export class AstAnalyser {
       isEcmaScriptModule: Boolean(module)
     });
 
-    const source = new SourceFile(str);
+    const source = new SourceFile(str, this.probesOptions);
 
     // we walk each AST Nodes, this is a purely synchronous I/O
     walk(body, {

src/SourceFile.js

@@ -20,14 +20,19 @@ export class SourceFile {
   encodedLiterals = new Map();
   warnings = [];
 
-  constructor(sourceCodeString) {
+  constructor(sourceCodeString, probesOptions = {}) {
     this.tracer = new VariableTracer()
       .enableDefaultTracing()
       .trace("crypto.createHash", {
         followConsecutiveAssignment: true, moduleName: "crypto"
       });
 
-    this.probesRunner = new ProbeRunner(this);
+    let probes = ProbeRunner.Defaults;
+    if (Array.isArray(probesOptions.customProbes) && probesOptions.customProbes.length > 0) {
+      probes = probesOptions.skipDefaultProbes === true ? probesOptions.customProbes : [...probes, ...probesOptions.customProbes];
+    }
+    this.probesRunner = new ProbeRunner(this, probes);
+
     if (trojan.verify(sourceCodeString)) {
       this.addWarning("obfuscated-code", "trojan-source");
     }

test/AstAnalyser.spec.js

@@ -206,7 +206,7 @@ describe("AstAnalyser", (t) => {
       const preparedSource = getAnalyser().prepareSource(`
       <!--
     // == fake comment == //
-  
+
     const yo = 5;
     //-->
     `, {
@@ -236,6 +236,13 @@ describe("AstAnalyser", (t) => {
         assert.deepEqual([...result.dependencies.keys()], []);
       });
     });
+
+    it("should instantiate with correct default ASTOptions", () => {
+      const analyser = new AstAnalyser();
+      assert(analyser.parser instanceof JsSourceParser || typeof analyser.parser.customParser === "object");
+      assert.deepStrictEqual(analyser.probesOptions.customProbes, []);
+      assert.strictEqual(analyser.probesOptions.skipDefaultProbes, false);
+    });
   });
 });
 

test/issues/221-inject-custom-probes.spec.js

@@ -0,0 +1,88 @@
+// Import Node.js Dependencies
+import { test } from "node:test";
+import assert from "node:assert";
+
+// Import Internal Dependencies
+import { JsSourceParser } from "../../src/JsSourceParser.js";
+import { AstAnalyser } from "../../src/AstAnalyser.js";
+import { ProbeSignals } from "../../src/ProbeRunner.js";
+import { runASTAnalysis } from "../../index.js";
+
+/**
+ * @see https://github.com/NodeSecure/js-x-ray/issues/221
+ */
+// CONSTANTS
+const kIncriminedCodeSample = "const danger = 'danger'; const stream = eval('require')('stream');";
+const kWarningUnsafeDanger = "unsafe-danger";
+const kWarningUnsafeImport = "unsafe-import";
+const kWarningUnsafeStmt = "unsafe-stmt";
+
+const customProbes = [
+  {
+    name: "customProbeUnsafeDanger",
+    validateNode: (node, sourceFile) => [node.type === "VariableDeclaration" && node.declarations[0].init.value === "danger"]
+    ,
+    main: (node, options) => {
+      const { sourceFile, data: calleeName } = options;
+      if (node.declarations[0].init.value === "danger") {
+        sourceFile.addWarning("unsafe-danger", calleeName, node.loc);
+
+        return ProbeSignals.Skip;
+      }
+
+      return null;
+    }
+  }
+];
+
+test("should append to list of probes (default)", () => {
+  const analyser = new AstAnalyser({ customParser: new JsSourceParser(), customProbes });
+  const result = analyser.analyse(kIncriminedCodeSample);
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+  assert.equal(result.warnings[1].kind, kWarningUnsafeImport);
+  assert.equal(result.warnings[2].kind, kWarningUnsafeStmt);
+  assert.equal(result.warnings.length, 3);
+});
+
+test("should replace list of probes", () => {
+  const analyser = new AstAnalyser({
+    parser: new JsSourceParser(),
+    customProbes,
+    skipDefaultProbes: true
+  });
+  const result = analyser.analyse(kIncriminedCodeSample);
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+  assert.equal(result.warnings.length, 1);
+});
+
+test("should append list of probes using runASTAnalysis", () => {
+  const result = runASTAnalysis(
+    kIncriminedCodeSample,
+    {
+      parser: new JsSourceParser(),
+      customProbes,
+      skipDefaultProbes: false
+    }
+  );
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+  assert.equal(result.warnings[1].kind, kWarningUnsafeImport);
+  assert.equal(result.warnings[2].kind, kWarningUnsafeStmt);
+  assert.equal(result.warnings.length, 3);
+});
+
+test("should replace list of probes using runASTAnalysis", () => {
+  const result = runASTAnalysis(
+    kIncriminedCodeSample,
+    {
+      parser: new JsSourceParser(),
+      customProbes,
+      skipDefaultProbes: true
+    }
+  );
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+  assert.equal(result.warnings.length, 1);
+});

types/api.d.ts

@@ -42,13 +42,35 @@ interface RuntimeOptions {
   /**
    * @default false
    */
-  isMinified?: boolean;
+  removeHTMLComments?: boolean;
   /**
    * @default false
    */
-  removeHTMLComments?: boolean;
-
+  isMinified?: boolean;
+}
+
+interface RuntimeFileOptions extends Omit<RuntimeOptions, "isMinified"> {
+  packageName?: string;
+}
+
+interface AstAnalyserOptions {
+  /**
+   * @default JsSourceParser
+   */
   customParser?: SourceParser;
+  /**
+   * @default []
+   */
+  customProbes?: Probe[];
+  /**
+   * @default false
+   */
+  skipDefaultProbes?: boolean;
+}
+
+interface Probe {
+  validateNode: Function | Function[];
+  main: Function;
 }
 
 interface Report {
@@ -59,10 +81,6 @@ interface Report {
   isOneLineRequire: boolean;
 }
 
-interface RuntimeFileOptions extends Omit<RuntimeOptions, "isMinified"> {
-  packageName?: string;
-}
-
 type ReportOnFile = {
   ok: true,
   warnings: Warning[];
@@ -78,10 +96,10 @@ interface SourceParser {
 }
 
 declare class AstAnalyser {
-  constructor(parser?: SourceParser);
-  analyse: (str: string, options?: Omit<RuntimeOptions, "customParser">) => Report;
-  analyzeFile(pathToFile: string, options?: Omit<RuntimeFileOptions, "customParser">): Promise<ReportOnFile>;
+  constructor(options?: AstAnalyserOptions);
+  analyse: (str: string, options?: RuntimeOptions) => Report;
+  analyzeFile(pathToFile: string, options?: RuntimeFileOptions): Promise<ReportOnFile>;
 }
 
-declare function runASTAnalysis(str: string, options?: RuntimeOptions): Report;
-declare function runASTAnalysisOnFile(pathToFile: string, options?: RuntimeFileOptions): Promise<ReportOnFile>;
+declare function runASTAnalysis(str: string, options?: RuntimeOptions & AstAnalyserOptions): Report;
+declare function runASTAnalysisOnFile(pathToFile: string, options?: RuntimeFileOptions & AstAnalyserOptions): Promise<ReportOnFile>;