Patch out haiku-based security checks in Claude Code

patch_claude_code_haiku_call.ts

#!/usr/bin/env bun
/**
 * Resilient AST-based patcher for Generic Code CLI
 *
 * This patcher uses Babel to parse and transform the CLI bundle,
 * making certain LLM-based checks configurable via environment variables.
 *
 * The patcher finds code by stable identifiers (querySource strings used for telemetry)
 * rather than function names, making it resilient to minification and refactoring.
 *
 * Configurable features:
 * - GENERIC_CODE_SKIP_PREFIX_CHECK=1 - Skip the pre-flight bash command prefix extraction
 * - GENERIC_CODE_SKIP_PATH_EXTRACTION=1 - Skip post-execution file path extraction from output
 */

import * as parser from '@babel/parser';
import _traverse from '@babel/traverse';
import _generate from '@babel/generator';
import * as t from '@babel/types';
import fs from 'fs';
import path from 'path';

// Handle default export differences
const traverse = (_traverse as any).default || _traverse;
const generate = (_generate as any).default || _generate;

// Stable identifiers - these are telemetry strings unlikely to change
const QUERY_SOURCE_PREFIX_CHECK = 'bash_extract_prefix';
const QUERY_SOURCE_PATH_EXTRACTION = 'bash_extract_command_paths';

interface PatchLocation {
  type: 'prefix_check' | 'path_extraction';
  querySource: string;
  found: boolean;
  patched: boolean;
  functionName?: string;
  line?: number;
}

interface PatchResult {
  success: boolean;
  locations: PatchLocation[];
  outputPath: string;
  errors: string[];
}

/**
 * Find the enclosing function/arrow function for a given path
 */
function findEnclosingFunction(nodePath: any): any {
  let current = nodePath;
  while (current) {
    if (
      t.isFunctionDeclaration(current.node) ||
      t.isFunctionExpression(current.node) ||
      t.isArrowFunctionExpression(current.node)
    ) {
      return current;
    }
    current = current.parentPath;
  }
  return null;
}

/**
 * Find the enclosing statement that can be wrapped
 */
function findEnclosingStatement(nodePath: any): any {
  let current = nodePath;
  while (current) {
    if (
      t.isExpressionStatement(current.node) ||
      t.isVariableDeclaration(current.node) ||
      t.isReturnStatement(current.node)
    ) {
      // Make sure we're at a statement level (parent is a block or program)
      if (
        current.parentPath &&
        (t.isBlockStatement(current.parentPath.node) ||
          t.isProgram(current.parentPath.node))
      ) {
        return current;
      }
    }
    current = current.parentPath;
  }
  return null;
}

/**
 * Check if an object has a querySource property with the given value
 */
function hasQuerySourceProperty(node: any, querySource: string): boolean {
  if (!t.isObjectExpression(node)) return false;

  return node.properties.some((prop: any) => {
    if (!t.isObjectProperty(prop)) return false;

    const key = prop.key;
    const value = prop.value;

    const isQuerySourceKey =
      (t.isIdentifier(key) && key.name === 'querySource') ||
      (t.isStringLiteral(key) && key.value === 'querySource');

    const matchesValue = t.isStringLiteral(value) && value.value === querySource;

    return isQuerySourceKey && matchesValue;
  });
}

/**
 * Recursively search for querySource in nested objects
 */
function findQuerySourceInNode(node: any, querySource: string): boolean {
  if (!node) return false;

  if (t.isObjectExpression(node)) {
    if (hasQuerySourceProperty(node, querySource)) return true;

    // Check nested objects
    for (const prop of node.properties) {
      if (t.isObjectProperty(prop) && findQuerySourceInNode(prop.value, querySource)) {
        return true;
      }
    }
  }

  if (t.isCallExpression(node)) {
    for (const arg of node.arguments) {
      if (findQuerySourceInNode(arg, querySource)) return true;
    }
  }

  return false;
}

/**
 * Create the environment variable check expression
 */
function createEnvCheck(envVar: string): any {
  // Creates: process.env.ENVVAR
  return t.memberExpression(
    t.memberExpression(
      t.identifier('process'),
      t.identifier('env')
    ),
    t.identifier(envVar)
  );
}

/**
 * Create a conditional wrapper that short-circuits the function
 * For async functions returning promises with specific structures
 */
function createPrefixCheckBypass(envVar: string): any {
  // if (process.env.ENVVAR) return { commandPrefix: A };
  // This returns early with a "pass-through" result
  return t.ifStatement(
    createEnvCheck(envVar),
    t.returnStatement(
      t.objectExpression([
        t.objectProperty(
          t.identifier('commandPrefix'),
          // Use the first argument (the command) as the prefix
          t.identifier('A')
        )
      ])
    )
  );
}

