js-confuser 1.2.2 → 1.4.2

package/src/transforms/controlFlowFlattening/controlFlowFlattening.ts

@@ -1,17 +1,25 @@
 import { ok } from "assert";
+import { compileJsSync } from "../../compiler";
 import { ObfuscateOrder } from "../../order";
 import { ComputeProbabilityMap } from "../../probability";
 import Template from "../../templates/template";
 import { getBlock, isBlock, walk } from "../../traverse";
 import {
+  ArrayExpression,
+  ArrayPattern,
   AssignmentExpression,
   BinaryExpression,
   BreakStatement,
+  CallExpression,
   ConditionalExpression,
   ExpressionStatement,
+  FunctionDeclaration,
   Identifier,
+  IfStatement,
+  LabeledStatement,
   Literal,
   Node,
+  ReturnStatement,
   SequenceExpression,
   SwitchCase,
   SwitchStatement,
@@ -26,6 +34,7 @@ import {
 import {
   clone,
   getBlockBody,
+  getFunction,
   getVarContext,
   isVarContext,
 } from "../../util/insert";
@@ -34,7 +43,14 @@ import Transform from "../transform";
 import ChoiceFlowObfuscation from "./choiceFlowObfuscation";
 import ControlFlowObfuscation from "./controlFlowObfuscation";
 import ExpressionObfuscation from "./expressionObfuscation";
-import SwitchCaseObfuscation from "./switchCaseObfucation";
+import SwitchCaseObfuscation from "./switchCaseObfuscation";
+
+var flattenStructures = new Set([
+  "IfStatement",
+  "ForStatement",
+  "WhileStatement",
+  "DoWhileStatement",
+]);
 
 /**
  * Breaks functions into DAGs (Directed Acyclic Graphs)
@@ -50,23 +66,36 @@ import SwitchCaseObfuscation from "./switchCaseObfucation";
  * - 3. The while loop continues until the the state variable is the end state.
  */
 export default class ControlFlowFlattening extends Transform {
+  isDebug = false;
+
   constructor(o) {
     super(o, ObfuscateOrder.ControlFlowFlattening);
 
-    this.before.push(new ExpressionObfuscation(o));
+    if (!this.isDebug) {
+      this.before.push(new ExpressionObfuscation(o));
 
-    this.before.push(new ControlFlowObfuscation(o));
-    this.after.push(new SwitchCaseObfuscation(o));
+      this.after.push(new ControlFlowObfuscation(o));
+      this.after.push(new SwitchCaseObfuscation(o));
+    } else {
+      console.warn("Debug mode enabled");
+    }
 
     // this.after.push(new ChoiceFlowObfuscation(o));
   }
 
   match(object, parents) {
-    return isBlock(object);
+    return (
+      isBlock(object) &&
+      (!parents[1] || !flattenStructures.has(parents[1].type)) &&
+      (!parents[2] || !flattenStructures.has(parents[2].type))
+    );
   }
 
   transform(object, parents) {
     return () => {
+      if (object.body.length < 3) {
+        return;
+      }
       if (containsLexicallyBoundVariables(object, parents)) {
         return;
       }
@@ -86,16 +115,42 @@ export default class ControlFlowFlattening extends Transform {
 
       var functionDeclarations: Set<Node> = new Set();
       var fnNames: Set<string> = new Set();
+      var illegalFnNames = new Set<string>();
+
+      /**
+       * The variable names
+       *
+       * index -> var name
+       */
+      var stateVars = Array(this.isDebug ? 1 : getRandomInteger(2, 5))
+        .fill(0)
+        .map(() => this.getPlaceholder());
 
       body.forEach((stmt, i) => {
         if (stmt.type == "FunctionDeclaration") {
           functionDeclarations.add(stmt);
-          fnNames.add(stmt.id && stmt.id.name);
+          var name = stmt.id && stmt.id.name;
+          fnNames.add(name);
+          if (stmt.body.type !== "BlockStatement") {
+            illegalFnNames.add(name);
+          } else {
+            walk(stmt, [body, object, ...parents], (o, p) => {
+              if (
+                o.type == "ThisExpression" ||
+                o.type == "SuperExpression" ||
+                (o.type == "Identifier" &&
+                  (o.name == "arguments" || o.name == "eval"))
+              ) {
+                illegalFnNames.add(name);
+                return "EXIT";
+              }
+            });
+          }
         }
       });
 
       walk(object, parents, (o, p) => {
-        if (o.type == "Identifier") {
+        if (o.type == "Identifier" && fnNames.has(o.name)) {
           var info = getIdentifierInfo(o, p);
           if (!info.spec.isReferenced) {
             return;
@@ -112,6 +167,50 @@ export default class ControlFlowFlattening extends Transform {
               fnNames.delete(o.name);
             }
           }
+
+          if (!info.spec.isDefined) {
+            var b = getBlock(o, p);
+            if (b !== object || !p[0] || p[0].type !== "CallExpression") {
+              illegalFnNames.add(o.name);
+            } else {
+              var isExtractable = false;
+              if (p[1]) {
+                if (
+                  p[1].type == "ExpressionStatement" &&
+                  p[1].expression == p[0] &&
+                  p[2] == object.body
+                ) {
+                  isExtractable = true;
+                  p[1].$callExpression = "ExpressionStatement";
+                  p[1].$fnName = o.name;
+                } else if (
+                  p[1].type == "VariableDeclarator" &&
+                  p[1].init == p[0] &&
+                  p[2].length === 1 &&
+                  p[4] == object.body
+                ) {
+                  isExtractable = true;
+                  p[3].$callExpression = "VariableDeclarator";
+                  p[3].$fnName = o.name;
+                } else if (
+                  p[1].type == "AssignmentExpression" &&
+                  p[1].operator == "=" &&
+                  p[1].right === p[0] &&
+                  p[2] &&
+                  p[2].type == "ExpressionStatement" &&
+                  p[3] == object.body
+                ) {
+                  isExtractable = true;
+                  p[2].$callExpression = "AssignmentExpression";
+                  p[2].$fnName = o.name;
+                }
+              }
+
+              if (!isExtractable) {
+                illegalFnNames.add(o.name);
+              }
+            }
+          }
         }
       });
 
@@ -120,7 +219,9 @@ export default class ControlFlowFlattening extends Transform {
         return;
       }
 
-      var chunks: Node[][] = [[]];
+      illegalFnNames.forEach((illegal) => {
+        fnNames.delete(illegal);
+      });
 
       var fraction = 0.9;
       if (body.length > 20) {
@@ -131,74 +232,514 @@ export default class ControlFlowFlattening extends Transform {
         fraction = 0.5;
       }
 
-      body.forEach((x, i) => {
-        if (functionDeclarations.has(x)) {
-          return;
-        }
+      var resultVar = this.getPlaceholder();
+      var argVar = this.getPlaceholder();
+      var needsResultAndArgVar = false;
+      var fnToLabel: { [fnName: string]: string } = Object.create(null);
 
-        var currentChunk = chunks[chunks.length - 1];
+      fnNames.forEach((fnName) => {
+        fnToLabel[fnName] = this.getPlaceholder();
+      });
 
-        if (!currentChunk.length || Math.random() < fraction) {
-          currentChunk.push(x);
-        } else {
-          // Start new chunk
-          chunks.push([x]);
+      const flattenBody = (
+        body: Node[],
+        startingLabel = this.getPlaceholder()
+      ): { label: string; body: Node[] }[] => {
+        var chunks = [];
+        var currentBody = [];
+        var currentLabel = startingLabel;
+        const finishCurrentChunk = (
+          pointingLabel?: string,
+          newLabel?: string,
+          addGotoStatement = true
+        ) => {
+          if (!newLabel) {
+            newLabel = this.getPlaceholder();
+          }
+          if (!pointingLabel) {
+            pointingLabel = newLabel;
+          }
+
+          if (addGotoStatement) {
+            currentBody.push({ type: "GotoStatement", label: pointingLabel });
+          }
+
+          chunks.push({
+            label: currentLabel,
+            body: [...currentBody],
+          });
+
+          currentLabel = newLabel;
+          currentBody = [];
+        };
+
+        body.forEach((stmt, i) => {
+          if (functionDeclarations.has(stmt)) {
+            return;
+          }
+
+          if (stmt.$exit) {
+            currentBody.push(stmt);
+            currentBody.push(BreakStatement(switchLabel));
+            finishCurrentChunk(null, null, false);
+            return;
+          }
+
+          if (stmt.$callExpression && fnToLabel[stmt.$fnName]) {
+            var afterPath = this.getPlaceholder();
+            var args = [];
+
+            switch (stmt.$callExpression) {
+              // var a = fn();
+              case "VariableDeclarator":
+                args = stmt.declarations[0].init.arguments;
+                stmt.declarations[0].init = Identifier(resultVar);
+                break;
+
+              // fn();
+              case "ExpressionStatement":
+                args = stmt.expression.arguments;
+                stmt.expression = Identifier("undefined");
+                break;
+
+              // a = fn();
+              case "AssignmentExpression":
+                args = stmt.expression.right.arguments;
+                stmt.expression.right = Identifier(resultVar);
+                break;
+            }
+
+            needsResultAndArgVar = true;
+
+            currentBody.push(
+              ExpressionStatement(
+                AssignmentExpression(
+                  "=",
+                  Identifier(argVar),
+                  ArrayExpression([
+                    {
+                      type: "StateIdentifier",
+                      label: afterPath,
+                    },
+                    ArrayExpression(args),
+                  ])
+                )
+              )
+            );
+            finishCurrentChunk(fnToLabel[stmt.$fnName], afterPath);
+          }
+
+          if (stmt.type == "GotoStatement" && i !== body.length - 1) {
+            finishCurrentChunk(stmt.label);
+            return;
+          }
+
+          if (stmt.type == "LabeledStatement") {
+            var lbl = stmt.label.name;
+            var control = stmt.body;
+
+            var isSwitchStatement = control.type === "SwitchStatement";
+
+            if (
+              isSwitchStatement ||
+              ((control.type == "ForStatement" ||
+                control.type == "WhileStatement" ||
+                control.type == "DoWhileStatement") &&
+                control.body.type == "BlockStatement")
+            ) {
+              if (isSwitchStatement) {
+                if (
+                  control.cases.length == 0 || // at least 1 case
+                  control.cases.find(
+                    (x) =>
+                      !x.test || // cant be default case
+                      !x.consequent.length || // must have body
+                      x.consequent.findIndex(
+                        (node) => node.type == "BreakStatement"
+                      ) !==
+                        x.consequent.length - 1 || // break statement must be at the end
+                      x.consequent[x.consequent.length - 1].type !== // must end with break statement
+                        "BreakStatement" ||
+                      !x.consequent[x.consequent.length - 1].label || // must be labeled and correct
+                      x.consequent[x.consequent.length - 1].label.name != lbl
+                  )
+                ) {
+                  currentBody.push(stmt);
+                  return;
+                }
+              }
+
+              var isLoop = !isSwitchStatement;
+              var supportContinueStatement = isLoop;
+
+              var testPath = this.getPlaceholder();
+              var updatePath = this.getPlaceholder();
+              var bodyPath = this.getPlaceholder();
+              var afterPath = this.getPlaceholder();
+              var possible = true;
+              var toReplace = [];
+
+              walk(control.body, [], (o, p) => {
+                if (
+                  o.type == "BreakStatement" ||
+                  (supportContinueStatement && o.type == "ContinueStatement")
+                ) {
+                  if (!o.label || o.label.name !== lbl) {
+                    possible = false;
+                    return "EXIT";
+                  }
+                  if (o.label.name === lbl) {
+                    return () => {
+                      toReplace.push([
+                        o,
+                        {
+                          type: "GotoStatement",
+                          label:
+                            o.type == "BreakStatement" ? afterPath : updatePath,
+                        },
+                      ]);
+                    };
+                  }
+                }
+              });
+              if (!possible) {
+                currentBody.push(stmt);
+                return;
+              }
+              toReplace.forEach((v) => this.replace(v[0], v[1]));
+
+              if (isSwitchStatement) {
+                var switchVarName = this.getPlaceholder();
+
+                currentBody.push(
+                  VariableDeclaration(
+                    VariableDeclarator(switchVarName, control.discriminant)
+                  )
+                );
+
+                var afterPath = this.getPlaceholder();
+                finishCurrentChunk();
+                control.cases.forEach((switchCase, i) => {
+                  var entryPath = this.getPlaceholder();
+
+                  currentBody.push(
+                    IfStatement(
+                      BinaryExpression(
+                        "===",
+                        Identifier(switchVarName),
+                        switchCase.test
+                      ),
+                      [
+                        {
+                          type: "GotoStatement",
+                          label: entryPath,
+                        },
+                      ]
+                    )
+                  );
+
+                  chunks.push(
+                    ...flattenBody(
+                      [
+                        ...switchCase.consequent.slice(
+                          0,
+                          switchCase.consequent.length - 1
+                        ),
+                        {
+                          type: "GotoStatement",
+                          label: afterPath,
+                        },
+                      ],
+                      entryPath
+                    )
+                  );
+
+                  if (i === control.cases.length - 1) {
+                  } else {
+                    finishCurrentChunk();
+                  }
+                });
+
+                finishCurrentChunk(afterPath, afterPath);
+                return;
+              } else if (isLoop) {
+                var isPostTest = control.type == "DoWhileStatement";
+
+                // add initializing section to current chunk
+                if (control.init) {
+                  if (control.init.type == "VariableDeclaration") {
+                    currentBody.push(control.init);
+                  } else {
+                    currentBody.push(ExpressionStatement(control.init));
+                  }
+                }
+
+                // create new label called `testPath` and have current chunk point to it (goto testPath)
+                finishCurrentChunk(isPostTest ? bodyPath : testPath, testPath);
+
+                currentBody.push(
+                  IfStatement(control.test || Literal(true), [
+                    {
+                      type: "GotoStatement",
+                      label: bodyPath,
+                    },
+                  ])
+                );
+
+                // create new label called `bodyPath` and have test body point to afterPath (goto afterPath)
+                finishCurrentChunk(afterPath, bodyPath);
+
+                var innerBothPath = this.getPlaceholder();
+                chunks.push(
+                  ...flattenBody(
+                    [
+                      ...control.body.body,
+                      {
+                        type: "GotoStatement",
+                        label: updatePath,
+                      },
+                    ],
+                    innerBothPath
+                  )
+                );
+
+                finishCurrentChunk(innerBothPath, updatePath);
+
+                if (control.update) {
+                  currentBody.push(ExpressionStatement(control.update));
+                }
+
+                finishCurrentChunk(testPath, afterPath);
+                return;
+              }
+            }
+          }
+
+          if (
+            stmt.type == "IfStatement" &&
+            stmt.consequent.type == "BlockStatement" &&
+            (!stmt.alternate || stmt.alternate.type == "BlockStatement")
+          ) {
+            finishCurrentChunk();
+
+            var hasAlternate = !!stmt.alternate;
+            ok(!(hasAlternate && stmt.alternate.type !== "BlockStatement"));
+
+            var yesPath = this.getPlaceholder();
+            var noPath = this.getPlaceholder();
+            var afterPath = this.getPlaceholder();
+
+            currentBody.push(
+              IfStatement(stmt.test, [
+                {
+                  type: "GotoStatement",
+                  label: yesPath,
+                },
+              ])
+            );
+
+            chunks.push(
+              ...flattenBody(
+                [
+                  ...stmt.consequent.body,
+                  {
+                    type: "GotoStatement",
+                    label: afterPath,
+                  },
+                ],
+                yesPath
+              )
+            );
+
+            if (hasAlternate) {
+              chunks.push(
+                ...flattenBody(
+                  [
+                    ...stmt.alternate.body,
+                    {
+                      type: "GotoStatement",
+                      label: afterPath,
+                    },
+                  ],
+                  noPath
+                )
+              );
+
+              finishCurrentChunk(noPath, afterPath);
+            } else {
+              finishCurrentChunk(afterPath, afterPath);
+            }
+
+            return;
+          }
+
+          if (!currentBody.length || Math.random() < fraction) {
+            currentBody.push(stmt);
+          } else {
+            // Start new chunk
+            finishCurrentChunk();
+            currentBody.push(stmt);
+          }
+        });
+
+        finishCurrentChunk();
+        chunks[chunks.length - 1].body.pop();
+
+        return chunks;
+      };
+
+      var chunks = [];
+
+      /**
+       * label: switch(a+b+c){...break label...}
+       */
+      var switchLabel = this.getPlaceholder();
+
+      functionDeclarations.forEach((node) => {
+        if (node.id && fnNames.has(node.id.name)) {
+          var exitStateName = this.getPlaceholder();
+          var argumentsName = this.getPlaceholder();
+
+          needsResultAndArgVar = true;
+
+          node.body.body.push(ReturnStatement());
+
+          walk(node.body, [], (o, p) => {
+            if (o.type == "ReturnStatement") {
+              if (!getFunction(o, p)) {
+                return () => {
+                  var exitExpr = SequenceExpression([
+                    AssignmentExpression(
+                      "=",
+                      ArrayPattern(stateVars.map(Identifier)),
+                      Identifier(exitStateName)
+                    ),
+                    AssignmentExpression(
+                      "=",
+                      Identifier(resultVar),
+                      o.argument || Identifier("undefined")
+                    ),
+                  ]);
+
+                  this.replace(o, ReturnStatement(exitExpr));
+                };
+              }
+            }
+          });
+
+          var declarations = [
+            VariableDeclarator(
+              ArrayPattern([
+                Identifier(exitStateName),
+                Identifier(argumentsName),
+              ]),
+              Identifier(argVar)
+            ),
+          ];
+
+          if (node.params.length) {
+            declarations.push(
+              VariableDeclarator(
+                ArrayPattern(node.params),
+                Identifier(argumentsName)
+              )
+            );
+          }
+
+          var innerName = this.getPlaceholder();
+
+          chunks.push(
+            ...flattenBody(
+              [
+                FunctionDeclaration(
+                  innerName,
+                  [],
+                  [VariableDeclaration(declarations), ...node.body.body]
+                ),
+                this.objectAssign(
+                  ExpressionStatement(
+                    CallExpression(Identifier(innerName), [])
+                  ),
+                  {
+                    $exit: true,
+                  } as any
+                ),
+              ],
+              fnToLabel[node.id.name]
+            )
+          );
         }
       });
 
-      if (!chunks[chunks.length - 1].length) {
-        chunks.pop();
-      }
-      if (chunks.length < 2) {
-        return;
-      }
+      var startLabel = this.getPlaceholder();
 
-      // Add empty chunks
-      Array(getRandomInteger(1, 10))
-        .fill(0)
-        .forEach(() => {
-          chunks.splice(getRandomInteger(0, chunks.length), 0, []);
-        });
+      chunks.push(...flattenBody(body, startLabel));
+      chunks[chunks.length - 1].body.push({
+        type: "GotoStatement",
+        label: "END_LABEL",
+      });
+      chunks.push({
+        label: "END_LABEL",
+        body: [],
+      });
 
       var caseSelection: Set<number> = new Set();
 
-      var uniqueStatesNeeded = chunks.length + 1;
-
-      for (var i = 0; i < uniqueStatesNeeded; i++) {
-        var newState;
-        do {
-          newState = getRandomInteger(1, chunks.length * 15);
-        } while (caseSelection.has(newState));
+      var uniqueStatesNeeded = chunks.length;
+      var endLabel = chunks[Object.keys(chunks).length - 1].label;
 
+      do {
+        var newState = getRandomInteger(1, chunks.length * 15);
+        if (this.isDebug) {
+          newState = caseSelection.size;
+        }
         caseSelection.add(newState);
-      }
+      } while (caseSelection.size !== uniqueStatesNeeded);
 
       ok(caseSelection.size == uniqueStatesNeeded);
 
+      /**
+       * The accumulated state values
+       *
+       * index -> total state value
+       */
       var caseStates = Array.from(caseSelection);
 
-      var startState = caseStates[0];
-      var endState = caseStates[caseStates.length - 1];
+      /**
+       * The individual state values for each label
+       *
+       * labels right now are just chunk indexes (numbers)
+       *
+       * but will expand to if statements and functions when `goto statement` obfuscation is added
+       */
+      var labelToStates: { [label: string]: number[] } = Object.create(null);
 
-      var stateVars = Array(getRandomInteger(2, 7))
-        .fill(0)
-        .map(() => this.getPlaceholder());
-      var stateValues = Array(stateVars.length)
-        .fill(0)
-        .map(() => getRandomInteger(-250, 250));
+      Object.values(chunks).forEach((chunk, i) => {
+        var state = caseStates[i];
 
-      const getCurrentState = () => {
-        return stateValues.reduce((a, b) => b + a, 0);
-      };
+        var stateValues = Array(stateVars.length)
+          .fill(0)
+          .map(() => getRandomInteger(-250, 250));
+
+        const getCurrentState = () => {
+          return stateValues.reduce((a, b) => b + a, 0);
+        };
 
-      var correctIndex = getRandomInteger(0, stateVars.length);
-      stateValues[correctIndex] =
-        startState - (getCurrentState() - stateValues[correctIndex]);
+        var correctIndex = getRandomInteger(0, stateValues.length);
+        stateValues[correctIndex] =
+          state - (getCurrentState() - stateValues[correctIndex]);
 
-      var initStateValues = [...stateValues];
+        labelToStates[chunk.label] = stateValues;
+      });
+
+      // console.log(labelToStates);
+
+      var initStateValues = [...labelToStates[startLabel]];
+      var endState = labelToStates[endLabel].reduce((a, b) => b + a, 0);
 
-      const numberLiteral = (num, depth) => {
-        if (depth > 12 || Math.random() > 0.9 / (depth * 4)) {
+      const numberLiteral = (num, depth, stateValues) => {
+        ok(Array.isArray(stateValues));
+        if (depth > 10 || Math.random() > 0.8 / (depth * 4)) {
           return Literal(num);
         }
 
@@ -211,13 +752,17 @@ export default class ControlFlowFlattening extends Transform {
             operator == "<"
               ? x < stateValues[opposing]
               : x > stateValues[opposing];
-          var correct = numberLiteral(num, depth + 1);
-          var incorrect = numberLiteral(getRandomInteger(-250, 250), depth + 1);
+          var correct = numberLiteral(num, depth + 1, stateValues);
+          var incorrect = numberLiteral(
+            getRandomInteger(-250, 250),
+            depth + 1,
+            stateValues
+          );
 
           return ConditionalExpression(
             BinaryExpression(
               operator,
-              numberLiteral(x, depth + 1),
+              numberLiteral(x, depth + 1, stateValues),
               Identifier(stateVars[opposing])
             ),
             answer ? correct : incorrect,
@@ -228,66 +773,84 @@ export default class ControlFlowFlattening extends Transform {
         return BinaryExpression(
           "+",
           Identifier(stateVars[opposing]),
-          numberLiteral(num - stateValues[opposing], depth + 1)
+          numberLiteral(num - stateValues[opposing], depth + 1, stateValues)
         );
       };
 
-      const createTransitionStatement = (index, add) => {
-        var newValue = stateValues[index] + add;
+      const createTransitionExpression = (
+        index: number,
+        add: number,
+        mutatingStateValues: number[]
+      ) => {
+        var newValue = mutatingStateValues[index] + add;
 
         var expr = null;
 
-        if (Math.random() > 0.5) {
-          expr = ExpressionStatement(
-            AssignmentExpression(
-              "+=",
-              Identifier(stateVars[index]),
-              numberLiteral(add, 0)
-            )
+        if (this.isDebug) {
+          expr = AssignmentExpression(
+            "=",
+            Identifier(stateVars[index]),
+            Literal(newValue)
+          );
+        } else if (Math.random() > 0.5) {
+          expr = AssignmentExpression(
+            "+=",
+            Identifier(stateVars[index]),
+            numberLiteral(add, 0, mutatingStateValues)
           );
         } else {
-          var double = stateValues[index] * 2;
+          var double = mutatingStateValues[index] * 2;
           var diff = double - newValue;
 
           var first = AssignmentExpression(
             "*=",
             Identifier(stateVars[index]),
-            numberLiteral(2, 0)
-          );
-          stateValues[index] = double;
-
-          expr = ExpressionStatement(
-            SequenceExpression([
-              first,
-              AssignmentExpression(
-                "-=",
-                Identifier(stateVars[index]),
-                numberLiteral(diff, 0)
-              ),
-            ])
+            numberLiteral(2, 0, mutatingStateValues)
           );
+          mutatingStateValues[index] = double;
+
+          expr = SequenceExpression([
+            first,
+            AssignmentExpression(
+              "-=",
+              Identifier(stateVars[index]),
+              numberLiteral(diff, 0, mutatingStateValues)
+            ),
+          ]);
         }
 
-        stateValues[index] = newValue;
+        mutatingStateValues[index] = newValue;
 
         return expr;
       };
 
       interface Case {
         state: number;
-        nextState: number;
         body: Node[];
         order: number;
-        transitionStatements: Node[];
       }
 
       var order = Object.create(null);
-      var cases: Case[] = chunks.map((body, i) => {
+      var cases: Case[] = [];
+
+      chunks.forEach((chunk, i) => {
+        // skip last case, its empty and never ran
+        if (chunk.label === endLabel) {
+          return;
+        }
+
+        ok(labelToStates[chunk.label]);
+        var state = caseStates[i];
         var made = 1;
 
-        body.forEach((stmt) => {
+        var breaksInsertion = [];
+        var staticStateValues = [...labelToStates[chunk.label]];
+
+        chunk.body.forEach((stmt, stmtIndex) => {
+          var addBreak = false;
           walk(stmt, [], (o, p) => {
             if (
+              !this.isDebug &&
               o.type == "Literal" &&
               typeof o.value === "number" &&
               Math.floor(o.value) === o.value &&
@@ -299,63 +862,91 @@ export default class ControlFlowFlattening extends Transform {
               return () => {
                 this.replaceIdentifierOrLiteral(
                   o,
-                  numberLiteral(o.value, 0),
+                  numberLiteral(o.value, 0, staticStateValues),
                   p
                 );
               };
             }
+
+            if (o.type == "StateIdentifier") {
+              return () => {
+                ok(labelToStates[o.label]);
+                this.replace(
+                  o,
+                  ArrayExpression(labelToStates[o.label].map(Literal))
+                );
+              };
+            }
+
+            if (o.type == "GotoStatement") {
+              return () => {
+                var blockIndex = p.findIndex((node) => isBlock(node));
+                if (blockIndex === -1) {
+                  addBreak = true;
+                } else {
+                  var child = p[blockIndex - 2] || o;
+                  var childIndex = p[blockIndex].body.indexOf(child);
+
+                  p[blockIndex].body.splice(
+                    childIndex + 1,
+                    0,
+                    BreakStatement(switchLabel)
+                  );
+                }
+
+                var mutatingStateValues = [...labelToStates[chunk.label]];
+                var nextStateValues = labelToStates[o.label];
+                ok(nextStateValues, o.label);
+                this.replace(
+                  o,
+                  ExpressionStatement(
+                    SequenceExpression(
+                      mutatingStateValues.map((_v, stateValueIndex) => {
+                        var diff =
+                          nextStateValues[stateValueIndex] -
+                          mutatingStateValues[stateValueIndex];
+                        return createTransitionExpression(
+                          stateValueIndex,
+                          diff,
+                          mutatingStateValues
+                        );
+                      })
+                    )
+                  )
+                );
+              };
+            }
           });
+
+          if (addBreak) {
+            breaksInsertion.push(stmtIndex);
+          }
         });
 
-        var state = caseStates[i];
-        var nextState = caseStates[i + 1];
-        var diff = nextState - state;
-        var transitionStatements = [];
-
-        ok(!isNaN(diff));
-        var modifying = getRandomInteger(0, stateVars.length);
-        var shift = 0;
-
-        // var c1 = Identifier("undefined");
-        // this.addComment(c1, stateValues.join(", "));
-        // transitionStatements.push(c1);
-
-        transitionStatements.push(
-          ...Array.from(
-            new Set(
-              Array(getRandomInteger(0, stateVars.length - 2))
-                .fill(0)
-                .map(() => getRandomInteger(0, stateVars.length))
-                .filter((x) => x != modifying)
-            )
-          ).map((x) => {
-            var randomNumber = getRandomInteger(-250, 250);
+        breaksInsertion.sort();
+        breaksInsertion.reverse();
 
-            shift += randomNumber;
-            return createTransitionStatement(x, randomNumber);
-          })
-        );
-        transitionStatements.push(
-          createTransitionStatement(modifying, diff - shift)
-        );
+        breaksInsertion.forEach((index) => {
+          chunk.body.splice(index + 1, 0, BreakStatement(switchLabel));
+        });
 
         // var c = Identifier("undefined");
         // this.addComment(c, stateValues.join(", "));
         // transitionStatements.push(c);
 
         var caseObject = {
-          body: body,
+          body: chunk.body,
           state: state,
-          nextState: nextState,
           order: i,
-          transitionStatements: transitionStatements,
         };
         order[i] = caseObject;
 
-        return caseObject;
+        cases.push(caseObject);
       });
 
-      shuffle(cases);
+      if (!this.isDebug) {
+        shuffle(cases);
+      }
 
       var discriminant = Template(`${stateVars.join("+")}`).single().expression;
 
@@ -363,7 +954,9 @@ export default class ControlFlowFlattening extends Transform {
 
       if (functionDeclarations.size) {
         functionDeclarations.forEach((x) => {
-          body.unshift(clone(x));
+          if (!x.id || illegalFnNames.has(x.id.name)) {
+            body.unshift(clone(x));
+          }
         });
       }
 
@@ -374,28 +967,36 @@ export default class ControlFlowFlattening extends Transform {
 
           statements.push(...x.body);
 
-          statements.push(...x.transitionStatements);
-
-          statements.push(BreakStatement());
-
           var test = Literal(x.state);
 
           return SwitchCase(test, statements);
         })
       );
 
+      var declarations = [];
+
+      if (needsResultAndArgVar) {
+        declarations.push(VariableDeclarator(resultVar));
+        declarations.push(VariableDeclarator(argVar));
+      }
+
+      declarations.push(
+        ...stateVars.map((stateVar, i) => {
+          return VariableDeclarator(stateVar, Literal(initStateValues[i]));
+        })
+      );
+
       body.push(
-        VariableDeclaration(
-          stateVars.map((stateVar, i) => {
-            return VariableDeclarator(stateVar, Literal(initStateValues[i]));
-          })
-        ),
+        VariableDeclaration(declarations),
 
         WhileStatement(
           BinaryExpression("!=", clone(discriminant), Literal(endState)),
-          [switchStatement]
+          [LabeledStatement(switchLabel, switchStatement)]
         )
       );
+
+      // mark this object for switch case obfuscation
+      switchStatement.$controlFlowFlattening = true;
     };
   }
 }