@kya-os/agentshield-nextjs/bin/setup-edge-wasm.js

Next.js middleware for AgentShield AI agent detection

#!/usr/bin/env node

/**
 * AgentShield Edge Runtime WASM Setup Script
 *
 * This script helps Next.js users set up WASM support for Edge Runtime (middleware).
 * It copies necessary files and creates the required configuration.
 */

const fs = require('fs');
const path = require('path');
const { execSync } = require('child_process');

// Colors for console output
const colors = {
  reset: '\x1b[0m',
  bright: '\x1b[1m',
  green: '\x1b[32m',
  yellow: '\x1b[33m',
  blue: '\x1b[34m',
  red: '\x1b[31m',
};

function log(message, color = colors.reset) {
  console.log(`${color}${message}${colors.reset}`);
}

function error(message) {
  console.error(`${colors.red}❌ ${message}${colors.reset}`);
}

function success(message) {
  log(`✅ ${message}`, colors.green);
}

function info(message) {
  log(`ℹ️  ${message}`, colors.blue);
}

function warning(message) {
  log(`⚠️  ${message}`, colors.yellow);
}

async function main() {
  log('\n🚀 AgentShield Edge Runtime WASM Setup\n', colors.bright);

  // Step 1: Check if we're in a Next.js project
  const packageJsonPath = path.join(process.cwd(), 'package.json');
  if (!fs.existsSync(packageJsonPath)) {
    error(
      'No package.json found. Please run this command from your Next.js project root.'
    );
    process.exit(1);
  }

  const packageJson = JSON.parse(fs.readFileSync(packageJsonPath, 'utf8'));
  const hasNext =
    packageJson.dependencies?.next || packageJson.devDependencies?.next;

  if (!hasNext) {
    warning('Next.js not detected in package.json. Continuing anyway...');
  }

  // Step 2: Check if AgentShield is installed
  const agentShieldPath = path.join(
    process.cwd(),
    'node_modules',
    '@kya-os',
    'agentshield-nextjs'
  );
  if (!fs.existsSync(agentShieldPath)) {
    error('@kya-os/agentshield-nextjs not found in node_modules.');
    info('Please install it first:');
    console.log('  npm install @kya-os/agentshield-nextjs');
    console.log('  # or');
    console.log('  pnpm add @kya-os/agentshield-nextjs');
    process.exit(1);
  }

  // Step 3: Create wasm directory
  const wasmDir = path.join(process.cwd(), 'wasm');
  if (!fs.existsSync(wasmDir)) {
    fs.mkdirSync(wasmDir, { recursive: true });
    success('Created wasm/ directory');
  }

  // Step 4: Copy WASM file
  const sourceWasm = path.join(
    agentShieldPath,
    'dist',
    'wasm',
    'agentshield_wasm_bg.wasm'
  );
  const destWasm = path.join(wasmDir, 'agentshield_wasm_bg.wasm');

  if (fs.existsSync(sourceWasm)) {
    fs.copyFileSync(sourceWasm, destWasm);
    const sizeKB = (fs.statSync(destWasm).size / 1024).toFixed(1);
    const sizeMB = (fs.statSync(destWasm).size / 1024 / 1024).toFixed(2);
    const sizeStr = sizeKB < 1024 ? `${sizeKB}KB` : `${sizeMB}MB`;
    success(`Copied WASM file to wasm/ (${sizeStr})`);
  } else {
    // Try alternative location
    const altSourceWasm = path.join(
      agentShieldPath,
      'wasm',
      'agentshield_wasm_bg.wasm'
    );
    if (fs.existsSync(altSourceWasm)) {
      fs.copyFileSync(altSourceWasm, destWasm);
      const sizeKB = (fs.statSync(destWasm).size / 1024).toFixed(1);
      const sizeMB = (fs.statSync(destWasm).size / 1024 / 1024).toFixed(2);
      const sizeStr = sizeKB < 1024 ? `${sizeKB}KB` : `${sizeMB}MB`;
      success(`Copied WASM file to wasm/ (${sizeStr})`);
    } else {
      error(
        'WASM file not found in the package. The package may be corrupted.'
      );
      process.exit(1);
    }
  }

  // Step 5: Create TypeScript definitions
  const dtsContent = `// Auto-generated by agentshield-setup-edge
// TypeScript definitions for AgentShield WASM in Edge Runtime

declare module '*.wasm?module' {
  const wasmModule: WebAssembly.Module;
  export default wasmModule;
}

export interface AgentShieldWasm {
  detect_agent(metadata: string): string;
  get_version(): string;
  memory: WebAssembly.Memory;
}

export interface DetectionResult {
  isAgent: boolean;
  confidence: number;
  agent?: string;
  verificationMethod: 'cryptographic' | 'pattern';
  riskLevel?: 'low' | 'medium' | 'high' | 'critical';
  timestamp: string;
}
`;

  const dtsPath = path.join(wasmDir, 'agentshield.d.ts');
  fs.writeFileSync(dtsPath, dtsContent);
  success('Created TypeScript definitions');

  // Step 6: Create lib directory and Edge loader
  const libDir = path.join(process.cwd(), 'lib');
  if (!fs.existsSync(libDir)) {
    fs.mkdirSync(libDir, { recursive: true });
  }

  const edgeLoaderContent = `/**
 * AgentShield Edge Runtime Loader
 * Auto-generated by agentshield-setup-edge
 * 
 * This module provides WASM-based AI agent detection for Next.js Edge Runtime.
 */

import type { NextRequest } from 'next/server';
import wasmModule from '../wasm/agentshield_wasm_bg.wasm?module';

let wasmInstance: WebAssembly.Instance | null = null;
let wasmMemory: WebAssembly.Memory | null = null;

interface WasmExports {
  detect_agent(metadataPtr: number, metadataLen: number): number;
  get_version(): number;
  __wbindgen_malloc(size: number): number;
  __wbindgen_free(ptr: number, size: number): void;
  __wbindgen_realloc(ptr: number, oldSize: number, newSize: number): number;
  memory: WebAssembly.Memory;
}

async function initWasm(): Promise<void> {
  if (wasmInstance) return;
  
  wasmMemory = new WebAssembly.Memory({ initial: 17, maximum: 256 });
  
  const imports = {
    wbg: {
      __wbindgen_throw: (ptr: number, len: number) => {
        throw new Error(readString(ptr, len));
      },
      __wbg_log_: (ptr: number, len: number) => {
        console.log(readString(ptr, len));
      },
    },
    env: {
      memory: wasmMemory,
    },
  };
  
  wasmInstance = await WebAssembly.instantiate(wasmModule, imports);
}

function readString(ptr: number, len: number): string {
  if (!wasmInstance || !wasmMemory) throw new Error('WASM not initialized');
  const memory = new Uint8Array(wasmMemory.buffer);
  const bytes = memory.slice(ptr, ptr + len);
  return new TextDecoder().decode(bytes);
}

function writeString(str: string): [number, number] {
  if (!wasmInstance) throw new Error('WASM not initialized');
  const exports = wasmInstance.exports as unknown as WasmExports;
  
  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];
}

export interface DetectionResult {
  isAgent: boolean;
  confidence: number;
  agent?: string;
  verificationMethod: 'cryptographic' | 'pattern';
  riskLevel?: 'low' | 'medium' | 'high' | 'critical';
  timestamp?: string;
}

export interface AgentShieldOptions {
  enableWasm?: boolean;
  onAgentDetected?: (result: DetectionResult) => void;
  blockAgents?: boolean;
  allowedAgents?: string[];
}

/**
 * Create an AgentShield middleware instance
 */
export function createAgentShieldMiddleware(options: AgentShieldOptions = {}) {
  let initialized = false;
  
  return {
    async init(): Promise<void> {
      if (initialized) return;
      
      if (options.enableWasm !== false) {
        try {
          await initWasm();
          initialized = true;
          console.log('✅ AgentShield WASM initialized in Edge Runtime');
        } catch (error) {
          console.warn('⚠️ WASM initialization failed, using pattern detection:', error);
          initialized = true;
        }
      } else {
        initialized = true;
      }
    },
    
    async detect(request: NextRequest): Promise<DetectionResult> {
      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()),
      };
      
      // Try WASM detection first
      if (wasmInstance && options.enableWasm !== false) {
        try {
          const exports = wasmInstance.exports as unknown as WasmExports;
          const [ptr, len] = writeString(JSON.stringify(metadata));
          
          const resultPtr = exports.detect_agent(ptr, len);
          const resultLen = 1024; // Assume max result size
          const resultStr = readString(resultPtr, resultLen);
          
          exports.__wbindgen_free(ptr, len);
          exports.__wbindgen_free(resultPtr, resultLen);
          
          const result = JSON.parse(resultStr);
          
          const detection: DetectionResult = {
            ...result,
            verificationMethod: 'cryptographic',
            timestamp: new Date().toISOString(),
          };
          
          if (options.onAgentDetected && detection.isAgent) {
            options.onAgentDetected(detection);
          }
          
          return detection;
        } catch (error) {
          console.error('WASM detection failed:', error);
          // Fall through to pattern detection
        }
      }
      
      // Fallback: Pattern-based detection
      return patternDetection(metadata, options);
    },
    
    isInitialized(): boolean {
      return initialized;
    },
    
    getVerificationMethod(): 'cryptographic' | 'pattern' {
      return wasmInstance ? 'cryptographic' : 'pattern';
    },
  };
}

/**
 * Pattern-based detection fallback
 */
function patternDetection(
  metadata: { userAgent: string; ipAddress: string; headers: Record<string, string> },
  options: AgentShieldOptions
): DetectionResult {
  const userAgent = metadata.userAgent.toLowerCase();
  
  // Known AI agent patterns with confidence scores
  const aiAgentPatterns = [
    { pattern: /chatgpt-user/i, name: 'ChatGPT', confidence: 0.95 },
    { pattern: /claude-web/i, name: 'Claude', confidence: 0.95 },
    { pattern: /anthropic/i, name: 'Anthropic', confidence: 0.90 },
    { pattern: /openai/i, name: 'OpenAI', confidence: 0.90 },
    { pattern: /gpt-crawler/i, name: 'GPT Crawler', confidence: 0.95 },
    { 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 },
  ];
  
  // Check headers for AI-specific patterns
  const suspiciousHeaders = [
    'x-openai-',
    'x-anthropic-',
    'x-ai-',
    'x-llm-',
  ];
  
  let headerBoost = 0;
  for (const [key, value] of Object.entries(metadata.headers)) {
    if (suspiciousHeaders.some(prefix => key.toLowerCase().startsWith(prefix))) {
      headerBoost = 0.1;
      break;
    }
  }
  
  // Check user agent patterns
  for (const { pattern, name, confidence } of aiAgentPatterns) {
    if (pattern.test(userAgent)) {
      const finalConfidence = Math.min(confidence + headerBoost, 1.0);
      
      const result: DetectionResult = {
        isAgent: true,
        confidence: finalConfidence,
        agent: name,
        verificationMethod: 'pattern',
        riskLevel: finalConfidence > 0.9 ? 'high' : 'medium',
        timestamp: new Date().toISOString(),
      };
      
      if (options.onAgentDetected) {
        options.onAgentDetected(result);
      }
      
      return result;
    }
  }
  
  // Not detected as AI agent
  return {
    isAgent: false,
    confidence: 0.85,
    verificationMethod: 'pattern',
    timestamp: new Date().toISOString(),
  };
}

export default createAgentShieldMiddleware;
`;

  const edgeLoaderPath = path.join(libDir, 'agentshield-edge.ts');
  fs.writeFileSync(edgeLoaderPath, edgeLoaderContent);
  success('Created Edge Runtime loader at lib/agentshield-edge.ts');

  // Step 7: Create example middleware if it doesn't exist
  const middlewarePath = path.join(process.cwd(), 'middleware.ts');
  if (!fs.existsSync(middlewarePath)) {
    const middlewareContent = `import { NextResponse } from 'next/server';
import type { NextRequest } from 'next/server';
import { createAgentShieldMiddleware } from './lib/agentshield-edge';

// Initialize AgentShield
const agentShield = createAgentShieldMiddleware({
  enableWasm: true,
  onAgentDetected: (result) => {
    console.log(\`🤖 AI Agent detected: \${result.agent} (confidence: \${(result.confidence * 100).toFixed(1)}%)\`);
  },
});

export async function middleware(request: NextRequest) {
  // Initialize WASM on first request
  await agentShield.init();
  
  // Detect AI agents
  const detection = await agentShield.detect(request);
  
  // Log detection results
  if (detection.isAgent) {
    console.log(\`AI Agent detected on \${request.url}:\`, detection);
    
    // Optional: Add headers to track AI agent traffic
    const response = NextResponse.next();
    response.headers.set('X-Agent-Detected', detection.agent || 'unknown');
    response.headers.set('X-Agent-Confidence', detection.confidence.toString());
    return response;
    
    // Optional: Block AI agents
    // if (detection.confidence > 0.9) {
    //   return NextResponse.json(
    //     { error: 'AI agent access denied' },
    //     { status: 403 }
    //   );
    // }
  }
  
  return NextResponse.next();
}

export const config = {
  matcher: [
    /*
     * Match all request paths except for the ones starting with:
     * - api/admin (protect admin APIs)
     * - _next/static (static files)
     * - _next/image (image optimization files)
     * - favicon.ico (favicon file)
     */
    '/((?!_next/static|_next/image|favicon.ico).*)',
  ],
};
`;

    fs.writeFileSync(middlewarePath, middlewareContent);
    success('Created example middleware.ts');
  } else {
    info('middleware.ts already exists - skipping creation');
    info('See lib/agentshield-edge.ts for integration instructions');
  }

  // Step 8: Update next.config.js if needed
  const nextConfigPath = path.join(process.cwd(), 'next.config.js');
  const nextConfigMjsPath = path.join(process.cwd(), 'next.config.mjs');
  const configPath = fs.existsSync(nextConfigPath)
    ? nextConfigPath
    : fs.existsSync(nextConfigMjsPath)
      ? nextConfigMjsPath
      : null;

  if (configPath) {
    const configContent = fs.readFileSync(configPath, 'utf8');

    if (!configContent.includes('wasm')) {
      warning(
        'Your next.config.js may need to be updated to handle WASM files.'
      );
      console.log('\nAdd this to your webpack config:\n');
      console.log(
        colors.yellow +
          `  webpack: (config) => {
    config.module.rules.push({
      test: /\\.wasm$/,
      type: 'asset/resource',
    });
    return config;
  },` +
          colors.reset
      );
    }
  }

  // Step 9: Final instructions
  console.log('\n' + colors.bright + '🎉 Setup Complete!' + colors.reset);
  console.log(
    '\n' +
      colors.green +
      'AgentShield Edge Runtime WASM is now configured!' +
      colors.reset
  );
  console.log('\nNext steps:');
  console.log('1. Review the generated files:');
  console.log('   - wasm/agentshield_wasm_bg.wasm (WASM binary)');
  console.log('   - lib/agentshield-edge.ts (Edge Runtime loader)');
  console.log('   - middleware.ts (Example implementation)');
  console.log('\n2. Test your setup:');
  console.log('   npm run dev');
  console.log('   # Visit your site and check the console for detection logs');
  console.log('\n3. Deploy to Vercel:');
  console.log('   git add .');
  console.log('   git commit -m "Add AgentShield WASM for Edge Runtime"');
  console.log('   git push');
  console.log(
    '\n' +
      colors.blue +
      'Documentation: https://github.com/kya-os/agentshield#edge-runtime' +
      colors.reset
  );
  console.log(
    colors.blue +
      'Issues: https://github.com/kya-os/agentshield/issues' +
      colors.reset
  );
}

// Run the setup
main().catch(err => {
  error('Setup failed: ' + err.message);
  console.error(err);
  process.exit(1);
});