Skip to content

Commit

Permalink
Merge pull request #324 from AikidoSec/new-hooks-system
Browse files Browse the repository at this point in the history 
Implement new hook system
  • Loading branch information
@hansott
hansott authored Oct 14, 2024
2 parents 0719612 + 8ae3e5f commit 02ac426
Show file tree
Hide file tree
Showing 107 changed files with 3,119 additions and 2,083 deletions.
2 changes: 1 addition & 1 deletion benchmarks/express/benchmark.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ const { promisify } = require("util");
const exec = promisify(require("child_process").exec);

// Accepted percentage of performance decrease
const AcceptedDecrease = 30; // %
const AcceptedDecrease = 40; // %

function generateWrkCommandForUrl(url) {
// Define the command with awk included
Expand Down
164 changes: 117 additions & 47 deletions library/agent/Agent.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import { LoggerForTesting } from "./logger/LoggerForTesting";
import { LoggerNoop } from "./logger/LoggerNoop";
import { Wrapper } from "./Wrapper";
import { Context } from "./Context";
import { createTestAgent } from "../helpers/createTestAgent";

t.test("it throws error if serverless is empty string", async () => {
t.throws(
Expand All @@ -30,12 +31,17 @@ t.test("it throws error if serverless is empty string", async () => {
});

t.test("it sends started event", async (t) => {
const logger = new LoggerNoop();
const logger = new LoggerForTesting();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([new MongoDB()]);

const mongodb = require("mongodb");

t.match(api.getEvents(), [
{
type: "started",
Expand All @@ -44,13 +50,11 @@ t.test("it sends started event", async (t) => {
hostname: hostname(),
version: "0.0.0",
ipAddress: ip(),
packages: {
mongodb: "6.8.0",
},
packages: {},
preventedPrototypePollution: false,
nodeEnv: "",
serverless: false,
stack: ["mongodb"],
stack: [],
os: {
name: platform(),
version: release(),
Expand All @@ -62,13 +66,22 @@ t.test("it sends started event", async (t) => {
},
},
]);

t.same(logger.getMessages(), [
"Starting agent...",
"Found token, reporting enabled!",
"[email protected] is supported!",
]);
});

t.test("it throws error if already started", async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([new MongoDB()]);
t.throws(() => agent.start([new MongoDB()]), "Agent already started!");
});
Expand All @@ -82,9 +95,15 @@ class WrapperForTesting implements Wrapper {
t.test("it logs if package is supported or not", async () => {
const logger = new LoggerForTesting();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([new WrapperForTesting()]);

agent.onPackageWrapped("shell-quote", { version: "1.8.1", supported: false });

t.same(logger.getMessages(), [
"Starting agent...",
"Found token, reporting enabled!",
Expand All @@ -95,22 +114,30 @@ t.test("it logs if package is supported or not", async () => {
t.test("it starts in non-blocking mode", async () => {
const logger = new LoggerForTesting();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(false, logger, api, token, undefined);
agent.start([new MongoDB()]);
const agent = createTestAgent({
block: false,
api,
logger,
token: new Token("123"),
});
agent.start([]);

t.same(logger.getMessages(), [
"Starting agent...",
"Dry mode enabled, no requests will be blocked!",
"Found token, reporting enabled!",
"[email protected] is supported!",
]);
});

t.test("when prevent prototype pollution is enabled", async (t) => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, "lambda");
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
serverless: "lambda",
});
agent.onPrototypePollutionPrevented();
agent.start([]);
t.match(api.getEvents(), [
Expand All @@ -126,18 +153,24 @@ t.test("when prevent prototype pollution is enabled", async (t) => {
t.test("it does not start interval in serverless mode", async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, "lambda");

const agent = createTestAgent({
api,
logger,
token: new Token("123"),
serverless: "lambda",
});
// This would otherwise keep the process running
agent.start([]);
});

t.test("when attack detected", async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.onDetectedAttack({
module: "mongodb",
kind: "nosql_injection",
Expand Down Expand Up @@ -196,8 +229,11 @@ t.test("when attack detected", async () => {
t.test("it checks if user agent is a string", async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.onDetectedAttack({
module: "mongodb",
kind: "nosql_injection",
Expand Down Expand Up @@ -267,8 +303,11 @@ t.test(
block: true,
receivedAnyStats: false,
});
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);
t.match(api.getEvents(), [
{
Expand Down Expand Up @@ -332,8 +371,11 @@ t.test(
block: true,
receivedAnyStats: false,
});
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);
t.match(api.getEvents(), [
{
Expand Down Expand Up @@ -368,8 +410,11 @@ t.test("it sends heartbeat when reached max timings", async () => {

const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);
for (let i = 0; i < 1000; i++) {
agent.getInspectionStatistics().onInspectedCall({
Expand Down Expand Up @@ -458,8 +503,11 @@ t.test("it logs when failed to report event", async () => {

const logger = new LoggerForTesting();
const api = new ReportingAPIThatThrows();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);

await waitForCalls();
Expand Down Expand Up @@ -514,8 +562,11 @@ t.test("unable to prevent prototype pollution", async () => {

const logger = new LoggerForTesting();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);
agent.unableToPreventPrototypePollution({ mongoose: "1.0.0" });
t.same(logger.getMessages(), [
Expand All @@ -542,9 +593,11 @@ t.test("unable to prevent prototype pollution", async () => {
t.test("when payload is object", async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);

const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.onDetectedAttack({
module: "mongodb",
kind: "nosql_injection",
Expand Down Expand Up @@ -642,8 +695,11 @@ t.test("it sends hostnames and routes along with heartbeat", async () => {

const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
agent.start([]);

agent.onConnectHostname("aikido.dev", 443);
Expand Down Expand Up @@ -738,8 +794,11 @@ t.test(
async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
t.same(agent.shouldBlock(), true);
agent.start([]);

Expand All @@ -755,8 +814,12 @@ t.test(
async () => {
const logger = new LoggerNoop();
const api = new ReportingAPIForTesting();
const token = new Token("123");
const agent = new Agent(false, logger, api, token, undefined);
const agent = createTestAgent({
block: false,
api,
logger,
token: new Token("123"),
});
t.same(agent.shouldBlock(), false);
agent.start([]);

Expand All @@ -778,8 +841,12 @@ t.test("it enables blocking mode after sending startup event", async () => {
allowedIPAddresses: [],
block: true,
});
const token = new Token("123");
const agent = new Agent(false, logger, api, token, undefined);
const agent = createTestAgent({
token: new Token("123"),
block: false,
api,
logger,
});
t.same(agent.shouldBlock(), false);
agent.start([]);

Expand All @@ -800,8 +867,11 @@ t.test("it goes into monitoring mode after sending startup event", async () => {
allowedIPAddresses: [],
block: false,
});
const token = new Token("123");
const agent = new Agent(true, logger, api, token, undefined);
const agent = createTestAgent({
api,
logger,
token: new Token("123"),
});
t.same(agent.shouldBlock(), true);
agent.start([]);

Expand Down
37 changes: 21 additions & 16 deletions library/agent/Agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -409,22 +409,7 @@ export class Agent {
}
}

this.wrappedPackages = wrapInstalledPackages(this, wrappers);

for (const pkg in this.wrappedPackages) {
const details = this.wrappedPackages[pkg];

/* c8 ignore next 3 */
if (!details.version) {
continue;
}

if (details.supported) {
this.logger.log(`${pkg}@${details.version} is supported!`);
} else {
this.logger.log(`${pkg}@${details.version} is not supported!`);
}
}
wrapInstalledPackages(wrappers);

// Send startup event and wait for config
// Then start heartbeats and polling for config changes
Expand All @@ -442,6 +427,26 @@ export class Agent {
this.logger.log(`Failed to wrap method ${name} in module ${module}`);
}

onFailedToWrapModule(module: string, error: Error) {
this.logger.log(`Failed to wrap module ${module}: ${error.message}`);
}

onPackageWrapped(name: string, details: WrappedPackage) {
if (this.wrappedPackages[name]) {
// Already reported as wrapped
return;
}
this.wrappedPackages[name] = details;

if (details.version) {
if (details.supported) {
this.logger.log(`${name}@${details.version} is supported!`);
} else {
this.logger.log(`${name}@${details.version} is not supported!`);
}
}
}

onFailedToWrapPackage(module: string) {
this.logger.log(`Failed to wrap package ${module}`);
}
Expand Down
Loading

0 comments on commit 02ac426

Please sign in to comment.