From b5fb883823feaf9088cf7caeb7d53f699db94d4e Mon Sep 17 00:00:00 2001
From: Duncan Ogilvie <mr.exodia.tpodt@gmail.com>
Date: Mon, 28 Oct 2024 02:10:44 +0100
Subject: [PATCH] Refactor analysis and output a liveness JSON

---
 obfuscator/CMakeLists.txt                   |  16 ++
 obfuscator/cmake.toml                       |   8 +-
 obfuscator/include/obfuscator/analyze.hpp   |  56 +++++-
 obfuscator/include/obfuscator/context.hpp   |   4 +-
 obfuscator/include/obfuscator/obfuscate.hpp |   2 +-
 obfuscator/src/obfuscate.cpp                | 157 +++++++++++++++-
 obfuscator/src/obfuscator/analyze.cpp       | 189 +++++++-------------
 obfuscator/src/obfuscator/context.cpp       |  73 +++++---
 obfuscator/src/obfuscator/obfuscate.cpp     |   2 +-
 9 files changed, 346 insertions(+), 161 deletions(-)

diff --git a/obfuscator/CMakeLists.txt b/obfuscator/CMakeLists.txt
index ad71042..49fdfc6 100644
--- a/obfuscator/CMakeLists.txt
+++ b/obfuscator/CMakeLists.txt
@@ -91,9 +91,24 @@ FetchContent_Declare(linux-pe
 )
 FetchContent_MakeAvailable(linux-pe)
 
