js-confuser 1.2.2 → 1.4.2

package/dist/transforms/controlFlowFlattening/controlFlowFlattening.js

@@ -29,10 +29,13 @@ var _controlFlowObfuscation = _interopRequireDefault(require("./controlFlowObfus
 
 var _expressionObfuscation = _interopRequireDefault(require("./expressionObfuscation"));
 
-var _switchCaseObfucation = _interopRequireDefault(require("./switchCaseObfucation"));
+var _switchCaseObfuscation = _interopRequireDefault(require("./switchCaseObfuscation"));
 
 function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }
 
+function _defineProperty(obj, key, value) { if (key in obj) { Object.defineProperty(obj, key, { value: value, enumerable: true, configurable: true, writable: true }); } else { obj[key] = value; } return obj; }
+
+var flattenStructures = new Set(["IfStatement", "ForStatement", "WhileStatement", "DoWhileStatement"]);
 /**
  * Breaks functions into DAGs (Directed Acyclic Graphs)
  *
@@ -46,20 +49,33 @@ function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { de
  * - 2. At the end of each case, the state variable is updated to the next block of code.
  * - 3. The while loop continues until the the state variable is the end state.
  */
+
 class ControlFlowFlattening extends _transform.default {
   constructor(o) {
     super(o, _order.ObfuscateOrder.ControlFlowFlattening);
-    this.before.push(new _expressionObfuscation.default(o));
-    this.before.push(new _controlFlowObfuscation.default(o));
-    this.after.push(new _switchCaseObfucation.default(o)); // this.after.push(new ChoiceFlowObfuscation(o));
+
+    _defineProperty(this, "isDebug", false);
+
+    if (!this.isDebug) {
+      this.before.push(new _expressionObfuscation.default(o));
+      this.after.push(new _controlFlowObfuscation.default(o));
+      this.after.push(new _switchCaseObfuscation.default(o));
+    } else {
+      console.warn("Debug mode enabled");
+    } // this.after.push(new ChoiceFlowObfuscation(o));
+
   }
 
   match(object, parents) {
-    return (0, _traverse.isBlock)(object);
+    return (0, _traverse.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 ((0, _identifiers.containsLexicallyBoundVariables)(object, parents)) {
         return;
       }
@@ -78,14 +94,34 @@ class ControlFlowFlattening extends _transform.default {
 
       var functionDeclarations = new Set();
       var fnNames = new Set();
+      var illegalFnNames = new Set();
+      /**
+       * The variable names
+       *
+       * index -> var name
+       */
+
+      var stateVars = Array(this.isDebug ? 1 : (0, _random.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 {
+            (0, _traverse.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";
+              }
+            });
+          }
         }
       });
       (0, _traverse.walk)(object, parents, (o, p) => {
-        if (o.type == "Identifier") {
+        if (o.type == "Identifier" && fnNames.has(o.name)) {
           var info = (0, _identifiers.getIdentifierInfo)(o, p);
 
           if (!info.spec.isReferenced) {
@@ -103,6 +139,36 @@ class ControlFlowFlattening extends _transform.default {
               fnNames.delete(o.name);
             }
           }
+
+          if (!info.spec.isDefined) {
+            var b = (0, _traverse.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);
+              }
+            }
+          }
         }
       }); // redefined function,
 
@@ -110,7 +176,9 @@ class ControlFlowFlattening extends _transform.default {
         return;
       }
 
-      var chunks = [[]];
+      illegalFnNames.forEach(illegal => {
+        fnNames.delete(illegal);
+      });
       var fraction = 0.9;
 
       if (body.length > 20) {
@@ -123,63 +191,339 @@ class ControlFlowFlattening extends _transform.default {
         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 = Object.create(null);
+      fnNames.forEach(fnName => {
+        fnToLabel[fnName] = this.getPlaceholder();
+      });
 
-        var currentChunk = chunks[chunks.length - 1];
+      const flattenBody = (body, startingLabel = this.getPlaceholder()) => {
+        var chunks = [];
+        var currentBody = [];
+        var currentLabel = startingLabel;
 
-        if (!currentChunk.length || Math.random() < fraction) {
-          currentChunk.push(x);
-        } else {
-          // Start new chunk
-          chunks.push([x]);
-        }
-      });
+        const finishCurrentChunk = (pointingLabel, newLabel, addGotoStatement = true) => {
+          if (!newLabel) {
+            newLabel = this.getPlaceholder();
+          }
 
-      if (!chunks[chunks.length - 1].length) {
-        chunks.pop();
-      }
+          if (!pointingLabel) {
+            pointingLabel = newLabel;
+          }
 
-      if (chunks.length < 2) {
-        return;
-      } // Add empty chunks
+          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((0, _gen.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 = (0, _gen.Identifier)(resultVar);
+                break;
+              // fn();
+
+              case "ExpressionStatement":
+                args = stmt.expression.arguments;
+                stmt.expression = (0, _gen.Identifier)("undefined");
+                break;
+              // a = fn();
+
+              case "AssignmentExpression":
+                args = stmt.expression.right.arguments;
+                stmt.expression.right = (0, _gen.Identifier)(resultVar);
+                break;
+            }
+
+            needsResultAndArgVar = true;
+            currentBody.push((0, _gen.ExpressionStatement)((0, _gen.AssignmentExpression)("=", (0, _gen.Identifier)(argVar), (0, _gen.ArrayExpression)([{
+              type: "StateIdentifier",
+              label: afterPath
+            }, (0, _gen.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 = [];
+              (0, _traverse.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((0, _gen.VariableDeclaration)((0, _gen.VariableDeclarator)(switchVarName, control.discriminant)));
+                var afterPath = this.getPlaceholder();
+                finishCurrentChunk();
+                control.cases.forEach((switchCase, i) => {
+                  var entryPath = this.getPlaceholder();
+                  currentBody.push((0, _gen.IfStatement)((0, _gen.BinaryExpression)("===", (0, _gen.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((0, _gen.ExpressionStatement)(control.init));
+                  }
+                } // create new label called `testPath` and have current chunk point to it (goto testPath)
+
+
+                finishCurrentChunk(isPostTest ? bodyPath : testPath, testPath);
+                currentBody.push((0, _gen.IfStatement)(control.test || (0, _gen.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((0, _gen.ExpressionStatement)(control.update));
+                }
+
+                finishCurrentChunk(testPath, afterPath);
+                return;
+              }
+            }
+          }
 
-      Array((0, _random.getRandomInteger)(1, 10)).fill(0).forEach(() => {
-        chunks.splice((0, _random.getRandomInteger)(0, chunks.length), 0, []);
+          if (stmt.type == "IfStatement" && stmt.consequent.type == "BlockStatement" && (!stmt.alternate || stmt.alternate.type == "BlockStatement")) {
+            finishCurrentChunk();
+            var hasAlternate = !!stmt.alternate;
+            (0, _assert.ok)(!(hasAlternate && stmt.alternate.type !== "BlockStatement"));
+            var yesPath = this.getPlaceholder();
+            var noPath = this.getPlaceholder();
+            var afterPath = this.getPlaceholder();
+            currentBody.push((0, _gen.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((0, _gen.ReturnStatement)());
+          (0, _traverse.walk)(node.body, [], (o, p) => {
+            if (o.type == "ReturnStatement") {
+              if (!(0, _insert.getFunction)(o, p)) {
+                return () => {
+                  var exitExpr = (0, _gen.SequenceExpression)([(0, _gen.AssignmentExpression)("=", (0, _gen.ArrayPattern)(stateVars.map(_gen.Identifier)), (0, _gen.Identifier)(exitStateName)), (0, _gen.AssignmentExpression)("=", (0, _gen.Identifier)(resultVar), o.argument || (0, _gen.Identifier)("undefined"))]);
+                  this.replace(o, (0, _gen.ReturnStatement)(exitExpr));
+                };
+              }
+            }
+          });
+          var declarations = [(0, _gen.VariableDeclarator)((0, _gen.ArrayPattern)([(0, _gen.Identifier)(exitStateName), (0, _gen.Identifier)(argumentsName)]), (0, _gen.Identifier)(argVar))];
+
+          if (node.params.length) {
+            declarations.push((0, _gen.VariableDeclarator)((0, _gen.ArrayPattern)(node.params), (0, _gen.Identifier)(argumentsName)));
+          }
+
+          var innerName = this.getPlaceholder();
+          chunks.push(...flattenBody([(0, _gen.FunctionDeclaration)(innerName, [], [(0, _gen.VariableDeclaration)(declarations), ...node.body.body]), this.objectAssign((0, _gen.ExpressionStatement)((0, _gen.CallExpression)((0, _gen.Identifier)(innerName), [])), {
+            $exit: true
+          })], fnToLabel[node.id.name]));
+        }
+      });
+      var startLabel = this.getPlaceholder();
+      chunks.push(...flattenBody(body, startLabel));
+      chunks[chunks.length - 1].body.push({
+        type: "GotoStatement",
+        label: "END_LABEL"
+      });
+      chunks.push({
+        label: "END_LABEL",
+        body: []
       });
       var caseSelection = new Set();
-      var uniqueStatesNeeded = chunks.length + 1;
+      var uniqueStatesNeeded = chunks.length;
+      var endLabel = chunks[Object.keys(chunks).length - 1].label;
 
-      for (var i = 0; i < uniqueStatesNeeded; i++) {
-        var newState;
+      do {
+        var newState = (0, _random.getRandomInteger)(1, chunks.length * 15);
 
-        do {
-          newState = (0, _random.getRandomInteger)(1, chunks.length * 15);
-        } while (caseSelection.has(newState));
+        if (this.isDebug) {
+          newState = caseSelection.size;
+        }
 
         caseSelection.add(newState);
-      }
+      } while (caseSelection.size !== uniqueStatesNeeded);
 
       (0, _assert.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];
-      var stateVars = Array((0, _random.getRandomInteger)(2, 7)).fill(0).map(() => this.getPlaceholder());
-      var stateValues = Array(stateVars.length).fill(0).map(() => (0, _random.getRandomInteger)(-250, 250));
+      /**
+       * 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 = Object.create(null);
+      Object.values(chunks).forEach((chunk, i) => {
+        var state = caseStates[i];
+        var stateValues = Array(stateVars.length).fill(0).map(() => (0, _random.getRandomInteger)(-250, 250));
 
-      const getCurrentState = () => {
-        return stateValues.reduce((a, b) => b + a, 0);
-      };
+        const getCurrentState = () => {
+          return stateValues.reduce((a, b) => b + a, 0);
+        };
+
+        var correctIndex = (0, _random.getRandomInteger)(0, stateValues.length);
+        stateValues[correctIndex] = state - (getCurrentState() - stateValues[correctIndex]);
+        labelToStates[chunk.label] = stateValues;
+      }); // console.log(labelToStates);
 
-      var correctIndex = (0, _random.getRandomInteger)(0, stateVars.length);
-      stateValues[correctIndex] = startState - (getCurrentState() - stateValues[correctIndex]);
-      var initStateValues = [...stateValues];
+      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) => {
+        (0, _assert.ok)(Array.isArray(stateValues));
+
+        if (depth > 10 || Math.random() > 0.8 / (depth * 4)) {
           return (0, _gen.Literal)(num);
         }
 
@@ -189,95 +533,142 @@ class ControlFlowFlattening extends _transform.default {
           var x = (0, _random.getRandomInteger)(-250, 250);
           var operator = (0, _random.choice)(["<", ">"]);
           var answer = operator == "<" ? x < stateValues[opposing] : x > stateValues[opposing];
-          var correct = numberLiteral(num, depth + 1);
-          var incorrect = numberLiteral((0, _random.getRandomInteger)(-250, 250), depth + 1);
-          return (0, _gen.ConditionalExpression)((0, _gen.BinaryExpression)(operator, numberLiteral(x, depth + 1), (0, _gen.Identifier)(stateVars[opposing])), answer ? correct : incorrect, answer ? incorrect : correct);
+          var correct = numberLiteral(num, depth + 1, stateValues);
+          var incorrect = numberLiteral((0, _random.getRandomInteger)(-250, 250), depth + 1, stateValues);
+          return (0, _gen.ConditionalExpression)((0, _gen.BinaryExpression)(operator, numberLiteral(x, depth + 1, stateValues), (0, _gen.Identifier)(stateVars[opposing])), answer ? correct : incorrect, answer ? incorrect : correct);
         }
 
-        return (0, _gen.BinaryExpression)("+", (0, _gen.Identifier)(stateVars[opposing]), numberLiteral(num - stateValues[opposing], depth + 1));
+        return (0, _gen.BinaryExpression)("+", (0, _gen.Identifier)(stateVars[opposing]), numberLiteral(num - stateValues[opposing], depth + 1, stateValues));
       };
 
-      const createTransitionStatement = (index, add) => {
-        var newValue = stateValues[index] + add;
+      const createTransitionExpression = (index, add, mutatingStateValues) => {
+        var newValue = mutatingStateValues[index] + add;
         var expr = null;
 
-        if (Math.random() > 0.5) {
-          expr = (0, _gen.ExpressionStatement)((0, _gen.AssignmentExpression)("+=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(add, 0)));
+        if (this.isDebug) {
+          expr = (0, _gen.AssignmentExpression)("=", (0, _gen.Identifier)(stateVars[index]), (0, _gen.Literal)(newValue));
+        } else if (Math.random() > 0.5) {
+          expr = (0, _gen.AssignmentExpression)("+=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(add, 0, mutatingStateValues));
         } else {
-          var double = stateValues[index] * 2;
+          var double = mutatingStateValues[index] * 2;
           var diff = double - newValue;
-          var first = (0, _gen.AssignmentExpression)("*=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(2, 0));
-          stateValues[index] = double;
-          expr = (0, _gen.ExpressionStatement)((0, _gen.SequenceExpression)([first, (0, _gen.AssignmentExpression)("-=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(diff, 0))]));
+          var first = (0, _gen.AssignmentExpression)("*=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(2, 0, mutatingStateValues));
+          mutatingStateValues[index] = double;
+          expr = (0, _gen.SequenceExpression)([first, (0, _gen.AssignmentExpression)("-=", (0, _gen.Identifier)(stateVars[index]), numberLiteral(diff, 0, mutatingStateValues))]);
         }
 
-        stateValues[index] = newValue;
+        mutatingStateValues[index] = newValue;
         return expr;
       };
 
       var order = Object.create(null);
-      var cases = chunks.map((body, i) => {
+      var cases = [];
+      chunks.forEach((chunk, i) => {
+        // skip last case, its empty and never ran
+        if (chunk.label === endLabel) {
+          return;
+        }
+
+        (0, _assert.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;
           (0, _traverse.walk)(stmt, [], (o, p) => {
-            if (o.type == "Literal" && typeof o.value === "number" && Math.floor(o.value) === o.value && Math.abs(o.value) < 100000 && Math.random() < 4 / made && !p.find(x => (0, _insert.isVarContext)(x))) {
+            if (!this.isDebug && o.type == "Literal" && typeof o.value === "number" && Math.floor(o.value) === o.value && Math.abs(o.value) < 100000 && Math.random() < 4 / made && !p.find(x => (0, _insert.isVarContext)(x))) {
               made++;
               return () => {
-                this.replaceIdentifierOrLiteral(o, numberLiteral(o.value, 0), p);
+                this.replaceIdentifierOrLiteral(o, numberLiteral(o.value, 0, staticStateValues), p);
+              };
+            }
+
+            if (o.type == "StateIdentifier") {
+              return () => {
+                (0, _assert.ok)(labelToStates[o.label]);
+                this.replace(o, (0, _gen.ArrayExpression)(labelToStates[o.label].map(_gen.Literal)));
+              };
+            }
+
+            if (o.type == "GotoStatement") {
+              return () => {
+                var blockIndex = p.findIndex(node => (0, _traverse.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, (0, _gen.BreakStatement)(switchLabel));
+                }
+
+                var mutatingStateValues = [...labelToStates[chunk.label]];
+                var nextStateValues = labelToStates[o.label];
+                (0, _assert.ok)(nextStateValues, o.label);
+                this.replace(o, (0, _gen.ExpressionStatement)((0, _gen.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 = [];
-        (0, _assert.ok)(!isNaN(diff));
-        var modifying = (0, _random.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((0, _random.getRandomInteger)(0, stateVars.length - 2)).fill(0).map(() => (0, _random.getRandomInteger)(0, stateVars.length)).filter(x => x != modifying))).map(x => {
-          var randomNumber = (0, _random.getRandomInteger)(-250, 250);
-          shift += randomNumber;
-          return createTransitionStatement(x, randomNumber);
-        }));
-        transitionStatements.push(createTransitionStatement(modifying, diff - shift)); // var c = Identifier("undefined");
+        breaksInsertion.sort();
+        breaksInsertion.reverse();
+        breaksInsertion.forEach(index => {
+          chunk.body.splice(index + 1, 0, (0, _gen.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
         };
         order[i] = caseObject;
-        return caseObject;
+        cases.push(caseObject);
       });
-      (0, _random.shuffle)(cases);
+
+      if (!this.isDebug) {
+        (0, _random.shuffle)(cases);
+      }
+
       var discriminant = (0, _template.default)("".concat(stateVars.join("+"))).single().expression;
       body.length = 0;
 
       if (functionDeclarations.size) {
         functionDeclarations.forEach(x => {
-          body.unshift((0, _insert.clone)(x));
+          if (!x.id || illegalFnNames.has(x.id.name)) {
+            body.unshift((0, _insert.clone)(x));
+          }
         });
       }
 
       var switchStatement = (0, _gen.SwitchStatement)(discriminant, cases.map((x, i) => {
         var statements = [];
         statements.push(...x.body);
-        statements.push(...x.transitionStatements);
-        statements.push((0, _gen.BreakStatement)());
         var test = (0, _gen.Literal)(x.state);
         return (0, _gen.SwitchCase)(test, statements);
       }));
-      body.push((0, _gen.VariableDeclaration)(stateVars.map((stateVar, i) => {
+      var declarations = [];
+
+      if (needsResultAndArgVar) {
+        declarations.push((0, _gen.VariableDeclarator)(resultVar));
+        declarations.push((0, _gen.VariableDeclarator)(argVar));
+      }
+
+      declarations.push(...stateVars.map((stateVar, i) => {
         return (0, _gen.VariableDeclarator)(stateVar, (0, _gen.Literal)(initStateValues[i]));
-      })), (0, _gen.WhileStatement)((0, _gen.BinaryExpression)("!=", (0, _insert.clone)(discriminant), (0, _gen.Literal)(endState)), [switchStatement]));
+      }));
+      body.push((0, _gen.VariableDeclaration)(declarations), (0, _gen.WhileStatement)((0, _gen.BinaryExpression)("!=", (0, _insert.clone)(discriminant), (0, _gen.Literal)(endState)), [(0, _gen.LabeledStatement)(switchLabel, switchStatement)])); // mark this object for switch case obfuscation
+
+      switchStatement.$controlFlowFlattening = true;
     };
   }