Add extraction of all global variables (PDB, ELF)

Adikso · Sep 29, 2018 · 228545f · 228545f
1 parent cc1fec9
commit 228545f
Show file tree

Hide file tree

Showing 9 changed files with 73 additions and 12 deletions.
diff --git a/src/debugextract.cpp b/src/debugextract.cpp
@@ -29,6 +29,7 @@ int main(int argc, char *argv[]) {
             ("a,all", "Dump all classes (Used with --output dumps compilable header files to directory)")
             ("l,list", "List all classes")
             ("o,output", "Output path", cxxopts::value<std::string>())
+            ("v,vars", "Show all global variables")
             ("positional", "...", cxxopts::value<std::vector<std::string>>());
     options.parse_positional({"input", "class", "positional"});
 
@@ -47,7 +48,7 @@ int main(int argc, char *argv[]) {
     }
 
     auto &v = args["positional"].as<std::vector<std::string>>();
-    if ((v.size() < 2 && !(args.count("all") || args.count("list"))) || v.empty()) {
+    if ((v.size() < 2 && !(args.count("all") || args.count("list") || args.count("vars"))) || v.empty()) {
         std::cout << options.help({"", "display"}) << std::endl;
         return 1;
     }
@@ -80,24 +81,35 @@ int main(int argc, char *argv[]) {
     }
 
     Analyser analyser{config};
+    ClassDumper * dumper = config.json ? (ClassDumper *) new JsonClassDumper : new CodeClassDumper;
 
     if (args.count("list")) {
         list(extractor, analyser);
         return 0;
     }
 
-    std::list<std::string> names;
-    if (args.count("all") > 0) {
-        names = extractor->getTypesList(false);
+    std::vector<Type *> types;
+    std::vector<std::string> dumpOut;
+
+    if (args.count("vars")) {
+        Type * type = new Type("GlobalVariables");
+        type->fields = extractor->getAllGlobalVariables();
+
+        types.push_back(type);
+        dumpOut.push_back(dumper->dump(type, config));
     } else {
-        names = split(v[1], ',');
-    }
+        std::list<std::string> names;
+        if (args.count("all") > 0) {
+            names = extractor->getTypesList(false);
+        } else {
+            names = split(v[1], ',');
+        }
 
-    std::vector<Type *> types = extractor->getTypes(std::move(names));
-    analyser.process(types);
+        types = extractor->getTypes(std::move(names));
+        analyser.process(types);
 
-    ClassDumper * dumper = config.json ? (ClassDumper *) new JsonClassDumper : new CodeClassDumper;
-    std::vector<std::string> dumpOut = dumper->dump(std::move(types), config);
+        dumpOut = dumper->dump(std::move(types), config);
+    }
 
     if (config.toDirectory) {
         std::string extension = (config.json ? "json" : "hpp");

diff --git a/src/extractor/Extractor.hpp b/src/extractor/Extractor.hpp
@@ -21,6 +21,7 @@ class Extractor {
     virtual std::vector<Type *> getTypes(std::list<std::string> typesList) = 0;
 
     virtual std::list<std::string> getTypesList(bool showStructs) = 0;
+    virtual std::vector<Field *> getAllGlobalVariables() = 0;
 
 protected:
     std::list<std::string> allDependentClasses;

diff --git a/src/extractor/dwarf/DWARFExtractor.cpp b/src/extractor/dwarf/DWARFExtractor.cpp
@@ -253,6 +253,10 @@ std::vector<Type *> DWARFExtractor::getTypes(std::list<std::string> typesList) {
     return types;
 }
 
+std::vector<Field *> DWARFExtractor::getAllGlobalVariables() {
+    return std::vector<Field *>();
+}
+
 
 }
 }
diff --git a/src/extractor/dwarf/DWARFExtractor.hpp b/src/extractor/dwarf/DWARFExtractor.hpp
@@ -18,6 +18,7 @@ class DWARFExtractor : public Extractor {
     std::vector<Type *> getTypes(std::list<std::string> typesList) override;
 
     std::list<std::string> getTypesList(bool showStructs) override;
+    std::vector<Field *> getAllGlobalVariables() override;
 
 private:
     ::elf::elf * elf;

diff --git a/src/extractor/elf/ELFExtractor.cpp b/src/extractor/elf/ELFExtractor.cpp
@@ -362,5 +362,30 @@ std::string ELFExtractor::getTypeFromSize(int size) {
     }
 }
 
+std::vector<Field *> ELFExtractor::getAllGlobalVariables() {
+    std::vector<Field *> fields;
+
+    for (auto sym : symtab) {
+        auto &data = sym.get_data();
+        auto demangled = demangleName(sym.get_name(), true);
+
+        if (data.type() != ::elf::stt::object || data.value == 0 || (sym.get_name().rfind("_Z", 0) == 0 && sym.get_name().rfind("_ZN", 0) != 0)) {
+            continue;
+        }
+
+        TypePtr * fieldType = new TypePtr(getTypeFromSize(data.size), false);
+
+        auto name = !demangled.empty() ? demangled : sym.get_name();
+        auto * field = new Field(name, fieldType, 0);
+        field->isStatic = true;
+        field->address = data.value;
+        field->accessibility = Accessibility::PUBLIC;
+
+        fields.push_back(field);
+    }
+
+    return fields;
+}
+
 }
 }