+message(STATUS "Fetching json...")
+FetchContent_Declare(json
+	URL
+		"https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz"
+	URL_HASH
+		SHA256=d6c65aca6b1ed68e7a182f4757257b107ae403032760ed6ef121c9d55e81757d
+)
+FetchContent_MakeAvailable(json)
+
 # Target: obfuscator
 set(obfuscator_SOURCES
 	cmake.toml
+	"include/obfuscator/analyze.hpp"
+	"include/obfuscator/context.hpp"
+	"include/obfuscator/disassemble.hpp"
+	"include/obfuscator/msvc-secure.hpp"
+	"include/obfuscator/obfuscate.hpp"
+	"include/obfuscator/utility.hpp"
 	"src/obfuscator/analyze.cpp"
 	"src/obfuscator/context.cpp"
 	"src/obfuscator/disassemble.cpp"
@@ -136,6 +151,7 @@ if(CMKR_ROOT_PROJECT) # root
 	target_link_libraries(obfuscate PRIVATE
 		riscvm::obfuscator
 		args::args
+		nlohmann_json::nlohmann_json
 	)
 
 	get_directory_property(CMKR_VS_STARTUP_PROJECT DIRECTORY ${PROJECT_SOURCE_DIR} DEFINITION VS_STARTUP_PROJECT)
diff --git a/obfuscator/cmake.toml b/obfuscator/cmake.toml
index 1530fc7..248bbeb 100644
--- a/obfuscator/cmake.toml
+++ b/obfuscator/cmake.toml
@@ -29,11 +29,15 @@ sha256 = "312151a2d13c8327f5c9c586ac6cf7cddc1658e8f53edae0ec56509c8fa516c9"
 git = "https://github.com/can1357/linux-pe"
 tag = "be6d1f6fc30fb8058b5220e0bb2652a6dc8ec0b0"
 
+[fetch-content.json]
+url = "https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz"
+sha256 = "d6c65aca6b1ed68e7a182f4757257b107ae403032760ed6ef121c9d55e81757d"
+
 [target.obfuscator]
 type = "static"
 alias = "riscvm::obfuscator"
 sources = ["src/obfuscator/*.cpp"]
-headers = ["src/obfuscator/*.hpp"]
+headers = ["include/obfuscator/*.hpp"]
 include-directories = ["include"]
 compile-features = ["cxx_std_20"]
 link-libraries = ["zasm::zasm", "linux-pe", "fmt::fmt"]
@@ -42,7 +46,7 @@ link-libraries = ["zasm::zasm", "linux-pe", "fmt::fmt"]
 condition = "root"
 type = "executable"
 sources = ["src/obfuscate.cpp"]
-link-libraries = ["riscvm::obfuscator", "args::args"]
+link-libraries = ["riscvm::obfuscator", "args::args", "nlohmann_json::nlohmann_json"]
 
 [target.tests]
 type = "executable"
diff --git a/obfuscator/include/obfuscator/analyze.hpp b/obfuscator/include/obfuscator/analyze.hpp
index 6692b49..ff91e4d 100644
--- a/obfuscator/include/obfuscator/analyze.hpp
+++ b/obfuscator/include/obfuscator/analyze.hpp
@@ -1,8 +1,60 @@
 #pragma once
 
+#include <map>
+#include <set>
+#include <vector>
+
 #include <obfuscator/context.hpp>
 
 namespace obfuscator
 {
-bool analyze(Context& ctx, bool verbose = false);
-}
+struct BasicBlock
+{
+    uint64_t                 address = 0;
+    zasm::Label              label;
+    zasm::Node*              begin = nullptr;
+    zasm::Node*              end   = nullptr;
+    std::vector<zasm::Label> successors;
+};
+
+struct BlockLiveness
+{
+    uint32_t regsGen     = 0;
+    uint32_t regsKill    = 0;
+    uint32_t regsLiveIn  = 0;
+    uint32_t regsLiveOut = 0;
+
+    zasm::InstrCPUFlags flagsGen     = 0;
+    zasm::InstrCPUFlags flagsKill    = 0;
+    zasm::InstrCPUFlags flagsLiveIn  = 0;
+    zasm::InstrCPUFlags flagsLiveOut = 0;
+};
+
+struct InstructionLiveness
+{
+    zasm::Node*         node      = nullptr;
+    uint32_t            regsLive  = 0;
+    zasm::InstrCPUFlags flagsLive = 0;
+};
+
+class CFG
+{
+    explicit CFG(const zasm::Program& program) : program(program)
+    {
+    }
+
+  public:
+    const zasm::Program&           program;
+    std::map<uint64_t, BasicBlock> blocks;
+    std::set<uint64_t>             exits;
+
+    static zasm::Expected<CFG, std::string>
+    analyze(const zasm::Program& program, zasm::Label entry, bool verbose = false);
+
+    std::map<uint64_t, std::set<uint64_t>> getPredecessors() const;
+    std::map<uint64_t, BlockLiveness>      getLivenessBlocks(bool verbose = false) const;
+    std::vector<InstructionLiveness>
+    getInstructionLiveness(const std::map<uint64_t, BlockLiveness>& blockLiveness, bool verbose = false) const;
+};
+
+} // namespace obfuscator
diff --git a/obfuscator/include/obfuscator/context.hpp b/obfuscator/include/obfuscator/context.hpp
index 8480bea..38d365c 100644
--- a/obfuscator/include/obfuscator/context.hpp
+++ b/obfuscator/include/obfuscator/context.hpp
@@ -24,6 +24,8 @@ template <> struct fmt::formatter<zasm::Label::Id> : fmt::formatter<std::string_
 namespace obfuscator
 {
 
+const char*                flagToString(uint32_t flag);
+std::vector<uint32_t>      maskToFlags(uint32_t mask);
 std::string                formatFlagsMask(uint32_t mask);
 std::string                formatRegsMask(uint64_t mask);
 std::vector<zasm::x86::Gp> maskToRegs(uint64_t mask);
@@ -38,8 +40,8 @@ struct InstructionData
     uint32_t                regsWritten   = 0;
     uint32_t                regsRead      = 0;
 
-    zasm::InstrCPUFlags flagsLive = 0;
     uint32_t            regsLive  = 0;
+    zasm::InstrCPUFlags flagsLive = 0;
 };
 
 // Stores additional data for nodes in the zasm::Program
diff --git a/obfuscator/include/obfuscator/obfuscate.hpp b/obfuscator/include/obfuscator/obfuscate.hpp
index 6e50f7d..7ef7f5b 100644
--- a/obfuscator/include/obfuscator/obfuscate.hpp
+++ b/obfuscator/include/obfuscator/obfuscate.hpp
@@ -4,5 +4,5 @@
 
 namespace obfuscator
 {
-bool obfuscate(Context& ctx);
+bool obfuscate(Context& ctx, bool verbose = false);
 }
