import { readFile as fsRead, writeFile as fsWrite } from "node:fs/promises";
import { buildSnapshot, type Exclusions } from "./snapshot.js";
import { fetchAccounts, type FetchLike, type RecordedAccounts } from "./fetch.js";

export const USAGE =
  "usage: kog-snapshot --mint <address> --exclusions <file> [--slot <n>] [--top <n>] [--out <file>] [--fixture <file>] [--save-fixture <file>]";

export interface RunDeps {
  env?: Record<string, string | undefined>;
  stdout?: (s: string) => void;
  stderr?: (s: string) => void;
  fetch?: FetchLike;
  readFile?: (path: string) => Promise<string>;
  writeFile?: (path: string, text: string) => Promise<void>;
}

class UsageError extends Error {}

const VALUE_FLAGS = new Set(["--mint", "--exclusions", "--slot", "--top", "--out", "--fixture", "--save-fixture"]);

function parseArgs(argv: string[]): Record<string, string> {
  const out: Record<string, string> = {};
  for (let i = 0; i < argv.length; i++) {
    let a = argv[i];
    let v: string | undefined;
    const eq = a.indexOf("=");
    if (a.startsWith("--") && eq > 0) {
      v = a.slice(eq + 1);
      a = a.slice(0, eq);
    }
    if (a === "--help" || a === "-h") throw new UsageError(USAGE);
    if (!VALUE_FLAGS.has(a)) throw new UsageError(`unknown option ${a}\n${USAGE}`);
    if (v === undefined) {
      v = argv[++i];
      if (v === undefined || v.startsWith("--")) throw new UsageError(`${a} needs a value`);
    }
    out[a.slice(2)] = v;
  }
  return out;
}

function wholeNumber(name: string, v: string | undefined, min: number): number | undefined {
  if (v === undefined) return undefined;
  if (!/^[0-9]+$/.test(v) || Number(v) < min || !Number.isSafeInteger(Number(v))) throw new UsageError(`${name} must be a whole number of at least ${min}`);
  return Number(v);
}

const BASE58 = /^[1-9A-HJ-NP-Za-km-z]{32,44}$/;

function readExclusions(text: string, path: string): Exclusions {
  let j: any;
  try {
    j = JSON.parse(text);
  } catch {
    throw new UsageError(`exclusions file ${path} is not valid JSON`);
  }
  if (!j || typeof j !== "object" || Array.isArray(j)) throw new UsageError(`exclusions file ${path} must be an object with owners and token_accounts`);
  for (const k of ["owners", "token_accounts"]) {
    if (j[k] !== undefined && !Array.isArray(j[k])) throw new UsageError(`exclusions ${k} must be a list`);
  }
  return { owners: j.owners ?? [], token_accounts: j.token_accounts ?? [] };
}

/** Runs the CLI and returns the exit code (0 on success, 1 on any error). Never prints the RPC URL. */
export async function run(argv: string[], deps: RunDeps = {}): Promise<number> {
  const env = deps.env ?? process.env;
  const stdout = deps.stdout ?? ((s) => process.stdout.write(s));
  const stderr = deps.stderr ?? ((s) => process.stderr.write(s));
  const readFile = deps.readFile ?? ((p) => fsRead(p, "utf8"));
  const writeFile = deps.writeFile ?? ((p, t) => fsWrite(p, t, "utf8"));
  const rpcUrl = env.SOLANA_RPC_URL;
  const scrub = (msg: string) => (rpcUrl ? msg.split(rpcUrl).join("<SOLANA_RPC_URL>") : msg);
  try {
    const args = parseArgs(argv);
    if (!args.mint) throw new UsageError(`missing --mint\n${USAGE}`);
    if (!BASE58.test(args.mint)) throw new UsageError("--mint is not a base58 Solana address");
    if (!args.exclusions) throw new UsageError(`missing --exclusions\n${USAGE}`);
    const slot = wholeNumber("--slot", args.slot, 0);
    const top = wholeNumber("--top", args.top, 1);

    let exText: string;
    try {
      exText = await readFile(args.exclusions);
    } catch {
      throw new UsageError(`cannot read exclusions file ${args.exclusions}`);
    }
    const exclusions = readExclusions(exText, args.exclusions);

    let recorded: RecordedAccounts;
    if (args.fixture) {
      let fx: any;
      try {
        fx = JSON.parse(await readFile(args.fixture));
      } catch {
        throw new UsageError(`cannot read fixture file ${args.fixture}`);
      }
      if (!fx || typeof fx.slot !== "number" || !Array.isArray(fx.accounts)) throw new UsageError(`fixture ${args.fixture} must hold {slot, accounts}`);
      if (slot !== undefined && fx.slot < slot) throw new UsageError(`fixture slot ${fx.slot} is below --slot ${slot}`);
      recorded = { slot: fx.slot, accounts: fx.accounts };
    } else {
      if (!rpcUrl) throw new UsageError("set SOLANA_RPC_URL or pass --fixture");
      recorded = await fetchAccounts({ rpcUrl, mint: args.mint, fetch: deps.fetch, minContextSlot: slot });
    }
    if (args["save-fixture"]) await writeFile(args["save-fixture"], JSON.stringify({ slot: recorded.slot, accounts: recorded.accounts }, null, 1) + "\n");

    const snap = buildSnapshot({ mint: args.mint, slot: recorded.slot, accounts: recorded.accounts, exclusions, top });
    const text = JSON.stringify(snap, null, 2) + "\n";
    if (args.out) await writeFile(args.out, text);
    else stdout(text);
    return 0;
  } catch (e) {
    stderr(`kog-snapshot: ${scrub(e instanceof Error ? e.message : String(e))}\n`);
    return 1;
  }
}

agentgate-source: code mirror on Cloudflare Pages. All repos