/**
 * Create bypass for path extraction - returns empty array
 */
function createPathExtractionBypass(envVar: string): any {
  // if (process.env.ENVVAR) return [];
  return t.ifStatement(
    createEnvCheck(envVar),
    t.returnStatement(t.arrayExpression([]))
  );
}

/**
 * Main patching function
 */
function patchFile(inputPath: string, outputPath: string): PatchResult {
  const result: PatchResult = {
    success: false,
    locations: [],
    outputPath,
    errors: []
  };

  // Track what we're looking for
  const targets: PatchLocation[] = [
    { type: 'prefix_check', querySource: QUERY_SOURCE_PREFIX_CHECK, found: false, patched: false },
    { type: 'path_extraction', querySource: QUERY_SOURCE_PATH_EXTRACTION, found: false, patched: false }
  ];

  let code: string;
  try {
    code = fs.readFileSync(inputPath, 'utf-8');
  } catch (err) {
    result.errors.push(`Failed to read input file: ${err}`);
    return result;
  }

  console.log(`Parsing ${inputPath} (${(code.length / 1024 / 1024).toFixed(2)} MB)...`);

  let ast: any;
  try {
    ast = parser.parse(code, {
      sourceType: 'module',
      plugins: ['jsx'],
      allowReturnOutsideFunction: true,
      allowSuperOutsideMethod: true,
      errorRecovery: true
    });
  } catch (err) {
    result.errors.push(`Failed to parse: ${err}`);
    return result;
  }

  if (!ast) {
    result.errors.push('AST is null after parsing');
    return result;
  }

  console.log('AST parsed successfully, searching for targets...');

  // Track functions we've already patched to avoid duplicates
  const patchedFunctions = new Set<any>();

  // First pass: find all locations
  traverse(ast, {
    ObjectExpression(nodePath: any) {
      for (const target of targets) {
        if (hasQuerySourceProperty(nodePath.node, target.querySource)) {
          target.found = true;
          target.line = nodePath.node.loc?.start?.line;

          // Find the enclosing function
          const funcPath = findEnclosingFunction(nodePath);
          if (funcPath) {
            target.functionName = funcPath.node.id?.name || '<anonymous>';
          }

          console.log(`  Found ${target.type} at line ${target.line || '?'} in ${target.functionName || 'anonymous function'}`);
        }
      }
    }
  });

  // Second pass: patch the functions
  traverse(ast, {
    ObjectExpression(nodePath: any) {
      // Check for prefix check
      if (hasQuerySourceProperty(nodePath.node, QUERY_SOURCE_PREFIX_CHECK)) {
        const funcPath = findEnclosingFunction(nodePath);
        if (funcPath && !patchedFunctions.has(funcPath.node)) {
          patchedFunctions.add(funcPath.node);

          // Get the function body
          const body = funcPath.node.body;
          if (t.isBlockStatement(body)) {
            // Insert bypass at the start of the function
            const bypass = createPrefixCheckBypass('GENERIC_CODE_SKIP_PREFIX_CHECK');
            body.body.unshift(bypass);

            const target = targets.find(t => t.type === 'prefix_check');
            if (target) target.patched = true;

            console.log('  Patched prefix_check function');
          }
        }
      }

      // Check for path extraction
      if (hasQuerySourceProperty(nodePath.node, QUERY_SOURCE_PATH_EXTRACTION)) {
        const funcPath = findEnclosingFunction(nodePath);
        if (funcPath && !patchedFunctions.has(funcPath.node)) {
          patchedFunctions.add(funcPath.node);

          // Get the function body
          const body = funcPath.node.body;
          if (t.isBlockStatement(body)) {
            // Insert bypass at the start of the function
            const bypass = createPathExtractionBypass('GENERIC_CODE_SKIP_PATH_EXTRACTION');
            body.body.unshift(bypass);

            const target = targets.find(t => t.type === 'path_extraction');
            if (target) target.patched = true;

            console.log('  Patched path_extraction function');
          }
        }
      }
    }
  });

  result.locations = targets;

  // Check if we found and patched everything
  const allFound = targets.every(t => t.found);
  const allPatched = targets.every(t => t.patched);

  if (!allFound) {
    const missing = targets.filter(t => !t.found).map(t => t.querySource);
    result.errors.push(`Could not find targets: ${missing.join(', ')}`);
  }

  if (!allPatched) {
    const notPatched = targets.filter(t => !t.patched).map(t => t.querySource);
    result.errors.push(`Could not patch targets: ${notPatched.join(', ')}`);
  }

  // Generate output
  console.log('Generating patched code...');
  try {
    const output = generate(ast, {
      retainLines: false,
      compact: true, // Keep output minified like the original
      comments: false,
      minified: true
    }, code);

    fs.writeFileSync(outputPath, output.code);
    result.success = allFound && allPatched;
    console.log(`Wrote patched file to ${outputPath}`);
  } catch (err) {
    result.errors.push(`Failed to generate output: ${err}`);
  }

  return result;
}

