UnflattenerHelper

UnflattenerHelper.cs

using de4dot.blocks;
using dnlib.DotNet.Emit;
using System;
using System.Collections.Generic;
using System.Linq;

namespace de4dot.code.deobfuscators.AgentTesla
{
    public class UnflattenerHelper
    {
        private Block _startBlock;
        private Block _loopCondition;

        private int _initialCase;

        Dictionary<int, Block> _casesDict = new Dictionary<int, Block>();
        Dictionary<int, Block> _settersDict = new Dictionary<int, Block>();

        HashSet<Block> visitedBlocks = new HashSet<Block>();

        public UnflattenerHelper(Block block)
        {
            if (block.Instructions.Count < 2 
                || block.Instructions[0].OpCode.Code != Code.Ldc_I4 
                || block.Instructions[1].OpCode.Code != Code.Stloc)
                return;

            _initialCase = (int)block.Instructions[0].Operand;

            _startBlock = block;

            if (_startBlock.FallThrough == null
                || _startBlock.FallThrough.Sources == null
                || _startBlock.FallThrough.Sources.Count < 2)
                return;

            _loopCondition = _startBlock.FallThrough.Sources[1];

            ExploreControlFlow(_startBlock.FallThrough);
        }

        void ExploreControlFlow(Block block)
        {
            if (visitedBlocks.Contains(block))
                return;

            visitedBlocks.Add(block);

            int StartBlockNum = IsCaseStartBlock(block);
            if (StartBlockNum != -1)
            {
                if(!_casesDict.ContainsKey(StartBlockNum))
                    _casesDict.Add(StartBlockNum, block.FallThrough);
            }

            int nextCase = IsCaseEndBlock(block);
            if (nextCase != -1)
            {
                if (!_settersDict.ContainsKey(nextCase))
                    _settersDict.Add(nextCase, block);
            }

            if (block.FallThrough != null)
            {
                ExploreControlFlow(block.FallThrough);
            }

            if (block.Targets != null)
            {
                foreach (Block targetBlock in block.Targets)
                    ExploreControlFlow(targetBlock);
            }
        }

        int IsCaseEndBlock(Block block)
        {
            for (int i = 0; i < block.Instructions.Count; i++)
            {
                if (block.Instructions[i].OpCode.Code == Code.Ldc_I4 
                    && block.Instructions[i + 1].OpCode.Code == Code.Stloc 
                    && block.Instructions[i + 1].Operand == _startBlock.Instructions[1].Operand)
                {
                    return (int)block.Instructions[i].Operand;
                }
            }

            return -1;
        }

        int IsCaseStartBlock(Block block)
        {
            for (int i = 0; i + 3 <= block.Instructions.Count; i++)
            {
                if (block.Instructions[i].OpCode.Code == Code.Ldloc
                    && block.Instructions[i].Operand == _startBlock.Instructions[1].Operand
                    && block.Instructions[i + 1].OpCode.Code == Code.Ldc_I4
                    && block.Instructions[i + 2].OpCode.Code == Code.Ceq
                    && block.Instructions[i + 3].OpCode.Code == Code.Brfalse)
                {
                    return (int)block.Instructions[i + 1].Operand;
                }
            }

            return -1;
        }

        void CleanBlock(Block block)
        {
            for (int i = 0; i < block.Instructions.Count; i++)
            {
                if (block.Instructions[i].OpCode.Code == Code.Ldc_I4 
                    && block.Instructions[i + 1].OpCode.Code == Code.Stloc
                    && block.Instructions[i + 1].Operand == _startBlock.Instructions[1].Operand)
                {
                    block.Instructions[i] = new Instr(OpCodes.Nop.ToInstruction());
                    block.Instructions[i + 1] = new Instr(OpCodes.Nop.ToInstruction());
                }
            }
        }

        public bool Unflatten()
        {
            if (_casesDict.Count == 0)
                return false;

            Block firstCase = _casesDict[_initialCase];

            _startBlock.SetNewFallThrough(firstCase);

            foreach (var setter in _settersDict)
            {
                if (!_casesDict.ContainsKey(setter.Key))
                {
                    throw new Exception();
                }
                else
                {
                    // Remove the code which sets the next case;
                    CleanBlock(setter.Value);
                    setter.Value.SetNewFallThrough(_casesDict[setter.Key]);
                }
            }

            _startBlock.Instructions[0] = new Instr(OpCodes.Nop.ToInstruction());
            _startBlock.Instructions[1] = new Instr(OpCodes.Nop.ToInstruction());

            return true;
        }
    }
}