diff --git a/obfuscator/src/obfuscate.cpp b/obfuscator/src/obfuscate.cpp
index c81a22e..cd870d7 100644
--- a/obfuscator/src/obfuscate.cpp
+++ b/obfuscator/src/obfuscate.cpp
@@ -11,6 +11,7 @@
 
 #include <fmt/format.h>
 #include <args.hpp>
+#include <nlohmann/json.hpp>
 
 using namespace zasm;
 using namespace obfuscator;
@@ -26,6 +27,7 @@ struct Arguments : ArgumentParser
     std::string output;
     std::string cleanOutput;
     std::string payload;
+    bool        verbose = false;
 
     Arguments(int argc, char** argv) : ArgumentParser("Obfuscates the riscvm_run function")
     {
@@ -33,13 +35,141 @@ struct Arguments : ArgumentParser
         addString("-output", output, "Obfuscated function binary blob");
         addString("-clean-output", cleanOutput, "Unobfuscated function binary blob");
         addString("-payload", payload, "Payload to execute (Windows only)");
+        addBool("-verbose", verbose, "Verbose output");
         parseOrExit(argc, argv);
     }
 };
 
+static nlohmann::json livenessJson(uint64_t regsLive, InstrCPUFlags flagsLive)
+{
+    nlohmann::json json;
+
+    std::vector<std::string> regs;
+    for (const auto& reg : maskToRegs(regsLive))
+    {
+        auto name = formatter::toString(reg);
+        for (auto& ch : name)
+        {
+            if (ch >= 'a' && ch <= 'z')
+            {
+                ch = toupper(ch);
+            }
+        }
+        regs.push_back(std::move(name));
+    }
+    json["regs"] = regs;
+
+    std::vector<std::string> flags;
+    for (const auto& flag : maskToFlags(flagsLive))
+    {
+        auto name = flagToString(flag);
+        if (name != nullptr)
+        {
+            flags.push_back(name);
+        }
+    }
+    json["flags"] = flags;
+
+    return json;
+}
+
+static void dumpLiveness(CFG& cfg, const std::map<uint64_t, BlockLiveness>& livenessBlocks)
+{
+    nlohmann::json json;
+    // Print the results
+    std::string script;
+    for (const auto& [address, block] : cfg.blocks)
+    {
+        auto& liveness = livenessBlocks.at(address);
+
+        nlohmann::json blockJson;
+        blockJson["Liveness in"]  = livenessJson(liveness.regsLiveIn, liveness.flagsLiveIn);
+        blockJson["Liveness out"] = livenessJson(liveness.regsLiveOut, liveness.flagsLiveOut);
+
+        nlohmann::json instrJson;
+        fmt::println("Results for block {:#x}\n==========", address);
+        for (auto node = block.begin; node != block.end; node = node->getNext())
+        {
+            auto data = node->getUserData<InstructionData>();
+            auto str  = formatter::toString(cfg.program, node, formatter::Options::HexImmediates);
+
+            instrJson[fmt::format("{:#x}", data->address)] = livenessJson(data->regsLive, data->flagsLive);
+
+            script += "commentset ";
+            char address[32];
+            sprintf_s(address, "0x%llX", data->address);
+            script += address;
+            script += ", \"";
+            if (data->regsLive || data->flagsLive)
+            {
+                script += formatRegsMask(data->regsLive);
+                if (data->flagsLive)
+                {
+                    script += "|";
+                    script += formatFlagsMask(data->flagsLive);
+                }
+            }
+            else
+            {
+                script += "no live (HA)";
+            }
+            script += "\"\n";
+
+            fmt::println(
+                "{:#x}|{}|{}|{}", data->address, str, formatRegsMask(data->regsLive), formatFlagsMask(data->flagsLive)
+            );
+
+            if (data->regsRead & ~data->regsLive)
+            {
+                fmt::println("\tdead regs read: %s\n", formatRegsMask(data->regsRead & ~data->regsLive).c_str());
+                __debugbreak();
+            }
+        }
+        fmt::println("==========");
+
+        auto blockLiveness = livenessBlocks.at(address);
+        fmt::println("\tregs_live_in: {}", formatRegsMask(blockLiveness.regsLiveIn));
+        fmt::println("\tregs_live_out: {}", formatRegsMask(blockLiveness.regsLiveOut));
+        fmt::println("\tflags_live_in: {}", formatFlagsMask(blockLiveness.flagsLiveIn));
+        fmt::println("\tflags_live_out: {}", formatFlagsMask(blockLiveness.flagsLiveOut));
+
+        blockJson["Instr Liveness"] = std::move(instrJson);
+
+        json[fmt::format("{:#x}", address)] = std::move(blockJson);
+    }
+
+    fmt::println("{}", script);
+
+    auto toHex = [](uint64_t value)
+    {
+        char buffer[64] = "";
+        sprintf_s(buffer, "\"0x%llX\"", value);
+        return std::string(buffer);
+    };
+
+    std::string dot = "digraph G {\n";
+    for (const auto& [address, block] : cfg.blocks)
+    {
+        dot += toHex(address) + " [label=\"" + cfg.program.getLabelData(block.label).value().name + "\"];\n";
+        for (const auto& successor : block.successors)
+        {
+            auto data = cfg.program.getLabelData(successor).value().node->getUserData<InstructionData>();
+            auto successorAddress = data->address;
+            dot += toHex(address) + " -> " + toHex(successorAddress) + ";\n";
+        }
+    }
+    dot += "}";
+
+    fmt::println("{}", dot);
+
+    std::ofstream ofs("liveness.json");
+    ofs << json.dump(2);
+}
+
 int main(int argc, char** argv)
 {
     Arguments args(argc, argv);
+    auto      verbose = args.verbose;
 
     std::vector<uint8_t> pe;
     if (!loadFile(args.input, pe))
@@ -60,18 +190,37 @@ int main(int argc, char** argv)
 
     Program program(MachineMode::AMD64);
     Context ctx(program);
-    if (!disassemble(ctx, riscvmRunAddress, riscvmRunCode))
+    if (!disassemble(ctx, riscvmRunAddress, riscvmRunCode, verbose))
     {
         fmt::println("Failed to disassemble riscvm_run function.");
         return EXIT_FAILURE;
     }
 
-    if (!analyze(ctx, true))
+    // Analyze the CFG
+    auto cfg = CFG::analyze(program, program.getEntryPoint(), verbose);
+    if (!cfg)
     {
-        fmt::println("Failed to analyze the riscvm_run function.");
+        fmt::println("Failed to analyze the riscvm_run function: {}", cfg.error());
         return EXIT_FAILURE;
     }
 
+    // Perform liveness analysis
+    auto livenessBlocks      = cfg->getLivenessBlocks(verbose);
+    auto instructionLiveness = cfg->getInstructionLiveness(livenessBlocks, verbose);
+
+    // Add liveness information to the instruction data
+    for (const auto& instruction : instructionLiveness)
+    {
+        auto data       = instruction.node->getUserData<InstructionData>();
+        data->regsLive  = instruction.regsLive;
+        data->flagsLive = instruction.flagsLive;
+    }
+
+    if (verbose)
+    {
+        dumpLiveness(*cfg, livenessBlocks);
+    }
+
     auto serializeToFile = [&program](const std::string& outputFile, uint64_t base = 0)
     {
         // Serialize the obfuscated function
@@ -96,7 +245,7 @@ int main(int argc, char** argv)
         return EXIT_FAILURE;
     }
 
-    if (!obfuscate(ctx))
+    if (!obfuscate(ctx, verbose))
     {
         fmt::println("Failed to obfuscate riscvm_run function.");
         return EXIT_FAILURE;
diff --git a/obfuscator/src/obfuscator/analyze.cpp b/obfuscator/src/obfuscator/analyze.cpp
index b41c780..2b56f24 100644
--- a/obfuscator/src/obfuscator/analyze.cpp
+++ b/obfuscator/src/obfuscator/analyze.cpp
@@ -13,38 +13,18 @@ namespace obfuscator
 
 using namespace zasm;
 
-bool analyze(Context& ctx, bool verbose)
+Expected<CFG, std::string> CFG::analyze(const zasm::Program& program, Label entry, bool verbose)
 {
-    Program& program = ctx.program;
-    auto     mode    = program.getMode();
+    CFG cfg(program);
+
     if (verbose)
+    {
         fmt::println("=== ANALYZE ===");
+    }
     std::vector<Label> queue;
-    queue.push_back(program.getEntryPoint());
+    queue.push_back(entry);
     std::set<Label::Id> visisted;
 
-    // Construct the control flow graph
-    struct BasicBlock
-    {
-        uint64_t           address = 0;
-        Label              label;
-        Node*              begin = nullptr;
-        Node*              end   = nullptr;
-        std::vector<Label> successors;
-
-        uint32_t regsGen     = 0;
-        uint32_t regsKill    = 0;
-        uint32_t regsLiveIn  = 0;
-        uint32_t regsLiveOut = 0;
-
-        InstrCPUFlags flagsGen     = 0;
-        InstrCPUFlags flagsKill    = 0;
-        InstrCPUFlags flagsLiveIn  = 0;
-        InstrCPUFlags flagsLiveOut = 0;
-    };
-
-    std::map<uint64_t, BasicBlock> blocks;
-    std::set<uint64_t>             exits;
     while (!queue.empty())
     {
         auto blockStartLabel = queue.back();
@@ -82,23 +62,29 @@ bool analyze(Context& ctx, bool verbose)
                 queue.push_back(*label);
                 bb.successors.push_back(*label);
                 if (verbose)
+                {
                     fmt::println("not instr!");
+                }
                 break;
             }
 
             auto data = node->getUserData<InstructionData>();
-            auto str  = formatter::toString(program, instr, formatter::Options::HexImmediates);
             if (verbose)
+            {
+                auto str = formatter::toString(program, instr, formatter::Options::HexImmediates);
                 fmt::println("{:#x}|{}", data->address, str);
+            }
 
-            auto info = *instr->getDetail(mode);
+            auto info = *instr->getDetail(program.getMode());
             switch (info.getCategory())
             {
             case x86::Category::UncondBR:
             {
                 auto dest = instr->getOperand<Label>(0);
                 if (verbose)
+                {
                     fmt::println("UncondBR: {}", dest.getId());
+                }
                 queue.push_back(dest);
                 bb.successors.push_back(dest);
                 finished = true;
@@ -110,7 +96,9 @@ bool analyze(Context& ctx, bool verbose)
                 auto brtrue  = instr->getOperand<Label>(0);
                 auto brfalse = node->getNext()->get<Label>();
                 if (verbose)
+                {
                     fmt::println("CondBr: {}, {}", brtrue.getId(), brfalse.getId());
+                }
                 queue.push_back(brfalse);
                 queue.push_back(brtrue);
                 bb.successors.push_back(brtrue);
@@ -124,8 +112,7 @@ bool analyze(Context& ctx, bool verbose)
                 auto dest = instr->getOperand(0);
                 if (dest.getIf<Label>() != nullptr)
                 {
-                    fmt::println("unsupported call imm {:#x}", data->address);
-                    return false;
+                    return makeUnexpected(fmt::format("unsupported call imm {:#x}", data->address));
                 }
             }
             break;
@@ -133,7 +120,7 @@ bool analyze(Context& ctx, bool verbose)
             case x86::Category::Ret:
             {
                 finished = true;
-                exits.insert(bb.address);
+                cfg.exits.insert(bb.address);
             }
             break;
 
@@ -153,10 +140,14 @@ bool analyze(Context& ctx, bool verbose)
             __debugbreak();
         }
 
-        blocks.emplace(bb.address, bb);
+        cfg.blocks.emplace(bb.address, bb);
     }
 
-    // Compute the predecessors for each block
+    return cfg;
+}
+
+std::map<uint64_t, std::set<uint64_t>> CFG::getPredecessors() const
+{
     std::map<uint64_t, std::set<uint64_t>> predecessors;
     for (const auto& [address, _] : blocks)
     {
@@ -173,6 +164,11 @@ bool analyze(Context& ctx, bool verbose)
         }
     }
 
+    return predecessors;
+}
+
+std::map<uint64_t, BlockLiveness> CFG::getLivenessBlocks(bool verbose) const
+{
     // Perform liveness analysis on the control flow graph
     // https://en.wikipedia.org/wiki/Live-variable_analysis
 
@@ -181,6 +177,7 @@ bool analyze(Context& ctx, bool verbose)
     // does not take into account the next iteration of the loop"
 
     // Compute the GEN and KILL sets for each block
+    std::map<uint64_t, BlockLiveness> livenessBlocks;
     for (auto& [address, block] : blocks)
     {
         if (verbose)
@@ -189,6 +186,7 @@ bool analyze(Context& ctx, bool verbose)
             fmt::println("Analyzing block {:#x}\n==========\n{}\n==========", address, str);
         }
 
+        auto& liveness = livenessBlocks[address];
         for (auto node = block.begin; node != block.end; node = node->getNext())
         {
             auto  data   = node->getUserData<InstructionData>();
@@ -204,21 +202,23 @@ bool analyze(Context& ctx, bool verbose)
                 fmt::println("\tflags modified: {}", formatFlagsMask(data->flagsModified));
             }
 
-            block.regsGen |= data->regsRead & ~block.regsKill;
-            block.regsKill |= data->regsWritten;
-            block.flagsGen  = block.flagsGen | (data->flagsTested & ~block.flagsKill);
-            block.flagsKill = block.flagsKill | data->flagsModified;
+            liveness.regsGen |= data->regsRead & ~liveness.regsKill;
+            liveness.regsKill |= data->regsWritten;
+            liveness.flagsGen  = liveness.flagsGen | (data->flagsTested & ~liveness.flagsKill);
+            liveness.flagsKill = liveness.flagsKill | data->flagsModified;
         }
 
         if (verbose)
         {
-            fmt::println("regs_gen: {}", formatRegsMask(block.regsGen));
-            fmt::println("regs_kill: {}", formatRegsMask(block.regsKill));
-            fmt::println("flags_gen: {}", formatFlagsMask(block.flagsGen));
-            fmt::println("flags_kill: {}", formatFlagsMask(block.flagsKill));
+            fmt::println("regs_gen: {}", formatRegsMask(liveness.regsGen));
+            fmt::println("regs_kill: {}", formatRegsMask(liveness.regsKill));
+            fmt::println("flags_gen: {}", formatFlagsMask(liveness.flagsGen));
+            fmt::println("flags_kill: {}", formatFlagsMask(liveness.flagsKill));
         }
     }
 
