Merge branch 'excl_func' (Issue #33)

hasherezade · Aug 8, 2023 · 31e29a9 · 31e29a9
2 parents 09ba22c + d9e5b9c
commit 31e29a9
Show file tree

Hide file tree

Showing 6 changed files with 134 additions and 10 deletions.
diff --git a/FuncWatch.cpp b/FuncWatch.cpp
@@ -10,12 +10,15 @@ bool WFuncInfo::load(const std::string &sline, char delimiter)
 {
     std::vector<std::string> args;
     util::splitList(sline, delimiter, args);
-    if (args.size() < 3) return false;
+    if (args.size() < 2) return false;
 
     this->dllName = args[0];
     this->funcName = args[1];
-    this->paramCount = util::loadInt(args[2]);
+    this->paramCount = 0;
 
+    if (args.size() >= 3) {
+        this->paramCount = util::loadInt(args[2]);
+    }
     return true;
 }
 
@@ -62,6 +65,74 @@ bool WSyscallInfo::update(const WSyscallInfo& syscall_info)
 
 //---
 
+bool FuncExcludeList::contains(const std::string& dll_name, const std::string& func)
+{
+    if (!dll_name.length() || !func.length()) return false;
+    if (this->isEmpty()) return false;
+
+    const std::string shortDll = util::getDllName(dll_name);
+    for (auto itr = funcs.begin(); itr != funcs.end(); ++itr) {
+        WFuncInfo& fInfo = *itr;
+        if (util::iequals(fInfo.dllName, shortDll)) {
+            if (fInfo.funcName == func) {
+                //std::cout << "Excluded Func: " << shortDll << "." << func << "\n";
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
+WFuncInfo* FuncExcludeList::findFunc(const std::string& dllName, const std::string& funcName)
+{
+    for (size_t i = 0; i < funcs.size(); i++)
+    {
+        WFuncInfo& info = funcs[i];
+        if (util::iequals(info.dllName, dllName)
+            && util::iequals(info.funcName, funcName))
+        {
+            return &info;
+        }
+    }
+    return NULL;
+}
+
+bool FuncExcludeList::appendFunc(WFuncInfo& func_info)
+{
+    if (!func_info.isValid()) {
+        return false;
+    }
+    WFuncInfo* found = findFunc(func_info.dllName, func_info.funcName);
+    if (!found) {
+        funcs.push_back(func_info);
+    }
+    return true;
+}
+
+size_t FuncExcludeList::loadList(const char* filename)
+{
+    std::ifstream myfile(filename);
+    if (!myfile.is_open()) {
+        std::cerr << "Coud not open file: " << filename << std::endl;
+        return 0;
+    }
+    const size_t MAX_LINE = 300;
+    char line[MAX_LINE] = { 0 };
+    while (!myfile.eof()) {
+        myfile.getline(line, MAX_LINE);
+
+        // Try to parse as a function
+        WFuncInfo func_info;
+
+        if (func_info.load(line, ';')) {
+            appendFunc(func_info);
+        }
+    }
+    return funcs.size();
+}
+
+//---
+
 WFuncInfo* FuncWatchList::findFunc(const std::string& dllName, const std::string &funcName)
 {
     for (size_t i = 0; i < funcs.size(); i++)
@@ -102,7 +173,7 @@ void FuncWatchList::appendSyscall(WSyscallInfo& syscall_info)
     }
 }
 
-size_t FuncWatchList::loadList(const char* filename)
+size_t FuncWatchList::loadList(const char* filename, FuncExcludeList* exclusions)
 {
     std::ifstream myfile(filename);
     if (!myfile.is_open()) {
@@ -124,6 +195,10 @@ size_t FuncWatchList::loadList(const char* filename)
         // Try to parse as a function
         WFuncInfo func_info;
         if (func_info.load(line, ';')) {
+            if (exclusions && exclusions->contains(func_info.dllName, func_info.funcName)) {
+                //std::cout << ">> Skipping: " << func_info.funcName << std::endl;
+                continue;
+            }
             appendFunc(func_info);
         }
     }

diff --git a/FuncWatch.h b/FuncWatch.h
@@ -49,6 +49,31 @@ struct WSyscallInfo
     size_t paramCount;
 };
 
+class FuncExcludeList {
+public:
+    FuncExcludeList()
+    {
+    }
+
+    ~FuncExcludeList()
+    {
+    }
+
+    bool isEmpty() { return this->funcs.size() > 0 ? false : true; }
+
+    bool contains(const std::string& dll_name, const std::string& func);
+
+    size_t loadList(const char* filename);
+
+    std::vector<WFuncInfo> funcs;
+
+private:
+    bool appendFunc(WFuncInfo& info);
+
+    WFuncInfo* findFunc(const std::string& dllName, const std::string& funcName);
+};
+
+
 class FuncWatchList {
 public:
     FuncWatchList()
@@ -59,7 +84,7 @@ class FuncWatchList {
     {
     }
 
-    size_t loadList(const char* filename);
+    size_t loadList(const char* filename, FuncExcludeList* exclusions);
 
     std::vector<WFuncInfo> funcs;
     std::map<uint32_t, WSyscallInfo> syscalls;

diff --git a/Settings.h b/Settings.h
@@ -90,4 +90,5 @@ class Settings {
 
     SyscallsTable syscallsTable; //Syscalls table: mapping the syscall ID to the function name
     FuncWatchList funcWatch; //List of functions, arguments of which are going to be logged
+    FuncExcludeList excludedFuncs; //List of functions that will NOT be logged
 };
diff --git a/TinyTracer.cpp b/TinyTracer.cpp
@@ -68,6 +68,9 @@ KNOB<std::string> KnobWatchListFile(KNOB_MODE_WRITEONCE, "pintool",
 KNOB<std::string> KnobSyscallsTable(KNOB_MODE_WRITEONCE, "pintool",
     "l", "", "Syscall table: a CSV file mapping a syscall ID (in hex) to a function name");
 
+KNOB<std::string> KnobExcludedListFile(KNOB_MODE_WRITEONCE, "pintool",
+    "x", "", "A list of functions excluded from watching");
+
 /* ===================================================================== */
 // Utilities
 /* ===================================================================== */
@@ -147,16 +150,17 @@ inline ADDRINT getReturnFromTheStack(const CONTEXT* ctx)
 
 VOID _SaveTransitions(const ADDRINT addrFrom, const ADDRINT addrTo, BOOL isIndirect, const CONTEXT* ctx = NULL)
 {
+    const WatchedType fromWType = isWatchedAddress(addrFrom); // is the call from the traced area?
+
     const bool isTargetMy = pInfo.isMyAddress(addrTo);
     const bool isCallerMy = pInfo.isMyAddress(addrFrom);
 
-    const WatchedType fromWType = isWatchedAddress(addrFrom); // is the call from the traced area?
-
     IMG targetModule = IMG_FindByAddress(addrTo);
     IMG callerModule = IMG_FindByAddress(addrFrom);
     const bool isCallerPeModule = IMG_Valid(callerModule);
     const bool isTargetPeModule = IMG_Valid(targetModule);
 
+
     /**
     is it a transition from the traced module to a foreign module?
     */
@@ -167,6 +171,9 @@ VOID _SaveTransitions(const ADDRINT addrFrom, const ADDRINT addrTo, BOOL isIndir
         if (isTargetPeModule) {
             const std::string func = get_func_at(addrTo);
             const std::string dll_name = IMG_Name(targetModule);
+            if (m_Settings.excludedFuncs.contains(dll_name, func)) {
+                return;
+            }
             traceLog.logCall(0, RvaFrom, true, dll_name, func);
         }
         else {
@@ -188,7 +195,9 @@ VOID _SaveTransitions(const ADDRINT addrFrom, const ADDRINT addrTo, BOOL isIndir
         if (isTargetPeModule) { // it is a call to a module
             const std::string func = get_func_at(addrTo);
             const std::string dll_name = IMG_Name(targetModule);
-
+            if (m_Settings.excludedFuncs.contains(dll_name, func)) {
+                return;
+            }
             traceLog.logCall(pageFrom, addrFrom, false, dll_name, func);
         }
         else if (pageFrom != pageTo) // it is a call to another shellcode
@@ -223,6 +232,9 @@ VOID _SaveTransitions(const ADDRINT addrFrom, const ADDRINT addrTo, BOOL isIndir
         if (toWType != WatchedType::NOT_WATCHED) {
             const std::string func = get_func_at(addrTo);
             const std::string dll_name = IMG_Name(targetModule);
+            if (m_Settings.excludedFuncs.contains(dll_name, func)) {
+                return;
+            }
             const ADDRINT pageRet = get_base(returnAddr);
             const ADDRINT RvaFrom = addr_to_rva(addrFrom);
             const ADDRINT base = isTargetMy ? 0 : get_base(addrFrom);
@@ -823,11 +835,19 @@ int main(int argc, char *argv[])
         std::cerr << "Coud not load the INI file: " << iniFilename << std::endl;
         m_Settings.saveINI(iniFilename);
     }
+
+    if (KnobExcludedListFile.Enabled()) {
+        std::string excludedList = KnobExcludedListFile.ValueString();
+        if (excludedList.length()) {
+            m_Settings.excludedFuncs.loadList(excludedList.c_str());
+            std::cout << "Excluded " << m_Settings.excludedFuncs.funcs.size() << " functions\n";
+        }
+    }
 
     if (KnobWatchListFile.Enabled()) {
         std::string watchListFile = KnobWatchListFile.ValueString();
         if (watchListFile.length()) {
-            m_Settings.funcWatch.loadList(watchListFile.c_str());
+            m_Settings.funcWatch.loadList(watchListFile.c_str(), &m_Settings.excludedFuncs);
             std::cout << "Watch " << m_Settings.funcWatch.funcs.size() << " functions\n";
             std::cout << "Watch " << m_Settings.funcWatch.syscalls.size() << " syscalls\n";
         }

diff --git a/install32_64/excluded.txt b/install32_64/excluded.txt
@@ -0,0 +1 @@
+kernelbase;InitializeCriticalSectionEx
diff --git a/install32_64/run_me.bat b/install32_64/run_me.bat
@@ -42,6 +42,8 @@ rem WATCH_BEFORE - a file with a list of functions which's parameters will be lo
 rem The file must be a list of records in a format: [dll_name];[func_name];[parameters_count]
 set WATCH_BEFORE=%PIN_TOOLS_DIR%\params.txt
 
+set EXCLUDED_FUNC=%PIN_TOOLS_DIR%\excluded.txt
+
 rem SYSCALLS_TABLE - a CSV file, mapping syscall ID to a function name. Format: [syscallID:hex],[functionName]
 set SYSCALLS_TABLE=%PIN_TOOLS_DIR%\syscalls.txt
 
@@ -89,8 +91,8 @@ if [%IS_ADMIN%] == [A] (
 
 set ADMIN_CMD=%PIN_TOOLS_DIR%\sudo.vbs
 
-set DLL_CMD=%PIN_DIR%\pin.exe -t %PINTOOL% -m "%TRACED_MODULE%" -o %TAG_FILE% -s %SETTINGS_FILE% -b "%WATCH_BEFORE%" -l "%SYSCALLS_TABLE%" -- "%DLL_LOAD%" "%TARGET_APP%" %DLL_EXPORTS%
-set EXE_CMD=%PIN_DIR%\pin.exe -t %PINTOOL% -m "%TRACED_MODULE%" -o %TAG_FILE% -s %SETTINGS_FILE% -b "%WATCH_BEFORE%" -l "%SYSCALLS_TABLE%" -- "%TARGET_APP%" %EXE_ARGS%
+set DLL_CMD=%PIN_DIR%\pin.exe -t %PINTOOL% -m "%TRACED_MODULE%" -o %TAG_FILE% -s %SETTINGS_FILE% -b "%WATCH_BEFORE%" -x "%EXCLUDED_FUNC%" -l "%SYSCALLS_TABLE%" -- "%DLL_LOAD%" "%TARGET_APP%" %DLL_EXPORTS%
+set EXE_CMD=%PIN_DIR%\pin.exe -t %PINTOOL% -m "%TRACED_MODULE%" -o %TAG_FILE% -s %SETTINGS_FILE% -b "%WATCH_BEFORE%" -x "%EXCLUDED_FUNC%" -l "%SYSCALLS_TABLE%" -- "%TARGET_APP%" %EXE_ARGS%
 
 ;rem "Trace EXE"
 if [%PE_TYPE%] == [exe] (