From 888c538deae6cdbc193e56dc375336a2452ec3e4 Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@gmail.com>
Date: Sat, 24 Aug 2024 06:30:53 -0800
Subject: [PATCH] [FEATURE] Allow logging all instructions in a defined range

---
 Settings.cpp   | 12 ++++++++++
 Settings.h     |  5 ++++-
 TinyTracer.cpp | 61 +++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/Settings.cpp b/Settings.cpp
index 898a2cc..ff5b1a1 100644
--- a/Settings.cpp
+++ b/Settings.cpp
@@ -24,6 +24,8 @@
 #define KEY_HYPREV_SET                  "EMULATE_HYPERV"
 #define KEY_STOP_OFFSET_TIME            "STOP_OFFSET_TIME"
 #define KEY_EMULATE_SINGLE_STEP         "EMULATE_SINGLE_STEP"
+#define KEY_DISASM_START                "DISASM_START"
+#define KEY_DISASM_STOP                 "DISASM_STOP"
 
 t_shellc_options ConvertShcOption(int value)
 {
@@ -189,6 +191,14 @@ bool fillSettings(Settings &s, std::string line)
         s.emulateSingleStep = loadBoolean(valStr, s.emulateSingleStep);
         isFilled = true;
     }
+    if (util::iequals(valName, KEY_DISASM_START)) {
+        s.disasmStart = util::loadInt(valStr, true);
+        isFilled = true;
+    }
+    if (util::iequals(valName, KEY_DISASM_STOP)) {
+        s.disasmStop = util::loadInt(valStr, true);
+        isFilled = true;
+    }
     return isFilled;
 }
 
@@ -246,6 +256,8 @@ bool Settings::saveINI(const std::string &filename)
     myfile << KEY_HYPREV_SET << DELIM << this->isHyperVSet << "\r\n";
     myfile << KEY_STOP_OFFSET_TIME << DELIM << std::dec << this->stopOffsetTime << "\r\n";
     myfile << KEY_EMULATE_SINGLE_STEP << DELIM << std::dec << this->emulateSingleStep << "\r\n";
+    myfile << KEY_DISASM_START << DELIM << std::hex << this->disasmStart << "\r\n";
+    myfile << KEY_DISASM_STOP << DELIM << std::hex << this->disasmStop << "\r\n";
     myfile.close();
     return true;
 }
diff --git a/Settings.h b/Settings.h
index 0c203f2..db625d5 100644
--- a/Settings.h
+++ b/Settings.h
@@ -98,7 +98,8 @@ class Settings {
         antivm(WATCH_DISABLED),
         useDebugSym(false),
         isHyperVSet(false),
-        emulateSingleStep(true)
+        emulateSingleStep(true),
+        disasmStart(0), disasmStop(0)
     {
     }
 
@@ -123,6 +124,8 @@ class Settings {
     bool useDebugSym;
     bool isHyperVSet; // emulate HyperV via CPUID (it changes execution path of some protectors, i.e. VMProtect). Works when antivm is enabled. 
     bool emulateSingleStep; // If the Trap Flag is set, throw a SINGLE_STEP exception emulating the typical behavior. Works when antidebug is enabled. 
+    int disasmStart;
+    int disasmStop;
 
     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
diff --git a/TinyTracer.cpp b/TinyTracer.cpp
index 32d09b2..5f19e5c 100644
--- a/TinyTracer.cpp
+++ b/TinyTracer.cpp
@@ -138,7 +138,7 @@ WatchedType isWatchedAddress(const ADDRINT Address)
             return WatchedType::WATCHED_SHELLCODE;
         }
     }
-    return WatchedType::NOT_WATCHED;;
+    return WatchedType::NOT_WATCHED;
 }
 
 /* ===================================================================== */
@@ -843,12 +843,71 @@ VOID MonitorFunctionArgs(IMG Image, const WFuncInfo &funcInfo)
     RTN_Close(funcRtn);
 }
 
+
+VOID LogInstruction(const CONTEXT* ctxt, THREADID tid, VOID *str)
+{
+    std::string* strPtr = (std::string*)str;
+    if (!strPtr) return;
+
+    PinLocker locker;
+    static BOOL traceStarted = FALSE;
+
+    const ADDRINT Address = (ADDRINT)PIN_GetContextReg(ctxt, REG_INST_PTR);
+    const WatchedType wType = isWatchedAddress(Address);
+
+    if (wType == WatchedType::NOT_WATCHED) {
+        return;
+    }
+
+    ADDRINT rva = UNKNOWN_ADDR;
+    ADDRINT base = UNKNOWN_ADDR;
+    if (wType == WatchedType::WATCHED_MY_MODULE) {
+        rva = addr_to_rva(Address); // convert to RVA
+        base = 0;
+        if (rva == m_Settings.disasmStart) {
+            traceStarted = TRUE;
+        }
+    }
+    if (!traceStarted) {
+        return;
+    }
+    if (wType == WatchedType::WATCHED_SHELLCODE) {
+        base = query_region_base(Address);
+        rva = Address - base;
+    }
+    if (base != UNKNOWN_ADDR && rva != UNKNOWN_ADDR) {
+        std::stringstream ss;
+        ss << strPtr->c_str();
+        if (!base && rva == m_Settings.disasmStart) {
+            ss << " # disasm start";
+        }
+        if (!base && rva == m_Settings.disasmStop) {
+            ss << " # disasm end";
+        }
+        traceLog.logInstruction(base, rva, ss.str());
+    }
+
+    if (wType == WatchedType::WATCHED_MY_MODULE && rva == m_Settings.disasmStop) {
+        traceStarted = FALSE;
+    }
+}
+
+
 /* ===================================================================== */
 // Instrumentation callbacks
 /* ===================================================================== */
 
 VOID InstrumentInstruction(INS ins, VOID *v)
 {
+    if (m_Settings.disasmStart || m_Settings.disasmStop)
+    INS_InsertCall(
+        ins,
+        IPOINT_BEFORE, (AFUNPTR)LogInstruction,
+        IARG_CONTEXT,
+        IARG_THREAD_ID,
+        IARG_PTR, new std::string(INS_Disassemble(ins)),
+        IARG_END
+    );
     if (m_Settings.stopOffsets.size() > 0 && m_Settings.stopOffsetTime) {
         INS_InsertCall(
             ins,