From 84e978f9e8ec2c2402d2ddade7ff4ba2ddc00c80 Mon Sep 17 00:00:00 2001
From: echo094 <20028238+echo094@users.noreply.github.com>
Date: Sat, 1 Feb 2025 00:33:34 +0800
Subject: [PATCH] visitor: add split-assignment

Signed-off-by: echo094 <20028238+echo094@users.noreply.github.com>
---
 src/plugin/obfuscator.js                      | 15 +-----
 src/visitor/split-assignment.js               | 50 +++++++++++++++++++
 test/helper.js                                | 17 +++++++
 test/visitor/split-assignment.test.js         | 31 ++++++++++++
 .../if-assignment-valid.fix.js                |  4 ++
 .../split-assignment/if-assignment-valid.js   |  3 ++
 test/visitor/split-assignment/if-invalid.js   |  3 ++
 .../split-assignment/if-member-valid.fix.js   |  8 +++
 .../split-assignment/if-member-valid.js       |  6 +++
 .../split-assignment/variable-invalid.js      |  3 ++
 .../split-assignment/variable-valid.fix.js    |  4 ++
 .../split-assignment/variable-valid.js        |  3 ++
 12 files changed, 134 insertions(+), 13 deletions(-)
 create mode 100644 src/visitor/split-assignment.js
 create mode 100644 test/helper.js
 create mode 100644 test/visitor/split-assignment.test.js
 create mode 100644 test/visitor/split-assignment/if-assignment-valid.fix.js
 create mode 100644 test/visitor/split-assignment/if-assignment-valid.js
 create mode 100644 test/visitor/split-assignment/if-invalid.js
 create mode 100644 test/visitor/split-assignment/if-member-valid.fix.js
 create mode 100644 test/visitor/split-assignment/if-member-valid.js
 create mode 100644 test/visitor/split-assignment/variable-invalid.js
 create mode 100644 test/visitor/split-assignment/variable-valid.fix.js
 create mode 100644 test/visitor/split-assignment/variable-valid.js

diff --git a/src/plugin/obfuscator.js b/src/plugin/obfuscator.js
index a5a48c0..59cee75 100755
--- a/src/plugin/obfuscator.js
+++ b/src/plugin/obfuscator.js
@@ -18,6 +18,7 @@ import lintIfStatement from '../visitor/lint-if-statement.js'
 import mergeObject from '../visitor/merge-object.js'
 import parseControlFlowStorage from '../visitor/parse-control-flow-storage.js'
 import pruneIfBranch from '../visitor/prune-if-branch.js'
+import splitAssignment from '../visitor/split-assignment.js'
 import splitSequence from '../visitor/split-sequence.js'
 import splitVarDeclaration from '../visitor/split-variable-declaration.js'
 
@@ -732,18 +733,6 @@ function cleanDeadCode(ast) {
   return ast
 }
 
-const splitVariableDeclarator = {
-  VariableDeclarator(path) {
-    const init = path.get('init')
-    if (!init.isAssignmentExpression()) {
-      return
-    }
-    path.parentPath.insertBefore(init.node)
-    init.replaceWith(init.node.left)
-    path.parentPath.scope.crawl()
-  },
-}
-
 function standardIfStatement(path) {
   const consequent = path.get('consequent')
   const alternate = path.get('alternate')
@@ -802,8 +791,8 @@ function purifyCode(ast) {
       path.remove()
     },
   })
+  traverse(ast, splitAssignment)
   // 删除未使用的变量
