@kya-os/agentshield-nextjs 0.1.32 → 0.1.34

package/dist/index.js

@@ -1057,6 +1057,466 @@ function createAgentShieldMiddleware2(config) {
   };
 }
 
+// src/wasm-confidence.ts
+async function checkWasmAvailability() {
+  try {
+    if (typeof WebAssembly === "undefined") {
+      return false;
+    }
+    const wasmCode = new Uint8Array([
+      0,
+      97,
+      115,
+      109,
+      1,
+      0,
+      0,
+      0
+    ]);
+    const module = await WebAssembly.compile(wasmCode);
+    await WebAssembly.instantiate(module);
+    return true;
+  } catch {
+    return false;
+  }
+}
+function shouldIndicateWasmVerification(confidence) {
+  return confidence >= 0.85 && confidence < 1;
+}
+function getWasmConfidenceBoost(baseConfidence, reasons = []) {
+  if (reasons.some(
+    (r) => r.includes("signature_agent") && !r.includes("signature_headers_present")
+  )) {
+    return 1;
+  }
+  if (baseConfidence >= 0.85) {
+    return 0.95;
+  }
+  if (baseConfidence >= 0.7) {
+    return Math.min(baseConfidence * 1.1, 0.9);
+  }
+  return baseConfidence;
+}
+function getVerificationMethod(confidence, reasons = []) {
+  if (reasons.some(
+    (r) => r.includes("signature_agent") && !r.includes("signature_headers_present")
+  )) {
+    return "signature";
+  }
+  if (shouldIndicateWasmVerification(confidence)) {
+    return "wasm-enhanced";
+  }
+  return "pattern";
+}
+
+// src/storage/memory-adapter.ts
+var MemoryStorageAdapter = class {
+  events = /* @__PURE__ */ new Map();
+  sessions = /* @__PURE__ */ new Map();
+  eventTimeline = [];
+  maxEvents = 1e3;
+  maxSessions = 100;
+  async storeEvent(event) {
+    const eventKey = `${event.timestamp}:${event.eventId}`;
+    this.events.set(eventKey, event);
+    this.eventTimeline.push({
+      timestamp: Date.parse(event.timestamp),
+      eventId: eventKey
+    });
+    this.eventTimeline.sort((a, b) => b.timestamp - a.timestamp);
+    if (this.eventTimeline.length > this.maxEvents) {
+      const removed = this.eventTimeline.splice(this.maxEvents);
+      removed.forEach((item) => this.events.delete(item.eventId));
+    }
+  }
+  async getRecentEvents(limit = 100) {
+    const recent = this.eventTimeline.slice(0, limit);
+    return recent.map((item) => this.events.get(item.eventId)).filter((event) => event !== void 0);
+  }
+  async getSessionEvents(sessionId) {
+    const events = [];
+    for (const event of this.events.values()) {
+      if (event.sessionId === sessionId) {
+        events.push(event);
+      }
+    }
+    return events.sort(
+      (a, b) => Date.parse(b.timestamp) - Date.parse(a.timestamp)
+    );
+  }
+  async storeSession(session) {
+    this.sessions.set(session.sessionId, session);
+    if (this.sessions.size > this.maxSessions) {
+      const sortedSessions = Array.from(this.sessions.entries()).sort((a, b) => Date.parse(b[1].lastSeen) - Date.parse(a[1].lastSeen));
+      const toRemove = sortedSessions.slice(this.maxSessions);
+      toRemove.forEach(([id]) => this.sessions.delete(id));
+    }
+  }
+  async getSession(sessionId) {
+    return this.sessions.get(sessionId) || null;
+  }
+  async getRecentSessions(limit = 10) {
+    const sorted = Array.from(this.sessions.values()).sort((a, b) => Date.parse(b.lastSeen) - Date.parse(a.lastSeen));
+    return sorted.slice(0, limit);
+  }
+  async cleanup(olderThan) {
+    const cutoff = olderThan.getTime();
+    this.eventTimeline = this.eventTimeline.filter((item) => {
+      if (item.timestamp < cutoff) {
+        this.events.delete(item.eventId);
+        return false;
+      }
+      return true;
+    });
+    for (const [id, session] of this.sessions.entries()) {
+      if (Date.parse(session.lastSeen) < cutoff) {
+        this.sessions.delete(id);
+      }
+    }
+  }
+};
+
+// src/storage/redis-adapter.ts
+var RedisStorageAdapter = class {
+  redis;
+  ttl;
+  keyPrefix = "agent-shield";
+  constructor(redis, ttl = 86400) {
+    this.redis = redis;
+    this.ttl = ttl;
+  }
+  eventKey(timestamp, eventId) {
+    return `${this.keyPrefix}:events:${timestamp}:${eventId}`;
+  }
+  sessionKey(sessionId) {
+    return `${this.keyPrefix}:sessions:${sessionId}`;
+  }
+  timelineKey() {
+    return `${this.keyPrefix}:events:timeline`;
+  }
+  async storeEvent(event) {
+    const key = this.eventKey(event.timestamp, event.eventId);
+    await this.redis.setex(key, this.ttl, JSON.stringify(event));
+    await this.redis.zadd(this.timelineKey(), {
+      score: Date.parse(event.timestamp),
+      member: key
+    });
+    const count = await this.redis.zcard(this.timelineKey());
+    if (count && count > 1e3) {
+      await this.redis.zremrangebyrank(this.timelineKey(), 0, -1001);
+    }
+  }
+  async getRecentEvents(limit = 100) {
+    const keys = await this.redis.zrevrange(this.timelineKey(), 0, limit - 1);
+    if (!keys || keys.length === 0) {
+      return [];
+    }
+    const events = [];
+    for (const key of keys) {
+      const data = await this.redis.get(key);
+      if (data) {
+        try {
+          const event = typeof data === "string" ? JSON.parse(data) : data;
+          events.push(event);
+        } catch (e) {
+          console.error(`Failed to parse event ${key}:`, e);
+        }
+      }
+    }
+    return events;
+  }
+  async getSessionEvents(sessionId) {
+    const keys = await this.redis.zrevrange(this.timelineKey(), 0, -1);
+    if (!keys || keys.length === 0) {
+      return [];
+    }
+    const events = [];
+    for (const key of keys) {
+      const data = await this.redis.get(key);
+      if (data) {
+        try {
+          const event = typeof data === "string" ? JSON.parse(data) : data;
+          if (event.sessionId === sessionId) {
+            events.push(event);
+          }
+        } catch (e) {
+          console.error(`Failed to parse event ${key}:`, e);
+        }
+      }
+    }
+    return events;
+  }
+  async storeSession(session) {
+    const key = this.sessionKey(session.sessionId);
+    const existing = await this.redis.get(key);
+    if (existing) {
+      const existingSession = typeof existing === "string" ? JSON.parse(existing) : existing;
+      const methods = /* @__PURE__ */ new Set([
+        ...existingSession.verificationMethods || [],
+        ...session.verificationMethods
+      ]);
+      const updatedSession = {
+        ...existingSession,
+        lastSeen: session.lastSeen,
+        eventCount: session.eventCount,
+        paths: Array.from(/* @__PURE__ */ new Set([...existingSession.paths, ...session.paths])),
+        averageConfidence: session.averageConfidence,
+        verificationMethods: Array.from(methods)
+      };
+      await this.redis.setex(key, this.ttl, JSON.stringify(updatedSession));
+    } else {
+      await this.redis.setex(key, this.ttl, JSON.stringify(session));
+    }
+  }
+  async getSession(sessionId) {
+    const key = this.sessionKey(sessionId);
+    const data = await this.redis.get(key);
+    if (!data) {
+      return null;
+    }
+    try {
+      return typeof data === "string" ? JSON.parse(data) : data;
+    } catch (e) {
+      console.error(`Failed to parse session ${sessionId}:`, e);
+      return null;
+    }
+  }
+  async getRecentSessions(limit = 10) {
+    const pattern = `${this.keyPrefix}:sessions:*`;
+    const sessions = [];
+    let cursor = 0;
+    do {
+      const [nextCursor, keys] = await this.redis.scan(cursor, {
+        match: pattern,
+        count: 100
+      });
+      cursor = parseInt(nextCursor);
+      for (const key of keys) {
+        const data = await this.redis.get(key);
+        if (data) {
+          try {
+            const session = typeof data === "string" ? JSON.parse(data) : data;
+            sessions.push(session);
+          } catch (e) {
+            console.error(`Failed to parse session from ${key}:`, e);
+          }
+        }
+      }
+    } while (cursor !== 0 && sessions.length < limit * 2);
+    return sessions.sort((a, b) => Date.parse(b.lastSeen) - Date.parse(a.lastSeen)).slice(0, limit);
+  }
+  async cleanup(olderThan) {
+    const cutoff = olderThan.getTime();
+    await this.redis.zremrangebyrank(
+      this.timelineKey(),
+      0,
+      Math.floor(cutoff / 1e3)
+    );
+  }
+};
+
+// src/storage/index.ts
+async function createStorageAdapter(config) {
+  if (!config || config.type === "memory") {
+    return new MemoryStorageAdapter();
+  }
+  if (config.type === "custom" && config.custom) {
+    return config.custom;
+  }
+  if (config.type === "redis" && config.redis) {
+    try {
+      const { Redis } = await import('@upstash/redis');
+      const redis = new Redis({
+        url: config.redis.url,
+        token: config.redis.token
+      });
+      return new RedisStorageAdapter(redis, config.ttl);
+    } catch (error) {
+      console.warn("[AgentShield] Failed to initialize Redis storage, falling back to memory:", error);
+      return new MemoryStorageAdapter();
+    }
+  }
+  return new MemoryStorageAdapter();
+}
+
+// src/enhanced-middleware.ts
+var SessionManager = class {
+  sessionLastActivity = /* @__PURE__ */ new Map();
+  generateSessionId(ipAddress, userAgent) {
+    const now = Date.now();
+    const timeWindow = Math.floor(now / (5 * 60 * 1e3));
+    const baseKey = `${ipAddress || "unknown"}:${userAgent || "unknown"}`;
+    const windowKey = `${baseKey}:${timeWindow}`;
+    const lastActivity = this.sessionLastActivity.get(windowKey);
+    const shouldCreateNewSession = !lastActivity || now - lastActivity > 3e4;
+    this.sessionLastActivity.set(windowKey, now);
+    if (this.sessionLastActivity.size > 100) {
+      const cutoff = now - 6e5;
+      for (const [key, time] of this.sessionLastActivity.entries()) {
+        if (time < cutoff) {
+          this.sessionLastActivity.delete(key);
+        }
+      }
+    }
+    const data = shouldCreateNewSession ? `${windowKey}:${now}` : `${windowKey}:${lastActivity}`;
+    let hash = 0;
+    for (let i = 0; i < data.length; i++) {
+      const char = data.charCodeAt(i);
+      hash = (hash << 5) - hash + char;
+      hash = hash & hash;
+    }
+    return Math.abs(hash).toString(16).padStart(12, "0").substring(0, 12);
+  }
+};
+function createEnhancedAgentShieldMiddleware(config = {}) {
+  let storageAdapter = null;
+  let storageInitPromise = null;
+  const getStorage = async () => {
+    if (storageAdapter) return storageAdapter;
+    if (storageInitPromise) return storageInitPromise;
+    storageInitPromise = createStorageAdapter(config.storage).then((adapter) => {
+      storageAdapter = adapter;
+      return adapter;
+    });
+    return storageInitPromise;
+  };
+  let wasmAvailable = false;
+  checkWasmAvailability().then((available) => {
+    wasmAvailable = available;
+    if (available) {
+      console.log("[AgentShield] \u2705 WASM support detected - enhanced detection enabled");
+    }
+  });
+  const detector = wasmAvailable ? new EdgeAgentDetectorWrapperWithWasm({ enableWasm: true }) : new EdgeAgentDetectorWrapper({});
+  const sessionManager = new SessionManager();
+  const sessionTrackingEnabled = config.sessionTracking?.enabled !== false;
+  return async (request) => {
+    const { pathname } = request.nextUrl;
+    if (config.skipPaths?.some((path) => pathname.startsWith(path))) {
+      return server.NextResponse.next();
+    }
+    const userAgent = request.headers.get("user-agent");
+    const ipAddress = request.ip ?? request.headers.get("x-forwarded-for");
+    const url = new URL(request.url);
+    const pathWithQuery = url.pathname + url.search;
+    const context = {
+      userAgent: userAgent || "",
+      ipAddress: ipAddress || "",
+      headers: Object.fromEntries(request.headers.entries()),
+      url: pathWithQuery,
+      method: request.method,
+      timestamp: /* @__PURE__ */ new Date()
+    };
+    const result = await detector.analyze(context);
+    let finalConfidence = result.confidence;
+    let verificationMethod = result.verificationMethod || "pattern";
+    if (result.isAgent) {
+      const reasons = result.reasons || [];
+      if (shouldIndicateWasmVerification(result.confidence)) {
+        finalConfidence = getWasmConfidenceBoost(result.confidence, reasons);
+        verificationMethod = getVerificationMethod(finalConfidence, reasons);
+      }
+    }
+    if (result.isAgent && finalConfidence >= (config.confidenceThreshold ?? 0.7)) {
+      if (sessionTrackingEnabled) {
+        const storage = await getStorage();
+        const sessionId = sessionManager.generateSessionId(ipAddress || void 0, userAgent || void 0);
+        const event = {
+          eventId: `agent_${Date.now()}_${Math.random().toString(36).substring(2, 11)}`,
+          sessionId,
+          timestamp: (/* @__PURE__ */ new Date()).toISOString(),
+          agentType: result.detectedAgent?.type || "unknown",
+          agentName: result.detectedAgent?.name || "Unknown",
+          confidence: finalConfidence,
+          path: pathWithQuery,
+          ...userAgent && { userAgent },
+          ...ipAddress && { ipAddress },
+          method: request.method,
+          detectionReasons: result.reasons || [],
+          verificationMethod,
+          detectionDetails: {
+            patterns: result.detectedAgent?.patterns,
+            behaviors: result.detectedAgent?.behaviors,
+            fingerprintMatches: result.detectedAgent?.fingerprints
+          }
+        };
+        try {
+          await storage.storeEvent(event);
+          let session = await storage.getSession(sessionId);
+          if (session) {
+            session.lastSeen = event.timestamp;
+            session.eventCount++;
+            if (!session.paths.includes(pathWithQuery)) {
+              session.paths.push(pathWithQuery);
+            }
+            session.averageConfidence = (session.averageConfidence * (session.eventCount - 1) + finalConfidence) / session.eventCount;
+            if (!session.verificationMethods.includes(verificationMethod)) {
+              session.verificationMethods.push(verificationMethod);
+            }
+          } else {
+            session = {
+              sessionId,
+              ...ipAddress && { ipAddress },
+              ...userAgent && { userAgent },
+              agentType: result.detectedAgent?.type || "unknown",
+              agentName: result.detectedAgent?.name || "Unknown",
+              firstSeen: event.timestamp,
+              lastSeen: event.timestamp,
+              eventCount: 1,
+              paths: [pathWithQuery],
+              averageConfidence: finalConfidence,
+              verificationMethods: [verificationMethod]
+            };
+          }
+          if (session) {
+            await storage.storeSession(session);
+          }
+        } catch (error) {
+          console.error("[AgentShield] Failed to store event:", error);
+        }
+      }
+      if (config.onDetection) {
+        await config.onDetection(result, context);
+      }
+      switch (config.onAgentDetected) {
+        case "block": {
+          const { status = 403, message = "Access denied: AI agent detected" } = config.blockedResponse || {};
+          const response2 = server.NextResponse.json(
+            { error: message, detected: true, confidence: finalConfidence },
+            { status }
+          );
+          response2.headers.set("x-agentshield-detected", "true");
+          response2.headers.set("x-agentshield-confidence", String(Math.round(finalConfidence * 100)));
+          response2.headers.set("x-agentshield-agent", result.detectedAgent?.name || "Unknown");
+          response2.headers.set("x-agentshield-verification", verificationMethod);
+          return response2;
+        }
+        case "log":
+          console.log(`[AgentShield] \u{1F916} AI Agent detected (${verificationMethod}):`, {
+            agent: result.detectedAgent?.name,
+            confidence: `${(finalConfidence * 100).toFixed(0)}%`,
+            path: pathWithQuery,
+            verification: verificationMethod
+          });
+          break;
+      }
+    }
+    const response = server.NextResponse.next();
+    if (result.isAgent) {
+      response.headers.set("x-agentshield-detected", "true");
+      response.headers.set("x-agentshield-confidence", String(Math.round(finalConfidence * 100)));
+      response.headers.set("x-agentshield-agent", result.detectedAgent?.name || "Unknown");
+      response.headers.set("x-agentshield-verification", verificationMethod);
+      if (finalConfidence > 0.9) {
+        response.headers.set("x-ai-visitor", "true");
+        response.headers.set("x-ai-confidence", finalConfidence.toString());
+        response.headers.set("x-ai-verification", verificationMethod);
+      }
+    }
+    return response;
+  };
+}
+
 // src/index.ts
 var VERSION = "0.1.0";
 /**
@@ -1070,6 +1530,7 @@ exports.StatelessSessionChecker = StatelessSessionChecker;
 exports.VERSION = VERSION;
 exports.createAgentShieldMiddleware = createAgentShieldMiddleware2;
 exports.createAgentShieldMiddlewareBase = createAgentShieldMiddleware;
+exports.createEnhancedAgentShieldMiddleware = createEnhancedAgentShieldMiddleware;
 exports.createMiddleware = createAgentShieldMiddleware2;
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map