feat: add plugin API for manipulating the AST

apps/docs/src/guide/api.md

@@ -91,3 +91,65 @@ New folder structure:
 ```
 
 See [@codemod/matchers](https://github.com/codemod-js/codemod/tree/main/packages/matchers#readme) for more information about matchers.
+
+## Plugins
+
+There are 5 stages you can hook into to manipulate the AST, which run in this order:
+
+- parse
+- prepare
+- deobfuscate
+- unminify
+- unpack
+
+See the [babel plugin handbook](https://github.com/jamiebuilds/babel-handbook/blob/master/translations/en/plugin-handbook.md#writing-your-first-babel-plugin) for more information about writing plugins.
+This API is pretty similar, but there are some differences:
+
+- The required `runAfter` property specifies the stage
+- Only `visitor`, `pre` and `post` are supported
+- [parse](https://babeljs.io/docs/babel-parser),
+  [types](https://babeljs.io/docs/babel-types),
+  [traverse](https://babeljs.io/docs/babel-traverse),
+  [template](https://babeljs.io/docs/babel-template) and
+  [matchers](https://github.com/codemod-js/codemod/tree/main/packages/matchers) are passed to the plugin function
+
+### Example
+
+```js
+import { webcrack } from 'webcrack';
+
+function myPlugin({ types: t, matchers: m }) {
+  return {
+    runAfter: 'parse', // change it to 'unminify' and see what happens
+    pre(state) {
+      this.cache = new Set();
+    },
+    visitor: {
+      StringLiteral(path) {
+        this.cache.add(path.node.value);
+      },
+    },
+    post(state) {
+      console.log(this.cache); // Set(2) {'a', 'b'}
+    },
+  };
+}
+
+const result = await webcrack('"a" + "b"', { plugins: [myPlugin] });
+```
+
+### Using Babel plugins
+
+It should be compatible with most Babel plugins as long as they only access the limited API specified above.
+They have to be wrapped to set the `runAfter` property.
+
+```js
+import removeConsole from 'babel-plugin-transform-remove-console';
+
+function removeConsoleWrapper(babel) {
+  return {
+    runAfter: 'deobfuscate',
+    ...removeConsole(babel),
+  };
+}
+```

packages/webcrack/src/index.ts

@@ -20,6 +20,8 @@ import debugProtection from './deobfuscate/debug-protection';
 import mergeObjectAssignments from './deobfuscate/merge-object-assignments';
 import selfDefending from './deobfuscate/self-defending';
 import varFunctions from './deobfuscate/var-functions';
+import type { Plugin } from './plugin';
+import { loadPlugins } from './plugin';
 import jsx from './transforms/jsx';
 import jsxNew from './transforms/jsx-new';
 import mangle from './transforms/mangle';
@@ -35,6 +37,7 @@ import { unpackAST } from './unpack';
 import { isBrowser } from './utils/platform';
 
 export { type Sandbox } from './deobfuscate';
+export type { Plugin, PluginAPI, PluginObject, Stage } from './plugin';
 
 type Matchers = typeof m;
 
@@ -74,6 +77,10 @@ export interface Options {
    * @default false
    */
   mangle?: boolean;
+  /**
+   * Run AST transformations after specific stages
+   */
+  plugins?: Plugin[];
   /**
    * Assigns paths to modules based on the given matchers.
    * This will also rewrite `require()` calls to use the new paths.
@@ -103,6 +110,7 @@ function mergeOptions(options: Options): asserts options is Required<Options> {
     unpack: true,
     deobfuscate: true,
     mangle: false,
+    plugins: [],
     mappings: () => ({}),
     onProgress: () => {},
     sandbox: isBrowser() ? createBrowserSandbox() : createNodeSandbox(),
@@ -130,6 +138,7 @@ export async function webcrack(
   let ast: ParseResult<t.File> = null!;
   let outputCode = '';
   let bundle: Bundle | undefined;
+  const plugins = loadPlugins(options.plugins);
 
   const stages = [
     () => {
@@ -139,19 +148,27 @@ export async function webcrack(
         plugins: ['jsx'],
       }));
     },
+    plugins.parse && (() => plugins.parse!(ast)),
+
     () => {
       return applyTransforms(
         ast,
         [blockStatements, sequence, splitVariableDeclarations, varFunctions],
         { name: 'prepare' },
       );
     },
+    plugins.prepare && (() => plugins.prepare!(ast)),
+
     options.deobfuscate &&
       (() => applyTransformAsync(ast, deobfuscate, options.sandbox)),
+    plugins.deobfuscate && (() => plugins.deobfuscate!(ast)),
+
     options.unminify &&
       (() => {
         applyTransforms(ast, [transpile, unminify]);
       }),
+    plugins.unminify && (() => plugins.unminify!(ast)),
+
     options.mangle && (() => applyTransform(ast, mangle)),
     // TODO: Also merge unminify visitor (breaks selfDefending/debugProtection atm)
     (options.deobfuscate || options.jsx) &&
@@ -171,6 +188,7 @@ export async function webcrack(
     // Unpacking modifies the same AST and may result in imports not at top level
     // so the code has to be generated before
     options.unpack && (() => (bundle = unpackAST(ast, options.mappings(m)))),
+    plugins.unpack && (() => plugins.unpack!(ast)),
   ].filter(Boolean) as (() => unknown)[];
 
   for (let i = 0; i < stages.length; i++) {

packages/webcrack/src/plugin.ts

@@ -0,0 +1,76 @@
+import { parse } from '@babel/parser';
+import template from '@babel/template';
+import traverse, { visitors, type Visitor } from '@babel/traverse';
+import * as t from '@babel/types';
+import * as m from '@codemod/matchers';
+
+const stages = [
+  'parse',
+  'prepare',
+  'deobfuscate',
+  'unminify',
+  'unpack',
+] as const;
+
+export type Stage = (typeof stages)[number];
+
+export type PluginState = { opts: Record<string, unknown> };
+
+export interface PluginObject {
+  name?: string;
+  runAfter: Stage;
+  pre?: (this: PluginState, state: PluginState) => Promise<void> | void;
+  post?: (this: PluginState, state: PluginState) => Promise<void> | void;
+  visitor?: Visitor<PluginState>;
+}
+
+export interface PluginAPI {
+  parse: typeof parse;
+  types: typeof t;
+  traverse: typeof traverse;
+  template: typeof template;
+  matchers: typeof m;
+}
+
+export type Plugin = (api: PluginAPI) => PluginObject;
+
+export function loadPlugins(plugins: Plugin[]) {
+  const groups = new Map<Stage, PluginObject[]>(
+    stages.map((stage) => [stage, []]),
+  );
+  for (const plugin of plugins) {
+    const obj = plugin({
+      parse,
+      types: t,
+      traverse,
+      template,
+      matchers: m,
+    });
+    groups.get(obj.runAfter)?.push(obj);
+  }
+  return Object.fromEntries(
+    [...groups].map(([stage, plugins]) => [
+      stage,
+      plugins.length
+        ? async (ast: t.File) => {
+            const state: PluginState = { opts: {} };
+            for (const transform of plugins) {
+              await transform.pre?.call(state, state);
+            }
+
+            const pluginVisitors = plugins.flatMap(
+              (plugin) => plugin.visitor ?? [],
+            );
+            if (pluginVisitors.length > 0) {
+              const mergedVisitor = visitors.merge(pluginVisitors);
+              traverse(ast, mergedVisitor, undefined, state);
+            }
+
+            for (const plugin of plugins) {
+              await plugin.post?.call(state, state);
+            }
+          }
+        : undefined,
+    ]),
+  ) as Record<Stage, (ast: t.File) => Promise<void>>;
+}

packages/webcrack/test/plugins.test.ts

@@ -0,0 +1,24 @@
+import { expect, test, vi } from 'vitest';
+import { webcrack } from '../src';
+import type { Plugin } from '../src/plugin';
+
+test('run plugin after parse', async () => {
+  const pre = vi.fn();
+  const post = vi.fn();
+
+  const plugin: Plugin = ({ types: t }) => ({
+    runAfter: 'parse',
+    pre,
+    post,
+    visitor: {
+      NumericLiteral(path) {
+        path.replaceWith(t.stringLiteral(path.node.value.toString()));
+      },
+    },
+  });
+  const result = await webcrack('1 + 1;', { plugins: [plugin] });
+
+  expect(pre).toHaveBeenCalledOnce();
+  expect(post).toHaveBeenCalledOnce();
+  expect(result.code).toBe('"11";');
+});