Add ObfuscationPassManager (calls the obfuscation funcs based on attrs)

Slattz · Mar 6, 2023 · 8dee26a · 8dee26a
1 parent fd89509
commit 8dee26a
Show file tree

Hide file tree

Showing 16 changed files with 93 additions and 53 deletions.
diff --git a/main.cpp b/main.cpp
@@ -4,44 +4,29 @@
 #include <llvm/Support/raw_ostream.h>
 #include "common.hpp"
 #include "FuncAnnotationsParser.hpp"
+#include "ObfuscationPassManager.hpp"
 
-struct HelloWorld2 : llvm::PassInfoMixin<HelloWorld2> {
-    // Takes IR unit to run the pass on Module and the corresponding manager
-    llvm::PreservedAnalyses run(llvm::Module& M, llvm::ModuleAnalysisManager &) {
-        for (auto&& F : M.getFunctionList()) {
-            llvm::outs() << "Hello from: "<< F.getName() << "\n";
-            llvm::outs() << "   number of arguments: " << F.arg_size() << "\n";
-            llvm::outs() << "   number of attributes: " << F.getAttributes().getNumAttrSets() << "\n";
-            for (auto&& A : F.getAttributes()) {
-                llvm::outs() << A.getAsString() << "\n";
+namespace obfusc {
+    llvm::PassPluginLibraryInfo getObfuscPluginInfo() {
+        llvm::outs() << "ObfusC Version: " << OBFUSC_VERSION_STR << "\n";
+        return {LLVM_PLUGIN_API_VERSION, "ObfusC", OBFUSC_VERSION_STR,
+            [](llvm::PassBuilder &PB) {
+                PB.registerPipelineEarlySimplificationEPCallback(
+                    [=](llvm::ModulePassManager &MPM, llvm::OptimizationLevel Level) {
+                        MPM.addPass(obfusc::FuncAnnotationsParser());
+                    }
+                );
+                PB.registerOptimizerLastEPCallback(
+                    [=](llvm::ModulePassManager &MPM, llvm::OptimizationLevel Level) {
+                        MPM.addPass(obfusc::ObfuscationPassManager());
+                    }
+                );
             }
-        }
-        return llvm::PreservedAnalyses::all();
+        };
     }
-
-    // If false, pass is skipped for functions decorated with the optnone attribute (e.g. -O0).
-    static bool isRequired() { return true; }
-};
-
-llvm::PassPluginLibraryInfo getHelloWorldPluginInfo() {
-    llvm::outs() << "ObfusC Version: " << OBFUSC_VERSION_STR << "\n";
-    return {LLVM_PLUGIN_API_VERSION, "ObfusC", OBFUSC_VERSION_STR,
-        [](llvm::PassBuilder &PB) {
-            PB.registerPipelineEarlySimplificationEPCallback(
-                [=](llvm::ModulePassManager &MPM, llvm::OptimizationLevel Level) {
-                    MPM.addPass(obfusc::FuncAnnotationsParser());
-                }
-            );
-            PB.registerOptimizerLastEPCallback(
-                [=](llvm::ModulePassManager &MPM, llvm::OptimizationLevel Level) {
-                    MPM.addPass(HelloWorld2());
-                }
-            );
-        }
-    };
 }
 
 // Dynamic Library Entrypoint
 extern "C" LLVM_ATTRIBUTE_WEAK llvm::PassPluginLibraryInfo llvmGetPassPluginInfo() {
-    return getHelloWorldPluginInfo();
+    return obfusc::getObfuscPluginInfo();
 }
diff --git a/source/BCF/BcfPass.cpp b/source/BCF/BcfPass.cpp
@@ -4,7 +4,7 @@ namespace obfusc {
     BcfPass::BcfPass() {}
     BcfPass::~BcfPass() {}
 
-    bool BcfPass::obfuscate(llvm::Module& mod) {
+    bool BcfPass::obfuscate(llvm::Module& mod, llvm::Function& func) {
         return true;
     }
 

diff --git a/source/BCF/BcfPass.hpp b/source/BCF/BcfPass.hpp
@@ -7,6 +7,6 @@ namespace obfusc {
         BcfPass();
         ~BcfPass() override;
 
-        bool obfuscate(llvm::Module& mod) override;
+        bool obfuscate(llvm::Module& mod, llvm::Function& func) override;
     };
 }
diff --git a/source/CFF/CffPass.cpp b/source/CFF/CffPass.cpp
@@ -4,7 +4,7 @@ namespace obfusc {
     CffPass::CffPass() {}
     CffPass::~CffPass() {}
 
-    bool CffPass::obfuscate(llvm::Module& mod) {
+    bool CffPass::obfuscate(llvm::Module& mod, llvm::Function& func) {
         return true;
     }
 

diff --git a/source/CFF/CffPass.hpp b/source/CFF/CffPass.hpp
@@ -7,6 +7,6 @@ namespace obfusc {
         CffPass();
         ~CffPass() override;
 
-        bool obfuscate(llvm::Module& mod) override;
+        bool obfuscate(llvm::Module& mod, llvm::Function& func) override;
     };
 }
diff --git a/source/FuncAttributeStore.cpp b/source/FuncAttributeStore.cpp
@@ -7,27 +7,39 @@ namespace obfusc {
         return x;
     }
 
-    bool FuncAttributeStore::IsAttrStored(const char* attr) {
-        if (!attr || strlen(attr) < 1) {
+    bool FuncAttributeStore::IsAttrStored(llvm::StringRef& attr) {
+        if (attr.size() < 2) {
             return false;
         }
 
+        const char* attrCStr = attr.data();
+
+        if (attr[0] == '\"') {
+            attrCStr = attr.substr(1, attr.size()-2).data();
+        }
+
         for (auto& info : m_info) {
-            if (strncasecmp(info.name, attr, info.size) == 0) {
+            if (strncasecmp(info.name, attrCStr, info.size) == 0) {
                 return true;
             }
         }
 
         return false;
     }
 
-    IObfuscationPass* FuncAttributeStore::GetAttrPass(const char* attr) {
-        if (!attr || strlen(attr) < 1) {
+    IObfuscationPass* FuncAttributeStore::GetAttrPass(llvm::StringRef&& attr) {
+        if (attr.size() < 2) {
             return nullptr;
         }
 
+        const char* attrCStr = attr.data();
+
+        if (attr[0] == '\"') {
+            attrCStr = attr.substr(1, attr.size()-2).data();
+        }
+
         for (auto& info : m_info) {
-            if (strncasecmp(info.name, attr, info.size) == 0) {
+            if (strncasecmp(info.name, attrCStr, info.size) == 0) {
                 return info.pass;
             }
         }

diff --git a/source/FuncAttributeStore.hpp b/source/FuncAttributeStore.hpp
@@ -1,17 +1,18 @@
 #pragma once
 #include <vector>
+#include <llvm/ADT/StringRef.h>
 #include "IObfuscationPass.hpp"
 
 namespace obfusc {
     class FuncAttributeStore {
     public:
         static FuncAttributeStore& GetInstance();
-        bool IsAttrStored(const char* attr);
-        IObfuscationPass* GetAttrPass(const char* attr);
+        bool IsAttrStored(llvm::StringRef& attr);
+        IObfuscationPass* GetAttrPass(llvm::StringRef&& attr);
         void StoreAttributeInfo(const char* name, IObfuscationPass* pass);
 
     private:
-        FuncAttributeStore();
+        FuncAttributeStore() {}
 
         struct AttrStoreInfo {
             const char* name;

diff --git a/source/ISUB/iSubPass.cpp b/source/ISUB/iSubPass.cpp
@@ -4,7 +4,7 @@ namespace obfusc {
     iSubPass::iSubPass() {}
     iSubPass::~iSubPass() {}
 
-    bool iSubPass::obfuscate(llvm::Module& mod) {
+    bool iSubPass::obfuscate(llvm::Module& mod, llvm::Function& func) {
         return true;
     }
 

diff --git a/source/ISUB/iSubPass.hpp b/source/ISUB/iSubPass.hpp
@@ -7,6 +7,6 @@ namespace obfusc {
         iSubPass();
         ~iSubPass() override;
 
-        bool obfuscate(llvm::Module& mod) override;
+        bool obfuscate(llvm::Module& mod, llvm::Function& func) override;
     };
 }
diff --git a/source/MBA/MbaPass.cpp b/source/MBA/MbaPass.cpp
@@ -4,7 +4,9 @@ namespace obfusc {
     MbaPass::MbaPass() {}
     MbaPass::~MbaPass() {}
 
-    bool MbaPass::obfuscate(llvm::Module& mod) {
+    bool MbaPass::obfuscate(llvm::Module& mod, llvm::Function& func) {
+        llvm::outs() << "Hello from: "<< __PRETTY_FUNCTION__ << "\n";
+
         return true;
     }
 

diff --git a/source/MBA/MbaPass.hpp b/source/MBA/MbaPass.hpp
@@ -7,6 +7,6 @@ namespace obfusc {
         MbaPass();
         ~MbaPass() override;
 
-        bool obfuscate(llvm::Module& mod) override;
+        bool obfuscate(llvm::Module& mod, llvm::Function& func) override;
     };
 }
diff --git a/source/Passes/CMakeLists.txt b/source/Passes/CMakeLists.txt
@@ -1,3 +1,5 @@
 target_sources(ObfusC PRIVATE FuncAnnotationsParser.hpp)
 target_sources(ObfusC PRIVATE FuncAnnotationsParser.cpp)
-target_sources(ObfusC PRIVATE IObfuscationPass.hpp)
+target_sources(ObfusC PRIVATE IObfuscationPass.hpp)
+target_sources(ObfusC PRIVATE ObfuscationPassManager.hpp)
+target_sources(ObfusC PRIVATE ObfuscationPassManager.cpp)
diff --git a/source/Passes/FuncAnnotationsParser.cpp b/source/Passes/FuncAnnotationsParser.cpp
@@ -23,7 +23,7 @@ namespace obfusc {
 
             if (llvm::ConstantDataArray* strArray = llvm::dyn_cast<llvm::ConstantDataArray>(globalStrPtr->getOperand(0))) { //Get Annotation str
                 llvm::StringRef str = strArray->getAsString();
-                if (FuncAttributeStore::GetInstance().IsAttrStored(str.data())) {
+                if (FuncAttributeStore::GetInstance().IsAttrStored(str)) {
                     func->addFnAttr(str); //add Annotation to function
                 }
             }

diff --git a/source/Passes/IObfuscationPass.hpp b/source/Passes/IObfuscationPass.hpp
@@ -3,7 +3,7 @@
 
 class IObfuscationPass {
 public:
-    virtual bool obfuscate(llvm::Module& mod) = 0;
+    virtual bool obfuscate(llvm::Module& mod, llvm::Function& func) = 0;
 protected:
     IObfuscationPass() {}
     virtual ~IObfuscationPass() {}

diff --git a/source/Passes/ObfuscationPassManager.cpp b/source/Passes/ObfuscationPassManager.cpp
@@ -0,0 +1,26 @@
+#include "ObfuscationPassManager.hpp"
+#include "IObfuscationPass.hpp"
+#include "FuncAttributeStore.hpp"
+
+namespace obfusc {
+    llvm::PreservedAnalyses ObfuscationPassManager::run(llvm::Module& mod, llvm::ModuleAnalysisManager&) {
+        bool changed = false;
+        for (auto& func : mod.getFunctionList()) { //Get all functions in module
+            for (auto& attrs : func.getAttributes()) { //Get each attribute lists attached to function
+                for (auto& attr : attrs) { //Get attributes one by one
+                    if (attr.isStringAttribute()) { //If attribute is a string
+                        IObfuscationPass* pass = FuncAttributeStore::GetInstance().GetAttrPass(llvm::StringRef(attr.getAsString().data())); //Check if attr is an obfuscation pass
+                        if (pass) {
+                            changed |= pass->obfuscate(mod, func); //Call obfuscate func
+                        }
+                    }
+                }
+            }
+        }
+
+        if (changed) {
+            return llvm::PreservedAnalyses::none();
+        }
+        return llvm::PreservedAnalyses::all();
+    }
+}
diff --git a/source/Passes/ObfuscationPassManager.hpp b/source/Passes/ObfuscationPassManager.hpp
@@ -0,0 +1,12 @@
+#pragma once
+#include <llvm/Passes/PassBuilder.h>
+
+namespace obfusc {
+    struct ObfuscationPassManager : llvm::PassInfoMixin<ObfuscationPassManager> {
+        // Takes IR unit to run the pass on Module and the corresponding manager
+        llvm::PreservedAnalyses run(llvm::Module& mod, llvm::ModuleAnalysisManager&);
+
+        // If false, pass is skipped for functions decorated with the optnone attribute (e.g. -O0).
+        static bool isRequired() { return true; }
+    };
+}