aiwaf-js 0.0.3 → 0.0.4

package/README.md

@@ -19,6 +19,7 @@
 - ✅ Anomaly Detection (Isolation Forest)
 - ✅ Redis Support for multiprocess environments
 - ✅ Offline Training from access logs
+- ✅ **Custom Cache Logic Support**
 
 ---
 
@@ -73,6 +74,31 @@ If Redis is unavailable, it gracefully falls back to in-memory mode.
 
 ---
 
+## Custom Cache Logic (Advanced)
+
+You can inject your own cache logic (in-memory, Redis, hybrid, or file-based) by passing a `cache` object implementing the following interface:
+
+```js
+const myCustomCache = {
+  get: async (key) => { /* return cached value */ },
+  set: async (key, value, options) => { /* store with optional TTL */ },
+  del: async (key) => { /* delete entry */ }
+}
+
+app.use(aiwaf({
+  cache: myCustomCache,
+  staticKeywords: ['.php'],
+  dynamicTopN: 5,
+  MAX_REQ: 10,
+  WINDOW_SEC: 15,
+  FLOOD_REQ: 20,
+}))
+```
+
+This overrides Redis/in-memory usage with your custom strategy for all cache operations.
+
+---
+
 ## Configuration
 
 ```js
@@ -83,6 +109,7 @@ app.use(aiwaf({
   MAX_REQ: 20,
   FLOOD_REQ: 10,
   HONEYPOT_FIELD: 'hp_field',
+  cache: myCustomCache, // optional custom cache injection
 }));
 ```
 