+    auto predecessors = getPredecessors();
+
     // Solve the dataflow equations
     std::queue<uint64_t> worklist;
     for (auto exit : exits)
@@ -230,19 +230,19 @@ bool analyze(Context& ctx, bool verbose)
         auto address = worklist.front();
         worklist.pop();
 
-        auto&         block          = blocks.at(address);
-        auto          newRegsLiveIn  = block.regsGen | (block.regsLiveOut & ~block.regsKill);
-        InstrCPUFlags newFlagsLiveIn = block.flagsGen | (block.flagsLiveOut & ~block.flagsKill);
-        if (newRegsLiveIn != block.regsLiveIn || newFlagsLiveIn != block.flagsLiveIn)
+        auto&         liveness       = livenessBlocks.at(address);
+        auto          newRegsLiveIn  = liveness.regsGen | (liveness.regsLiveOut & ~liveness.regsKill);
+        InstrCPUFlags newFlagsLiveIn = liveness.flagsGen | (liveness.flagsLiveOut & ~liveness.flagsKill);
+        if (newRegsLiveIn != liveness.regsLiveIn || newFlagsLiveIn != liveness.flagsLiveIn)
         {
             // Update the LIVEin sets
-            block.regsLiveIn  = newRegsLiveIn;
-            block.flagsLiveIn = newFlagsLiveIn;
+            liveness.regsLiveIn  = newRegsLiveIn;
+            liveness.flagsLiveIn = newFlagsLiveIn;
 
             // Update the LIVEout sets in the predecessors and add them to the worklist
             for (const auto& predecessor : predecessors.at(address))
             {
-                auto& predecessorBlock = blocks.at(predecessor);
+                auto& predecessorBlock = livenessBlocks.at(predecessor);
                 predecessorBlock.regsLiveOut |= newRegsLiveIn;
                 predecessorBlock.flagsLiveOut = predecessorBlock.flagsLiveOut | newFlagsLiveIn;
                 worklist.push(predecessor);
@@ -250,7 +250,14 @@ bool analyze(Context& ctx, bool verbose)
         }
     }
 
