Add new logger with log levels

end2end/tests/lambda-mongodb.test.js

@@ -18,7 +18,7 @@ t.before(async () => {
 
 t.test("it does not block by default", async (t) => {
   const { stdout, stderr } = await execAsync(
-    "npx --node-options='--no-deprecation' --loglevel=error serverless@3.38.0 invoke local --function login --path payloads/nosql-injection-request.json",
+    "npx --node-options='--no-deprecation' --loglevel=error serverless@3.38.0 invoke local -e AIKIDO_LOG_LEVEL=error --function login --path payloads/nosql-injection-request.json",
     {
       cwd: directory,
     }
@@ -39,7 +39,7 @@ t.test("it does not block by default", async (t) => {
 
 t.test("it blocks when AIKIDO_BLOCKING is true", async (t) => {
   const { stdout, stderr } = await execAsync(
-    "npx --node-options='--no-deprecation' --loglevel=error serverless@3.38.0 invoke local -e AIKIDO_BLOCKING=true --function login --path payloads/nosql-injection-request.json",
+    "npx --node-options='--no-deprecation' --loglevel=error serverless@3.38.0 invoke local -e AIKIDO_BLOCKING=true -e AIKIDO_LOG_LEVEL=error --function login --path payloads/nosql-injection-request.json",
     {
       cwd: directory,
     }

library/agent/Agent.test.ts

@@ -98,8 +98,9 @@ t.test("it sends started event", async (t) => {
   ]);
 
   t.same(logger.getMessages(), [
-    "Starting agent...",
+    "Starting agent v0.0.0",
     "Found token, reporting enabled!",
+    "Enabled blocking mode, attacks will be blocked!",
     "mongodb@6.9.0 is supported!",
   ]);
 });
@@ -137,8 +138,9 @@ t.test("it logs if package is supported or not", async () => {
   agent.onPackageWrapped("shell-quote", { version: "1.8.1", supported: false });
 
   t.same(logger.getMessages(), [
-    "Starting agent...",
+    "Starting agent v0.0.0",
     "Found token, reporting enabled!",
+    "Enabled blocking mode, attacks will be blocked!",
     "shell-quote@1.8.1 is not supported!",
   ]);
 });
@@ -156,9 +158,9 @@ t.test("it starts in non-blocking mode", async () => {
   agent.start([]);
 
   t.same(logger.getMessages(), [
-    "Starting agent...",
-    "Dry mode enabled, no requests will be blocked!",
+    "Starting agent v0.0.0",
     "Found token, reporting enabled!",
+    "Monitoring mode enabled, no attacks will be blocked!",
   ]);
 });
 
@@ -664,11 +666,12 @@ t.test("it logs when failed to report event", async () => {
   await setTimeout(0);
 
   t.same(logger.getMessages(), [
-    "Starting agent...",
+    "Starting agent v0.0.0",
     "Found token, reporting enabled!",
-    "Failed to start agent",
+    "Enabled blocking mode, attacks will be blocked!",
+    "Failed to start agent: Failed to report event",
     "Heartbeat...",
-    "Failed to do heartbeat",
+    "Failed to send heartbeat event to Aikido platform: Failed to report event",
     "Failed to report attack",
   ]);
 });
@@ -687,8 +690,9 @@ t.test("unable to prevent prototype pollution", async () => {
   agent.start([]);
   agent.unableToPreventPrototypePollution({ mongoose: "1.0.0" });
   t.same(logger.getMessages(), [
-    "Starting agent...",
+    "Starting agent v0.0.0",
     "Found token, reporting enabled!",
+    "Enabled blocking mode, attacks will be blocked!",
     "Unable to prevent prototype pollution, incompatible packages found: mongoose@1.0.0",
   ]);
 

library/agent/Agent.ts

@@ -85,13 +85,13 @@
       list.push(`${pkg}@${incompatiblePackages[pkg]}`);
     }
 
-    this.logger.log(
+    this.logger.warn(
       `Unable to prevent prototype pollution, incompatible packages found: ${list.join(" ")}`
     );
   }
 
   onPrototypePollutionPrevented() {
-    this.logger.log("Prevented prototype pollution!");
+    this.logger.debug("Prevented prototype pollution!");
 
     // Will be sent in the next heartbeat
     this.preventedPrototypePollution = true;
@@ -112,6 +112,17 @@
         },
         this.timeoutInMS
       );
+      if (!result.success) {
+        if (result.error === "invalid_token") {
+          this.logger.error(
+            "Unable to access the Aikido platform, please check your token."
+          );
+        } else {
+          this.logger.error(
+            `Failed to connect to the Aikido platform: ${result.error}`
+          );
+        }
+      }
 
       this.updateServiceConfig(result);
 
@@ -128,7 +139,7 @@
     module: string;
     method: string;
   }) {
-    this.logger.log(
+    this.logger.error(
       `Internal error in module "${module}" in method "${method}"\n${error.stack}`
     );
   }
@@ -194,7 +205,7 @@
 
     if (this.token) {
       this.api.report(this.token, attack, this.timeoutInMS).catch(() => {
-        this.logger.log("Failed to report attack");
+        this.logger.error("Failed to report attack");
       });
     }
   }
@@ -203,8 +214,10 @@
    * Sends a heartbeat via the API to the server (only when not in serverless mode)
    */
   private heartbeat(timeoutInMS = this.timeoutInMS) {
-    this.sendHeartbeat(timeoutInMS).catch(() => {
-      this.logger.log("Failed to do heartbeat");
+    this.sendHeartbeat(timeoutInMS).catch((err) => {
+      this.logger.error(
+        `Failed to send heartbeat event to Aikido platform: ${err.message}`
+      );
     });
   }
 
@@ -221,9 +234,13 @@
       if (typeof response.block === "boolean") {
         if (response.block !== this.block) {
           this.block = response.block;
-          this.logger.log(
-            `Block mode has been set to ${this.block ? "on" : "off"}`
-          );
+          if (this.block) {
+            this.logger.info(`Enabled blocking mode, attacks will be blocked!`);
+          } else {
+            this.logger.warn(
+              `Monitoring mode enabled, no attacks will be blocked!`
+            );
+          }
         }
       }
 
@@ -261,7 +278,7 @@
 
   private async sendHeartbeat(timeoutInMS: number) {
     if (this.token) {
-      this.logger.log("Heartbeat...");
+      this.logger.debug("Heartbeat...");
       const stats = this.statistics.getStats();
       const routes = this.routes.asArray();
       const outgoingDomains = this.hostnames.asArray();
@@ -291,6 +308,12 @@
         timeoutInMS
       );
 
+      if (!response.success && response.error === "invalid_token") {
+        this.logger.error(
+          "Unable to access the Aikido platform, please check your token."
+        );
+      }
+
       this.updateServiceConfig(response);
     }
   }
@@ -300,14 +323,14 @@
    */
   private startHeartbeats() {
     if (this.serverless) {
-      this.logger.log(
+      this.logger.debug(
         "Running in serverless environment, not starting heartbeats"
       );
       return;
     }
 
     if (!this.token) {
-      this.logger.log("No token provided, not starting heartbeats");
+      this.logger.debug("No token provided, not starting heartbeats");
       return;
     }
 
@@ -354,7 +377,7 @@
       this.serviceConfig.updateBlockedIPAddresses(blockedIPAddresses);
       this.serviceConfig.updateBlockedUserAgents(blockedUserAgents);
     } catch (error: any) {
-      this.logger.log(`Failed to update blocked lists: ${error.message}`);
+      this.logger.error(`Failed to update blocked lists: ${error.message}`);
     }
   }
 
@@ -367,7 +390,7 @@
       onConfigUpdate: (config) => {
         this.updateServiceConfig({ success: true, ...config });
         this.updateBlockedLists().catch((error) => {
-          this.logger.log(`Failed to update blocked lists: ${error.message}`);
+          this.logger.error(`Failed to update blocked lists: ${error.message}`);
         });
       },
     });
@@ -420,25 +443,30 @@
 
     this.started = true;
 
-    this.logger.log("Starting agent...");
-
-    if (!this.block) {
-      this.logger.log("Dry mode enabled, no requests will be blocked!");
-    }
+    this.logger.info(`Starting agent v${getAgentVersion()}`);
 
     if (this.token) {
-      this.logger.log("Found token, reporting enabled!");
+      this.logger.info("Found token, reporting enabled!");
+
+      if (!this.block) {
+        this.logger.warn(
+          "Monitoring mode enabled, no attacks will be blocked!"
+        );
+      }
     } else {
-      this.logger.log("No token provided, disabling reporting.");
+      this.logger.info("No token found, reporting disabled!");
 
-      if (!this.block && !isAikidoCI()) {
-        // eslint-disable-next-line no-console
-        console.log(
-          "AIKIDO: Running in monitoring only mode without reporting to Aikido Cloud. Set AIKIDO_BLOCK=true to enable blocking."
+      if (!this.block) {
+        this.logger.warn(
+          "Running in monitoring only mode without reporting to Aikido Cloud. Set AIKIDO_BLOCK=true to enable blocking."
         );
       }
     }
 
+    if (this.block) {
+      this.logger.info("Enabled blocking mode, attacks will be blocked!");
+    }
+
     wrapInstalledPackages(wrappers);
 
     // Send startup event and wait for config
@@ -448,17 +476,19 @@
         this.startHeartbeats();
         this.startPollingForConfigChanges();
       })
-      .catch(() => {
-        this.logger.log("Failed to start agent");
+      .catch((err) => {
+        this.logger.error(`Failed to start agent: ${err.message}`);
       });
   }
 
-  onFailedToWrapMethod(module: string, name: string) {
-    this.logger.log(`Failed to wrap method ${name} in module ${module}`);
+  onFailedToWrapMethod(module: string, name: string, error: Error) {
+    this.logger.error(
+      `Failed to wrap method ${name} in module ${module}: ${error.message}`
+    );
   }
 
   onFailedToWrapModule(module: string, error: Error) {
-    this.logger.log(`Failed to wrap module ${module}: ${error.message}`);
+    this.logger.error(`Failed to wrap module ${module}: ${error.message}`);
   }
 
   onPackageWrapped(name: string, details: WrappedPackage) {
@@ -470,19 +500,19 @@
 
     if (details.version) {
       if (details.supported) {
-        this.logger.log(`${name}@${details.version} is supported!`);
+        this.logger.debug(`${name}@${details.version} is supported!`);
       } else {
-        this.logger.log(`${name}@${details.version} is not supported!`);
+        this.logger.warn(`${name}@${details.version} is not supported!`);
       }
     }
   }
 
   onFailedToWrapPackage(module: string) {
-    this.logger.log(`Failed to wrap package ${module}`);
+    this.logger.error(`Failed to wrap package ${module}`);
   }
 
   onFailedToWrapFile(module: string, filename: string) {
-    this.logger.log(`Failed to wrap file ${filename} in module ${module}`);
+    this.logger.error(`Failed to wrap file ${filename} in module ${module}`);
   }
 
   onConnectHostname(hostname: string, port: number | undefined) {
@@ -516,13 +546,17 @@
     return this.routes;
   }
 
-  log(message: string) {
-    this.logger.log(message);
+  getLogger() {
+    return this.logger;
   }
 
   async flushStats(timeoutInMS: number) {
     this.statistics.forceCompress();
-    await this.sendHeartbeat(timeoutInMS);
+    await this.sendHeartbeat(timeoutInMS).catch((err) => {
+      this.logger.error(
+        `Failed to send heartbeat event to Aikido platform: ${err.message}`
+      );
+    });
   }
 
   getRateLimiter() {

library/agent/api/fetchBlockedLists.ts

@@ -25,6 +25,11 @@ export async function fetchBlockedLists(token: Token): Promise<{
   });
 
   if (statusCode !== 200) {
+    if (statusCode === 401) {
+      throw new Error(
+        `Unable to access the Aikido platform, please check your token.`
+      );
+    }
     throw new Error(`Failed to fetch blocked lists: ${statusCode}`);
   }
 

library/agent/context/user.test.ts

@@ -1,5 +1,4 @@
 import * as t from "tap";
-import { wrap } from "../../helpers/wrap";
 import {
   type Context,
   getContext,
@@ -126,31 +125,23 @@ t.test("it logs when setUser has invalid input", async () => {
 t.test(
   "it does not log warning when setUser is called before middleware",
   async () => {
-    const logs: string[] = [];
-    wrap(console, "warn", function warn() {
-      return function warn(message: string) {
-        logs.push(message);
-      };
-    });
+    const logger = new LoggerForTesting();
+    createTestAgent({ logger });
 
     const context = createContext();
     runWithContext(context, () => {
       setUser({ id: "id" });
     });
 
-    t.same(logs, []);
+    t.same(logger.getMessages(), []);
   }
 );
 
 t.test(
   "it logs warning when setUser is called after middleware (once)",
   async () => {
-    const logs: string[] = [];
-    wrap(console, "warn", function warn() {
-      return function warn(message: string) {
-        logs.push(message);
-      };
-    });
+    const logger = new LoggerForTesting();
+    createTestAgent({ logger });
 
     const context = createContext();
     runWithContext(context, () => {
@@ -159,7 +150,7 @@ t.test(
       setUser({ id: "id" });
     });
 
-    t.same(logs, [
+    t.same(logger.getMessages(), [
       "setUser(...) must be called before the Zen middleware is executed.",
     ]);
   }