From a6e86fd2f22d612bd173f176dbf46dfa82214a3c Mon Sep 17 00:00:00 2001
From: x86byte <111459558+x86byte@users.noreply.github.com>
Date: Mon, 29 Jun 2026 12:53:31 +0100
Subject: [PATCH] add --subs-only flag for reachable-function-only extraction

---
 cfg/builder.cpp | 18 ++++++++++++------
 cfg/builder.hpp |  3 ++-
 main.cpp        | 23 ++++++++++++++++++++---
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/cfg/builder.cpp b/cfg/builder.cpp
index dd05f16..aa4254d 100644
--- a/cfg/builder.cpp
+++ b/cfg/builder.cpp
@@ -17,8 +17,8 @@ static bool regMatch(const string& op, const string& target)
     return norm(op) == norm(target) || op == target;
 }
 
-CFGBuilder::CFGBuilder(Binary* binary, Disassembler* disasm)
-    : m_bin(binary), m_dis(disasm) {}
+CFGBuilder::CFGBuilder(Binary* binary, Disassembler* disasm, bool subsOnly)
+    : m_bin(binary), m_dis(disasm), m_subs_only(subsOnly) {}
 
 bool CFGBuilder::isVisited(addr_t addr)
 {
@@ -487,8 +487,11 @@ CFG CFGBuilder::build()
         process(ep, ep_name);
     }
 
-    for (const auto& exp : m_bin->getExportedFunctions())
-        pending.push(exp.first);
+    if (!m_subs_only)
+    {
+        for (const auto& exp : m_bin->getExportedFunctions())
+            pending.push(exp.first);
+    }
 
     map<addr_t, string> imp_map;
     for (const auto& imp : m_bin->getImportedFunctions())
@@ -521,8 +524,11 @@ CFG CFGBuilder::build()
                         if (!built.count(t)) pending.push(t);
                 }
 
-        for (auto p : m_prologs)
-            if (!built.count(p)) pending.push(p);
+        if (!m_subs_only)
+        {
+            for (auto p : m_prologs)
+                if (!built.count(p)) pending.push(p);
+        }
     }
 
     cfg.indirect_targets = m_targets;
diff --git a/cfg/builder.hpp b/cfg/builder.hpp
index 106d95e..3061071 100644
--- a/cfg/builder.hpp
+++ b/cfg/builder.hpp
@@ -9,7 +9,7 @@
 class CFGBuilder
 {
     public:
-        CFGBuilder(Binary* binary, Disassembler* disasm);
+        CFGBuilder(Binary* binary, Disassembler* disasm, bool subsOnly = false);
         CFG build();
 
     private:
@@ -27,6 +27,7 @@ class CFGBuilder
 
         Binary* m_bin;
         Disassembler* m_dis;
+        bool m_subs_only;
         set<addr_t> m_funcs;
         set<addr_t> m_blocks;
         set<addr_t> m_prologs;
diff --git a/main.cpp b/main.cpp
index d49e411..25f7fed 100644
--- a/main.cpp
+++ b/main.cpp
@@ -7,11 +7,27 @@ int main(int argc, char* argv[])
 {
     if (argc < 2)
     {
-        cerr << "usage: cfgrip <binary>" << endl;
+        cerr << "usage: cfgrip [--subs-only] <binary>" << endl;
         return 1;
     }
 
-    string path = argv[1];
+    bool subsOnly = false;
+    string path;
+
+    for (int i = 1; i < argc; i++)
+    {
+        string arg = argv[i];
+        if (arg == "--subs-only")
+            subsOnly = true;
+        else
+            path = arg;
+    }
+
+    if (path.empty())
+    {
+        cerr << "usage: cfgrip [--subs-only] <binary>" << endl;
+        return 1;
+    }
 
     auto binary = Binary::create(path);
     if (!binary)
@@ -55,7 +71,8 @@ int main(int argc, char* argv[])
         return 1;
     }
 
-    CFGBuilder builder(binary.get(), &dis);
+    CFGBuilder builder(binary.get(), &dis, subsOnly);
+    if (subsOnly) cout << "mode: subs-only" << endl;
     CFG cfg = builder.build();
 
     cout << "functions: " << cfg.functions.size() << endl;