+    return livenessBlocks;
+}
+
+std::vector<InstructionLiveness>
+CFG::getInstructionLiveness(const std::map<uint64_t, BlockLiveness>& livenessBlocks, bool verbose) const
+{
     // Compute liveness backwards for each block individually
+    std::vector<InstructionLiveness> result;
     for (auto& [address, block] : blocks)
     {
         if (verbose)
@@ -260,8 +267,9 @@ bool analyze(Context& ctx, bool verbose)
         }
 
         // We start with the live-out set of the block
-        InstrCPUFlags flagsLive = block.flagsLiveOut;
-        uint32_t      regsLive  = block.regsLiveOut;
+        auto&         liveness  = livenessBlocks.at(address);
+        InstrCPUFlags flagsLive = liveness.flagsLiveOut;
+        uint32_t      regsLive  = liveness.regsLiveOut;
         for (auto node = block.end->getPrev(); node != block.begin->getPrev(); node = node->getPrev())
         {
             auto  data   = node->getUserData<InstructionData>();
@@ -308,8 +316,7 @@ bool analyze(Context& ctx, bool verbose)
             }
 
             // Store the liveness state for the instruction
-            data->flagsLive = flagsLive;
-            data->regsLive  = regsLive;
+            result.push_back({node, regsLive, flagsLive});
 
             if (flagsModified)
             {
@@ -333,81 +340,7 @@ bool analyze(Context& ctx, bool verbose)
         }
     }
 