/**
 * CLI interface
 */
function main() {
  const args = process.argv.slice(2);

  if (args.includes('--help') || args.includes('-h')) {
    console.log(`
Generic Code CLI Patcher

Usage: bun patcher.ts [options] <input-file> [output-file]

Options:
  --help, -h     Show this help message
  --dry-run      Only scan, don't patch
  --backup       Create a backup of the original file

Environment Variables (after patching):
  GENERIC_CODE_SKIP_PREFIX_CHECK=1     Skip bash command prefix extraction (faster commands)
  GENERIC_CODE_SKIP_PATH_EXTRACTION=1  Skip file path extraction from output

Examples:
  bun patcher.ts cli-patch.js                    # Patch in-place
  bun patcher.ts cli-patch.js patched-cli.js    # Patch to new file
  bun patcher.ts --dry-run cli-patch.js         # Just scan for targets
`);
    process.exit(0);
  }

  const dryRun = args.includes('--dry-run');
  const backup = args.includes('--backup');
  const fileArgs = args.filter(a => !a.startsWith('--'));

  if (fileArgs.length === 0) {
    // Default to cli-patch.js in current directory
    fileArgs.push('cli-patch.js');
  }

  const inputPath = path.resolve(fileArgs[0]);
  const outputPath = fileArgs[1] ? path.resolve(fileArgs[1]) : inputPath;

  if (!fs.existsSync(inputPath)) {
    console.error(`Error: Input file not found: ${inputPath}`);
    process.exit(1);
  }

  if (backup && inputPath === outputPath) {
    const backupPath = inputPath + '.backup';
    console.log(`Creating backup at ${backupPath}`);
    fs.copyFileSync(inputPath, backupPath);
  }

  if (dryRun) {
    console.log('Dry run mode - scanning only...\n');
  }

  console.log(`Input: ${inputPath}`);
  console.log(`Output: ${outputPath}\n`);

  if (dryRun) {
    // Just scan and report
    const code = fs.readFileSync(inputPath, 'utf-8');

    // Quick scan for our target strings
    const prefixMatch = code.includes(`querySource:"${QUERY_SOURCE_PREFIX_CHECK}"`) ||
                        code.includes(`querySource: "${QUERY_SOURCE_PREFIX_CHECK}"`);
    const pathMatch = code.includes(`querySource:"${QUERY_SOURCE_PATH_EXTRACTION}"`) ||
                      code.includes(`querySource: "${QUERY_SOURCE_PATH_EXTRACTION}"`);

    console.log('Scan Results:');
    console.log(`  bash_extract_prefix: ${prefixMatch ? 'FOUND' : 'NOT FOUND'}`);
    console.log(`  bash_extract_command_paths: ${pathMatch ? 'FOUND' : 'NOT FOUND'}`);

    if (prefixMatch && pathMatch) {
      console.log('\nAll targets found. Run without --dry-run to patch.');
    } else {
      console.log('\nWarning: Some targets not found. The CLI bundle may have changed.');
    }
    process.exit(0);
  }

  const result = patchFile(inputPath, outputPath);

  console.log('\n=== Patch Results ===');
  console.log(`Success: ${result.success}`);
  console.log('\nLocations:');
  for (const loc of result.locations) {
    console.log(`  ${loc.type}:`);
    console.log(`    Query Source: ${loc.querySource}`);
    console.log(`    Found: ${loc.found}`);
    console.log(`    Patched: ${loc.patched}`);
    if (loc.line) console.log(`    Line: ${loc.line}`);
    if (loc.functionName) console.log(`    Function: ${loc.functionName}`);
  }

  if (result.errors.length > 0) {
    console.log('\nErrors:');
    for (const err of result.errors) {
      console.log(`  - ${err}`);
    }
  }

  if (result.success) {
    console.log(`
=== Usage ===
After patching, you can disable the LLM checks with environment variables:

  # Skip prefix extraction (faster bash commands)
  export GENERIC_CODE_SKIP_PREFIX_CHECK=1

  # Skip path extraction from output
  export GENERIC_CODE_SKIP_PATH_EXTRACTION=1

  # Or both
  export GENERIC_CODE_SKIP_PREFIX_CHECK=1 GENERIC_CODE_SKIP_PATH_EXTRACTION=1
`);
  }

  process.exit(result.success ? 0 : 1);
}

main();