@@ -97,6 +124,7 @@ app.use(aiwaf({
 | `anomalyThreshold` | `ANOMALY_THRESHOLD` | 0.5                         | Threshold for IsolationForest-based anomaly detection.  |
 | `logPath`          | `NODE_LOG_PATH`     | "/var/log/nginx/access.log" | Path to access log file.                                |
 | `logGlob`          | `NODE_LOG_GLOB`     | "${logPath}.*"              | Glob pattern to include rotated/gzipped logs.           |
+| `cache`            | —                   | undefined                   | Custom cache implementation (overrides Redis/memory)    |
 
 ---
 

package/lib/featureUtils.js

@@ -1,23 +1,38 @@
 const STATIC_KW = ['.php', '.xmlrpc', 'wp-', '.env', '.git', '.bak', 'shell'];
 const STATUS_IDX = ['200', '403', '404', '500'];
-const localCache = new Map();
+const defaultCache = new Map();
 const { getClient } = require('./redisClient');
 
+let customCache = null;
+
+function init(opts = {}) {
+  customCache = opts.cache || null;
+}
+
 async function extractFeatures(req) {
   const uri = req.path.toLowerCase();
   const redis = getClient();
 
+  // Custom cache logic
+  if (customCache?.get && customCache?.set) {
+    const cached = await customCache.get(uri);
+    if (cached) return cached;
+  }
+
+  // Redis check
   if (redis) {
     try {
       const cached = await redis.get(`features:${uri}`);
       if (cached) return JSON.parse(cached);
     } catch (err) {
-      console.warn('⚠️ Redis read failed. Using fallback.', err);
+      console.warn('⚠️ Redis read failed. Falling back.', err);
     }
   }
 
-  if (localCache.has(uri)) return localCache.get(uri);
+  // Local cache fallback
+  if (defaultCache.has(uri)) return defaultCache.get(uri);
 
+  // Compute features
   const pathLen = uri.length;
   const kwHits = STATIC_KW.reduce((count, kw) => count + (uri.includes(kw) ? 1 : 0), 0);
   const statusIdx = STATUS_IDX.indexOf(String(req.res?.statusCode || 200));
@@ -26,7 +41,10 @@ async function extractFeatures(req) {
   const total404 = 0;
   const features = [pathLen, kwHits, statusIdx, rt, burst, total404];
 
-  if (redis) {
+  // Write back
+  if (customCache?.set) {
+    await customCache.set(uri, features);
+  } else if (redis) {
     try {
       await redis.set(`features:${uri}`, JSON.stringify(features), { EX: 60 });
     } catch (err) {
@@ -34,8 +52,8 @@ async function extractFeatures(req) {
     }
   }
 
-  localCache.set(uri, features);
+  defaultCache.set(uri, features);
   return features;
 }
 
-module.exports = { extractFeatures };
+module.exports = { extractFeatures, init };

package/lib/rateLimiter.js

@@ -1,56 +1,63 @@
 const NodeCache = require('node-cache');
 const blacklistManager = require('./blacklistManager');
-const { getClient } = require('./redisClient');
 
-const memoryCache = new NodeCache();
+let cacheBackend;
 let opts;
 
+const defaultMemoryCache = new NodeCache();
+const fallbackCache = {
+  get: (key) => Promise.resolve(defaultMemoryCache.get(key)),
+  set: (key, value, ttl) => Promise.resolve(defaultMemoryCache.set(key, value, ttl)),
+  lPush: async (key, value) => {
+    const arr = defaultMemoryCache.get(key) || [];
+    arr.unshift(value);
+    defaultMemoryCache.set(key, arr, opts.WINDOW_SEC);
+  },
+  expire: (key, ttl) => Promise.resolve(defaultMemoryCache.ttl(key, ttl)),
+  lLen: async (key) => {
+    const arr = defaultMemoryCache.get(key) || [];
+    return arr.length;
+  },
+  lRange: async (key, start, end) => {
+    const arr = defaultMemoryCache.get(key) || [];
+    return arr.slice(start, end + 1);
+  }
+};
+
 module.exports = {
   async init(o) {
     opts = o;
+    cacheBackend = o.cache || fallbackCache;
   },
 
   async record(ip) {
     const now = Date.now().toString();
-    const redis = getClient();
 
-    if (redis) {
-      try {
-        const key = `ratelimit:${ip}`;
-        await redis.lPush(key, now);
-        await redis.expire(key, opts.WINDOW_SEC);
-        const count = await redis.lLen(key);
-        if (count > opts.FLOOD_REQ) await blacklistManager.block(ip, 'flood');
-        return;
-      } catch (err) {
-        console.warn('⚠️ Redis error in record(). Using fallback.', err);
+    try {
+      const key = `ratelimit:${ip}`;
+      await cacheBackend.lPush(key, now);
+      await cacheBackend.expire(key, opts.WINDOW_SEC);
+      const count = await cacheBackend.lLen(key);
+      if (count > opts.FLOOD_REQ) {
+        await blacklistManager.block(ip, 'flood');
       }
+    } catch (err) {
+      console.warn('⚠️ Cache error in record().', err);
     }
-
-    const logs = memoryCache.get(ip) || [];
-    logs.push(now);
-    memoryCache.set(ip, logs, opts.WINDOW_SEC);
-    if (logs.length > opts.FLOOD_REQ) await blacklistManager.block(ip, 'flood');
   },
 
   async isBlocked(ip) {
     if (await blacklistManager.isBlocked(ip)) return true;
-    const now = Date.now();
-    const redis = getClient();
 
-    if (redis) {
-      try {
-        const key = `ratelimit:${ip}`;
-        const timestamps = (await redis.lRange(key, 0, -1)).map(Number);
-        const within = timestamps.filter(t => now - t < opts.WINDOW_SEC * 1000);
-        return within.length > opts.MAX_REQ;
-      } catch (err) {
-        console.warn('⚠️ Redis error in isBlocked(). Using fallback.', err);
-      }
+    try {
+      const key = `ratelimit:${ip}`;
+      const now = Date.now();
+      const timestamps = (await cacheBackend.lRange(key, 0, -1)).map(Number);
+      const within = timestamps.filter(t => now - t < opts.WINDOW_SEC * 1000);
+      return within.length > opts.MAX_REQ;
+    } catch (err) {
+      console.warn('⚠️ Cache error in isBlocked().', err);
+      return false;
     }
-
-    const logs = memoryCache.get(ip) || [];
-    const within = logs.filter(t => now - t < opts.WINDOW_SEC * 1000);
-    return within.length > opts.MAX_REQ;
   }
-};
+};

package/lib/wafMiddleware.js

@@ -8,6 +8,7 @@ const anomalyDetector  = require('./anomalyDetector');
 const { extractFeatures } = require('./featureUtils');
 
 module.exports = function aiwaf(opts = {}) {
+  // Allow injecting custom cache backend
   rateLimiter.init(opts);
   keywordDetector.init(opts);
   dynamicKeyword.init(opts);
@@ -77,4 +78,4 @@ module.exports = function aiwaf(opts = {}) {
 
     next();
   };
-};
+};

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "aiwaf-js",
-    "version": "0.0.3",
+    "version": "0.0.4",
     "description": "Adaptive Web Application Firewall middleware for Node.js (Express, Fastify, Hapi, Next.js)",
     "main": "index.js",
     "scripts": {

package/test/waf.test.js

@@ -1,17 +1,28 @@
 const request = require('supertest');
 const express = require('express');
-
-process.env.NODE_ENV = 'test';
-jest.setTimeout(20000);
-
 const db = require('../utils/db');
 const aiwaf = require('../index');
 const dynamicKeyword = require('../lib/dynamicKeyword');
 const redisManager = require('../lib/redisClient');
 const { init: initRateLimiter } = require('../lib/rateLimiter');
 
+process.env.NODE_ENV = 'test';
+jest.setTimeout(20000);
+
 let redisAvailable = false;
 
+const testCache = (() => {
+  const store = new Map();
+  return {
+    async get(key) {
+      return store.has(key) ? store.get(key) : null;
+    },
+    async set(key, val) {
+      store.set(key, val);
+    }
+  };
+})();
+
 beforeAll(async () => {
   const hasTable = await db.schema.hasTable('blocked_ips');
   if (!hasTable) {
@@ -29,7 +40,8 @@ beforeAll(async () => {
   await initRateLimiter({
     WINDOW_SEC: 1,
     MAX_REQ: 5,
-    FLOOD_REQ: 10
+    FLOOD_REQ: 10,
+    cache: testCache // inject test cache
   });
 });
 
@@ -48,7 +60,9 @@ describe('AIWAF-JS Middleware', () => {
       WINDOW_SEC: 1,
       MAX_REQ: 5,
       FLOOD_REQ: 10,
-      HONEYPOT_FIELD: 'hp_field'
+      HONEYPOT_FIELD: 'hp_field',
+      cache: testCache,
+      logger: console
     }));
 
     app.get('/', (req, res) => res.send('OK'));
@@ -64,17 +78,18 @@ describe('AIWAF-JS Middleware', () => {
       .set('x-response-time', '15')
       .expect(403, { error: 'blocked' })
   );
+
   it('continues working if Redis goes down', async () => {
     const redis = redisManager.getClient();
-    if (redis) await redis.quit(); // force disconnect
-  
+    if (redis) await redis.quit();
+
     const segment = `/simulate-${Date.now().toString(36)}`;
     for (let i = 0; i < 3; i++) {
-      await request(app).get(segment).set('X-Forwarded-For', ip); // fallback path
+      await request(app).get(segment).set('X-Forwarded-For', ip);
     }
     await request(app).get(segment).set('X-Forwarded-For', ip).expect(403);
   });
-  
+
   it('allows safe paths', () =>
     request(app)
       .get('/')
@@ -116,9 +131,9 @@ describe('AIWAF-JS Middleware', () => {
   it('learns and blocks dynamic keywords', async () => {
     const segment = `/secret-${Date.now().toString(36)}`;
     for (let i = 0; i < 3; i++) {
-      await request(app).get(segment).set('X-Forwarded-For', ip); // allow learning phase
+      await request(app).get(segment).set('X-Forwarded-For', ip);
     }
-    await request(app).get(segment).set('X-Forwarded-For', ip).expect(403); // block expected
+    await request(app).get(segment).set('X-Forwarded-For', ip).expect(403);
     await request(app).get(segment).set('X-Forwarded-For', ip).expect(403, { error: 'blocked' });
   });
 
@@ -137,4 +152,4 @@ afterAll(async () => {
     await redisManager.getClient().quit();
   }
   await db.destroy();
-});
+});