-    // Print the results
-    if (verbose)
-    {
-        std::string script;
-        for (const auto& [address, block] : blocks)
-        {
-            fmt::println("Results for block {:#x}\n==========", address);
-            for (auto node = block.begin; node != block.end; node = node->getNext())
-            {
-                auto data = node->getUserData<InstructionData>();
-                auto str  = formatter::toString(program, node, formatter::Options::HexImmediates);
-
-                script += "commentset ";
-                char address[32];
-                sprintf_s(address, "0x%llX", data->address);
-                script += address;
-                script += ", \"";
-                if (data->regsLive || data->flagsLive)
-                {
-                    script += formatRegsMask(data->regsLive);
-                    if (data->flagsLive)
-                    {
-                        script += "|";
-                        script += formatFlagsMask(data->flagsLive);
-                    }
-                }
-                else
-                {
-                    script += "no live (HA)";
-                }
-                script += "\"\n";
-
-                fmt::println(
-                    "{:#x}|{}|{}|{}", data->address, str, formatRegsMask(data->regsLive), formatFlagsMask(data->flagsLive)
-                );
-
-                if (data->regsRead & ~data->regsLive)
-                {
-                    fmt::println("\tdead regs read: %s\n", formatRegsMask(data->regsRead & ~data->regsLive).c_str());
-                    __debugbreak();
-                }
-            }
-            fmt::println("==========");
-
-            fmt::println("\tregs_live_in: {}", formatRegsMask(block.regsLiveIn));
-            fmt::println("\tregs_live_out: {}", formatRegsMask(block.regsLiveOut));
-            fmt::println("\tflags_live_in: {}", formatFlagsMask(block.flagsLiveIn));
-            fmt::println("\tflags_live_out: {}", formatFlagsMask(block.flagsLiveOut));
-        }
-
-        fmt::println("{}", script);
-
-        auto toHex = [](uint64_t value)
-        {
-            char buffer[64] = "";
-            sprintf_s(buffer, "\"0x%llX\"", value);
-            return std::string(buffer);
-        };
-
-        std::string dot = "digraph G {\n";
-        for (const auto& [address, block] : blocks)
-        {
-            dot += toHex(address) + " [label=\"" + program.getLabelData(block.label).value().name + "\"];\n";
-            for (const auto& successor : block.successors)
-            {
-                auto data = program.getLabelData(successor).value().node->getUserData<InstructionData>();
-                auto successorAddress = data->address;
-                dot += toHex(address) + " -> " + toHex(successorAddress) + ";\n";
-            }
-        }
-        dot += "}";
-
-        fmt::println("{}", dot);
-    }
-
-    return true;
+    return result;
 }