-  traverse(ast, splitVariableDeclarator)
   traverse(ast, deleteUnusedVar)
   // 替换索引器
   function FormatMember(path) {
diff --git a/src/visitor/split-assignment.js b/src/visitor/split-assignment.js
new file mode 100644
index 0000000..ce5a189
--- /dev/null
+++ b/src/visitor/split-assignment.js
@@ -0,0 +1,50 @@
+import * as t from '@babel/types'
+
+/**
+ * Split the AssignmentExpressions. For example:
+ *
+ * - In the test of IfStatement
+ * - In the VariableDeclaration
+ */
+export default {
+  IfStatement(path) {
+    if (!path.parentPath.isBlockStatement() && !path.parentPath.isProgram()) {
+      return
+    }
+    let test = path.get('test')
+    if (test.isAssignmentExpression()) {
+      path.insertBefore(t.expressionStatement(test.node))
+      test.replaceWith(test.node.left)
+      path.scope.crawl()
+      return
+    }
+    if (test.isMemberExpression()) {
+      let object = test.get('object')
+      if (object.isAssignmentExpression()) {
+        path.insertBefore(t.expressionStatement(object.node))
+        object.replaceWith(object.node.left)
+      }
+      let property = test.get('property')
+      if (property.isAssignmentExpression()) {
+        path.insertBefore(t.expressionStatement(property.node))
+        property.replaceWith(property.node.left)
+      }
+      path.scope.crawl()
+    }
+  },
+  VariableDeclaration(path) {
+    if (!path.parentPath.isBlockStatement() && !path.parentPath.isProgram()) {
+      return
+    }
+    for (let i = 0; i < path.node.declarations.length; ++i) {
+      const declaration = path.get(`declarations.${i}`)
+      const init = declaration.node.init
+      if (!t.isAssignmentExpression(init)) {
+        continue
+      }
+      path.insertBefore(t.ExpressionStatement(init))
+      declaration.get('init').replaceWith(init.left)
+    }
+    path.scope.crawl()
+  },
+}
diff --git a/test/helper.js b/test/helper.js
new file mode 100644
index 0000000..dea3e77
--- /dev/null
+++ b/test/helper.js
@@ -0,0 +1,17 @@
+import fs from 'fs'
+import { expect } from 'vitest'
+import { parse } from '@babel/parser'
+import generate from '@babel/generator'
+import traverse from '@babel/traverse'
+
+export default function (visitor, fix, input) {
+  const sourceCode = fs.readFileSync(input + '.js', { encoding: 'utf-8' })
+  const ast = parse(sourceCode)
+  traverse(ast, visitor)
+  if (fix) {
+    const cmpCode = fs.readFileSync(input + '.fix.js', { encoding: 'utf-8' })
+    expect(generate(ast).code).toBe(cmpCode)
+  } else {
+    expect(generate(ast).code).toBe(sourceCode)
+  }
+}
diff --git a/test/visitor/split-assignment.test.js b/test/visitor/split-assignment.test.js
new file mode 100644
index 0000000..5484e26
--- /dev/null
+++ b/test/visitor/split-assignment.test.js
@@ -0,0 +1,31 @@
+import { join } from 'path'
+import { test } from 'vitest'
+import getResult from '../helper.js'
+import splitAssignment from '#visitor/split-assignment'
+
+const root = join(__dirname, 'split-assignment')
+
+test('if-invalid', () => {
+  const tc = 'if-invalid'
+  getResult(splitAssignment, false, join(root, tc))
+})
+
+test('if-assignment-valid', () => {
+  const tc = 'if-assignment-valid'
+  getResult(splitAssignment, true, join(root, tc))
+})
+
+test('if-member-valid', () => {
+  const tc = 'if-member-valid'
+  getResult(splitAssignment, true, join(root, tc))
+})
+
+test('variable-invalid', () => {
+  const tc = 'variable-invalid'
+  getResult(splitAssignment, false, join(root, tc))
+})
+
+test('variable-valid', () => {
+  const tc = 'variable-valid'
+  getResult(splitAssignment, true, join(root, tc))
+})
diff --git a/test/visitor/split-assignment/if-assignment-valid.fix.js b/test/visitor/split-assignment/if-assignment-valid.fix.js
new file mode 100644
index 0000000..e741bd0
--- /dev/null
+++ b/test/visitor/split-assignment/if-assignment-valid.fix.js
@@ -0,0 +1,4 @@
+let a = 1;
+let b;
+b = a;
+if (b) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/if-assignment-valid.js b/test/visitor/split-assignment/if-assignment-valid.js
new file mode 100644
index 0000000..71321a2
--- /dev/null
+++ b/test/visitor/split-assignment/if-assignment-valid.js
@@ -0,0 +1,3 @@
+let a = 1;
+let b;
+if (b = a) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/if-invalid.js b/test/visitor/split-assignment/if-invalid.js
new file mode 100644
index 0000000..cd73dff
--- /dev/null
+++ b/test/visitor/split-assignment/if-invalid.js
@@ -0,0 +1,3 @@
+let a = 1;
+let b;
+while (1) if (b = a) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/if-member-valid.fix.js b/test/visitor/split-assignment/if-member-valid.fix.js
new file mode 100644
index 0000000..a5ff8e7
--- /dev/null
+++ b/test/visitor/split-assignment/if-member-valid.fix.js
@@ -0,0 +1,8 @@
+let a = {
+  b: 1
+};
+let b;
+let c;
+b = a;
+c = 'b';
+if (b[c]) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/if-member-valid.js b/test/visitor/split-assignment/if-member-valid.js
new file mode 100644
index 0000000..3d1db6b
--- /dev/null
+++ b/test/visitor/split-assignment/if-member-valid.js
@@ -0,0 +1,6 @@
+let a = {
+  b: 1
+};
+let b;
+let c;
+if ((b = a)[c = 'b']) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/variable-invalid.js b/test/visitor/split-assignment/variable-invalid.js
new file mode 100644
index 0000000..6084b1e
--- /dev/null
+++ b/test/visitor/split-assignment/variable-invalid.js
@@ -0,0 +1,3 @@
+let a = 1;
+let b;
+for (let c = a = b;;) {}
\ No newline at end of file
diff --git a/test/visitor/split-assignment/variable-valid.fix.js b/test/visitor/split-assignment/variable-valid.fix.js
new file mode 100644
index 0000000..e8dade0
--- /dev/null
+++ b/test/visitor/split-assignment/variable-valid.fix.js
@@ -0,0 +1,4 @@
+let a = 1;
+let b;
+a = b;
+let c = a;
\ No newline at end of file
diff --git a/test/visitor/split-assignment/variable-valid.js b/test/visitor/split-assignment/variable-valid.js
new file mode 100644
index 0000000..033f085
--- /dev/null
+++ b/test/visitor/split-assignment/variable-valid.js
@@ -0,0 +1,3 @@
+let a = 1;
+let b;
+let c = a = b;
\ No newline at end of file