add --subs-only flag for reachable-function-only extraction

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)
+CFGBuilder::CFGBuilder(Binary* binary, Disassembler* disasm, bool subsOnly)
-    : m_bin(binary), m_dis(disasm) {}
+    : 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())
+    if (!m_subs_only)
-        pending.push(exp.first);
+    {
+        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 (!m_subs_only)
-            if (!built.count(p)) pending.push(p);
+        {
+            for (auto p : m_prologs)
+                if (!built.count(p)) pending.push(p);
+        }
     }
 
     cfg.indirect_targets = m_targets;

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;

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;