+
 } // namespace obfuscator
diff --git a/obfuscator/src/obfuscator/context.cpp b/obfuscator/src/obfuscator/context.cpp
index c4d1139..195c886 100644
--- a/obfuscator/src/obfuscator/context.cpp
+++ b/obfuscator/src/obfuscator/context.cpp
@@ -1,3 +1,5 @@
+#include <map>
+
 #include <obfuscator/context.hpp>
 
 #include <zasm/formatter/formatter.hpp>
@@ -8,31 +10,58 @@ namespace obfuscator
 
 using namespace zasm;
 
+static std::map<uint32_t, const char*> flagToStringMap = {
+    {ZYDIS_CPUFLAG_CF, "CF"},
+    {ZYDIS_CPUFLAG_PF, "PF"},
+    {ZYDIS_CPUFLAG_AF, "AF"},
+    {ZYDIS_CPUFLAG_ZF, "ZF"},
+    {ZYDIS_CPUFLAG_SF, "SF"},
+    {ZYDIS_CPUFLAG_TF, "TF"},
+    {ZYDIS_CPUFLAG_IF, "IF"},
+    {ZYDIS_CPUFLAG_DF, "DF"},
+    {ZYDIS_CPUFLAG_OF, "OF"},
+    {ZYDIS_CPUFLAG_NT, "NT"},
+    {ZYDIS_CPUFLAG_RF, "RF"},
+    {ZYDIS_CPUFLAG_VM, "VM"},
+    {ZYDIS_CPUFLAG_AC, "AC"},
+    {ZYDIS_CPUFLAG_VIF, "VIF"},
+    {ZYDIS_CPUFLAG_VIP, "VIP"},
+    {ZYDIS_CPUFLAG_ID, "ID"},
+};
+
+const char* flagToString(uint32_t flag)
+{
+    auto itr = flagToStringMap.find(flag);
+    return itr == flagToStringMap.end() ? nullptr : itr->second;
+}
+
+std::vector<uint32_t> maskToFlags(uint32_t mask)
+{
+    std::vector<uint32_t> result;
+    for (auto& [flag, name] : flagToStringMap)
+    {
+        if (mask & flag)
+        {
+            result.push_back(flag);
+        }
+    }
+    return result;
+}
+
 std::string formatFlagsMask(uint32_t mask)
 {
     std::string result;
-#define FLAG(x)   \
-    if (mask & x) \
-    result += (&(#x)[14]), result += " "
-    FLAG(ZYDIS_CPUFLAG_CF);
-    FLAG(ZYDIS_CPUFLAG_PF);
-    FLAG(ZYDIS_CPUFLAG_AF);
-    FLAG(ZYDIS_CPUFLAG_ZF);
-    FLAG(ZYDIS_CPUFLAG_SF);
-    FLAG(ZYDIS_CPUFLAG_TF);
-    FLAG(ZYDIS_CPUFLAG_IF);
-    FLAG(ZYDIS_CPUFLAG_DF);
-    FLAG(ZYDIS_CPUFLAG_OF);
-    FLAG(ZYDIS_CPUFLAG_NT);
-    FLAG(ZYDIS_CPUFLAG_RF);
-    FLAG(ZYDIS_CPUFLAG_VM);
-    FLAG(ZYDIS_CPUFLAG_AC);
-    FLAG(ZYDIS_CPUFLAG_VIF);
-    FLAG(ZYDIS_CPUFLAG_VIP);
-    FLAG(ZYDIS_CPUFLAG_ID);
-#undef FLAG
-    if (!result.empty())
-        result.pop_back();
+    for (auto& [flag, name] : flagToStringMap)
+    {
+        if (mask & flag)
+        {
+            if (!result.empty())
+            {
+                result += " ";
+            }
+            result += name;
+        }
+    }
     return "(" + result + ")";
 }
 
diff --git a/obfuscator/src/obfuscator/obfuscate.cpp b/obfuscator/src/obfuscator/obfuscate.cpp
index 63c947a..acac661 100644
--- a/obfuscator/src/obfuscator/obfuscate.cpp
+++ b/obfuscator/src/obfuscator/obfuscate.cpp
@@ -5,7 +5,7 @@ namespace obfuscator
 {
 using namespace zasm;
 
-bool obfuscate(Context& ctx)
+bool obfuscate(Context& ctx, bool verbose)
 {
     Program&       program = ctx.program;
     x86::Assembler assembler(program);