diff --git a/src/extractor/elf/ELFExtractor.hpp b/src/extractor/elf/ELFExtractor.hpp
@@ -15,6 +15,7 @@ class ELFExtractor : public Extractor {
     Type *getType(std::string name) override;
     std::vector<Type *> getTypes(std::list<std::string> typesList) override;
     std::list<std::string> getTypesList(bool showStructs) override;
+    std::vector<Field *> getAllGlobalVariables() override;
 
 private:
     Method *getMethod(::elf::sym * element);

diff --git a/src/extractor/pdb/PDBExtractor.cpp b/src/extractor/pdb/PDBExtractor.cpp
@@ -258,6 +258,22 @@ std::vector<Type *> PDBExtractor::getTypes(std::list<std::string> typesList) {
     return types;
 }
 
+std::vector<Field *> PDBExtractor::getAllGlobalVariables() {
+    std::vector<Field *> fields;
+
+    for (auto &var : *pdb.get_global_variables()) {
+        auto * field = new Field();
+        field->name = var.second.name;
+        field->address = var.second.address;
+        field->accessibility = Accessibility::PUBLIC;
+        field->typePtr = getReturnType(var.second.type_def);
+
+        fields.push_back(field);
+    }
+
+    return fields;
+}
+
 // PDBUniversalType methods
 
 std::string PDBUniversalType::getName() {

diff --git a/src/extractor/pdb/PDBExtractor.hpp b/src/extractor/pdb/PDBExtractor.hpp
@@ -14,6 +14,7 @@ class PDBExtractor : public Extractor {
     Type *getType(std::string name) override;
     std::vector<Type *> getTypes(std::list<std::string> typesList) override;
     std::list<std::string> getTypesList(bool showStructs) override;
+    std::vector<Field *> getAllGlobalVariables();
 
 private:
     retdec::pdbparser::PDBFile pdb;

diff --git a/src/utils/utils.hpp b/src/utils/utils.hpp
@@ -51,7 +51,7 @@ inline int isDirectory(const std::string &path) {
 }
 
 // This is hack
-inline std::string demangleName(const std::string &mangled) {
+inline std::string demangleName(const std::string &mangled, bool leaveFull = false) {
     std::string prefix = "_ZT";
     std::string prefix2 = "_ZN";
     if (mangled.rfind(prefix, 0) == 0 || mangled.rfind(prefix2, 0) == 0) {
@@ -87,7 +87,7 @@ inline std::string demangleName(const std::string &mangled) {
         }
 
         // Leave class name only
-        if (parts.size() > 1 && mangled.rfind(prefix2, 0) == 0) {
+        if (parts.size() > 1 && mangled.rfind(prefix2, 0) == 0 && !leaveFull) {
             parts.pop_back();
         }