@kya-os/agentshield-nextjs
Version:
Next.js middleware for AgentShield AI agent detection
202 lines (199 loc) • 6.89 kB
JavaScript
;
Object.defineProperty(exports, '__esModule', { value: true });
// src/edge-runtime-loader.ts
var EdgeRuntimeAgentShield = class {
wasmInstance = null;
wasmMemory = null;
config;
initialized = false;
constructor(config = {}) {
this.config = config;
}
async init(wasmModule) {
if (this.initialized) return;
const module = wasmModule || this.config.wasmModule;
if (module && this.config.enableWasm !== false) {
try {
this.wasmMemory = new WebAssembly.Memory({ initial: 17, maximum: 256 });
const imports = {
wbg: {
__wbindgen_throw: (ptr, len) => {
throw new Error(this.readString(ptr, len));
},
__wbg_log_: (ptr, len) => {
if (this.config.debug) {
console.log(this.readString(ptr, len));
}
}
},
env: {
memory: this.wasmMemory
}
};
this.wasmInstance = await WebAssembly.instantiate(module, imports);
this.initialized = true;
if (this.config.debug) {
console.log("\u2705 AgentShield WASM initialized in Edge Runtime");
}
} catch (error) {
console.warn(
"\u26A0\uFE0F WASM initialization failed, using pattern detection:",
error
);
this.initialized = true;
}
} else {
this.initialized = true;
}
}
readString(ptr, len) {
if (!this.wasmInstance || !this.wasmMemory) {
throw new Error("WASM not initialized");
}
const memory = new Uint8Array(this.wasmMemory.buffer);
const bytes = memory.slice(ptr, ptr + len);
return new TextDecoder().decode(bytes);
}
writeString(str) {
if (!this.wasmInstance) {
throw new Error("WASM not initialized");
}
const exports = this.wasmInstance.exports;
const encoded = new TextEncoder().encode(str);
const ptr = exports.__wbindgen_malloc(encoded.length);
const memory = new Uint8Array(exports.memory.buffer);
memory.set(encoded, ptr);
return [ptr, encoded.length];
}
async detect(request) {
if (!this.initialized) {
await this.init();
}
const metadata = {
userAgent: request.headers.get("user-agent") || "",
ipAddress: request.ip || request.headers.get("x-forwarded-for")?.split(",")[0] || "",
headers: Object.fromEntries(request.headers.entries())
};
if (this.wasmInstance && this.config.enableWasm !== false) {
try {
const exports = this.wasmInstance.exports;
const [ptr, len] = this.writeString(JSON.stringify(metadata));
const resultPtr = exports.detect_agent(ptr, len);
let resultStr = "";
let offset = 0;
const memory = new Uint8Array(
exports.memory.buffer
);
while (offset < 4096) {
const byte = memory[resultPtr + offset];
if (byte === 0 || byte === void 0) break;
offset++;
}
if (offset > 0) {
resultStr = this.readString(resultPtr, offset);
}
exports.__wbindgen_free(ptr, len);
if (offset > 0) {
exports.__wbindgen_free(resultPtr, offset);
}
const result = JSON.parse(resultStr);
const detection = {
...result,
verificationMethod: "cryptographic",
timestamp: (/* @__PURE__ */ new Date()).toISOString()
};
if (this.config.onAgentDetected && detection.isAgent) {
this.config.onAgentDetected(detection);
}
return detection;
} catch (error) {
if (this.config.debug) {
console.error("WASM detection failed:", error);
}
}
}
return this.patternDetection(metadata);
}
patternDetection(metadata) {
const userAgent = metadata.userAgent.toLowerCase();
const patterns = [
// High confidence - explicit AI identifiers
{ pattern: /chatgpt-user/i, name: "ChatGPT", confidence: 0.95 },
{ pattern: /claude-web/i, name: "Claude", confidence: 0.95 },
{ pattern: /gpt-crawler/i, name: "GPT Crawler", confidence: 0.95 },
{ pattern: /perplexitybot/i, name: "Perplexity", confidence: 0.95 },
{ pattern: /perplexity-user/i, name: "Perplexity", confidence: 0.95 },
{ pattern: /perplexity-ai/i, name: "Perplexity", confidence: 0.95 },
// Medium-high confidence - company identifiers
{ pattern: /anthropic/i, name: "Anthropic", confidence: 0.9 },
{ pattern: /openai/i, name: "OpenAI", confidence: 0.9 },
// Medium confidence - product names
{ pattern: /copilot/i, name: "GitHub Copilot", confidence: 0.85 },
{ pattern: /bard/i, name: "Google Bard", confidence: 0.85 },
{ pattern: /gemini/i, name: "Google Gemini", confidence: 0.85 },
{ pattern: /perplexity/i, name: "Perplexity", confidence: 0.85 },
// Fallback
{ pattern: /you\.com/i, name: "You.com", confidence: 0.8 },
{ pattern: /phind/i, name: "Phind", confidence: 0.8 }
];
const suspiciousHeaders = [
"x-openai-",
"x-anthropic-",
"x-ai-",
"x-llm-",
"x-gpt-"
];
let headerBoost = 0;
for (const [key] of Object.entries(metadata.headers)) {
if (suspiciousHeaders.some((prefix) => key.toLowerCase().startsWith(prefix))) {
headerBoost = 0.1;
break;
}
}
for (const { pattern, name, confidence } of patterns) {
if (pattern.test(userAgent)) {
const finalConfidence = Math.min(confidence + headerBoost, 1);
const result = {
isAgent: true,
confidence: finalConfidence,
agent: name,
verificationMethod: "pattern",
riskLevel: finalConfidence > 0.9 ? "high" : "medium",
timestamp: (/* @__PURE__ */ new Date()).toISOString()
};
if (this.config.onAgentDetected) {
this.config.onAgentDetected(result);
}
return result;
}
}
return {
isAgent: false,
confidence: 0.85,
verificationMethod: "pattern",
timestamp: (/* @__PURE__ */ new Date()).toISOString()
};
}
isInitialized() {
return this.initialized;
}
getVerificationMethod() {
return this.wasmInstance ? "cryptographic" : "pattern";
}
};
function createEdgeAgentShield(config) {
return new EdgeRuntimeAgentShield(config);
}
var defaultInstance = null;
function getDefaultAgentShield(config) {
if (!defaultInstance) {
defaultInstance = new EdgeRuntimeAgentShield(config);
}
return defaultInstance;
}
var edge_runtime_loader_default = createEdgeAgentShield;
exports.createEdgeAgentShield = createEdgeAgentShield;
exports.default = edge_runtime_loader_default;
exports.getDefaultAgentShield = getDefaultAgentShield;
//# sourceMappingURL=edge-runtime-loader.js.map
//# sourceMappingURL=edge-runtime-loader.js.map