append/replace custom probe in ASTAnalyzer class

src/AstAnalyser.js

@@ -16,9 +16,10 @@ export class AstAnalyser {
    * @constructor
    * @param { SourceParser } [parser]
    */
-  constructor(parser = new JsSourceParser(), customProbes = []) {
+  constructor(parser = new JsSourceParser(), customProbes = [], mergeMode = "append") {
     this.parser = parser;
     this.customProbes = customProbes;
+    this.mergeMode = mergeMode;
   }
 
   analyse(str, options = Object.create(null)) {
@@ -32,7 +33,7 @@ export class AstAnalyser {
       isEcmaScriptModule: Boolean(module)
     });
 
-    const source = new SourceFile(str, this.customProbes);
+    const source = new SourceFile(str, this.customProbes, this.mergeMode);
 
     // we walk each AST Nodes, this is a purely synchronous I/O
     walk(body, {

src/SourceFile.js

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

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

@@ -11,30 +11,45 @@ import { ProbeSignals } from "../../src/ProbeRunner.js";
  * @see https://github.com/NodeSecure/js-x-ray/issues/221
  */
 // CONSTANTS
-const kIncriminedCodeSample = "const danger = 'danger';";
+const kIncriminedCodeSample = "const danger = 'danger'; const stream = eval('require')('stream');";
 const kWarningUnsafeDanger = "unsafe-danger";
+const kWarningUnsafeImport = "unsafe-import";
+const kWarningUnsafeStmt = "unsafe-stmt";
 
-test("should detect a custom probe alert unsafe-danger", () => {
-  const customProbes = [
-    {
-      name: "customProbeUnsafeDanger",
-      validateNode: (node, sourceFile) => [true]
-      ,
-      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;
+const customProbes = [
+  {
+    name: "customProbeUnsafeDanger",
+    validateNode: (node, sourceFile) => {
+      return [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(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(new JsSourceParser(), customProbes, "replace");
+  const result = analyser.analyse(kIncriminedCodeSample);
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+  assert.equal(result.warnings.length, 1);
 });