1.fix mistake extern class when make category;2.support .app as path

flexih · May 16, 2020 · ca6cd5a · ca6cd5a
1 parent 6ee620e
commit ca6cd5a
Show file tree

Hide file tree

Showing 10 changed files with 90 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -22,6 +22,7 @@ See [SnakeKit][1]
 - [x] Unsued classes.
 - [x] Unused protocols.
 - [x] Duplicate selectors.
+- [x] All classes.
 - [x] Selector/Classes/Protocols sort by library, and selector size, if Linkmap file provided.
 - [x] Fast, a 460.6M binary and a 134.3M linkmap file costs 1.62s(3.7 GHz 6-Core Intel Core i5;40 GB 2667 MHz DDR4).
 
@@ -37,12 +38,13 @@ Usage:
   -c, --class        Unused classes
   -p, --protocol     Unused protocoles
   -d, --duplicate    Duplicate selectors
+  -a, --allclass     All Classes
   -l, --linkmap arg  Linkmap file, which has selector size, library name
   -j, --json         Output json format
       --help         Print help
 ```
 
-snake -l path/to/linkmap path/to/binary [-dscp]
+snake -l path/to/linkmap path/to/binary [-dscpa]
 
 ## Example
 bin/snake  -l demo/release/demo-LinkMap-normal-x86\_64.txt demo/release/demo.app/demo -c

diff --git a/bin/snake b/bin/snake
diff --git a/snake/mach-o/Arch.cpp b/snake/mach-o/Arch.cpp
@@ -329,6 +329,7 @@ namespace snake {
                     } else if (auto n = bindinfoSymAt(catRefPtr[i] + sizeof(uintptr_t))) {
                         if (auto i = n->find("_OBJC_CLASS_$_"); i != std::string::npos) {
                             objcClass = ObjCClassForName((n->substr(i + 14)).c_str());
+                            objcClass->externed = true;
                         } else {
                             objcClass = ObjCClassForName(n->c_str());
                         }
@@ -401,7 +402,9 @@ namespace snake {
         std::vector<std::string> keys;
         keys.reserve(allClasses.size());
         for (auto &pair : allClasses) {
-            keys.push_back(pair.first);
+            if (!pair.second.externed) {
+                keys.push_back(pair.first);
+            }
         }
         return keys;
     }
@@ -456,12 +459,13 @@ namespace snake {
         return result;
     }
     std::vector<std::string> Arch::ObjCClassesUnused() const {
+        auto classes = ObjCClasses();
         std::vector<std::string> unused;
-        unused.reserve(allClasses.size());
+        unused.reserve(classes.size());
         auto usedClasses = ObjCClassesUsed();
-        for (auto &c : allClasses) {
-            if (!contains(usedClasses, c.first)) {
-                unused.push_back(c.first);
+        for (auto &c : classes) {
+            if (!contains(usedClasses, c)) {
+                unused.push_back(c);
             }
         }
         return unused;

diff --git a/snake/mach-o/Arch.h b/snake/mach-o/Arch.h
@@ -24,6 +24,8 @@ namespace snake {
         std::vector<std::string> ObjCClassesUnused() const;
         std::vector<std::string> ObjCProtocolsUnused() const;
         std::vector<std::string> ObjCDuplicateSelectors() const;
+        std::set<std::string> ObjCProtocolsUsed() const;
+        std::set<std::string> ObjCClassesUsed() const;
         std::map<std::string, ObjCClass> ObjCSelectorsUnused() const;
     private:
         void parseSections();
@@ -36,8 +38,6 @@ namespace snake {
         void handleSymtab();
         std::vector<std::string> handleDyld();
         void handleObjCSections();
-        std::set<std::string> ObjCProtocolsUsed() const;
-        std::set<std::string> ObjCClassesUsed() const;
         const char *POINTER(uintptr_t x);
         const char *OFFSET(uintptr_t x);
         ObjCClass* ObjCClassForName(const char *name);

diff --git a/snake/mach-o/Bin.cpp b/snake/mach-o/Bin.cpp
@@ -9,6 +9,7 @@
 #include "Bin.h"
 #include <sys/mman.h>
 #include <sys/stat.h>
+#include <sys/types.h>
 #include <unistd.h>
 #include <fcntl.h>
 #include <inttypes.h>
@@ -85,4 +86,22 @@ namespace snake {
         }
         return false;
     }
+    bool Bin::isMachO(std::string &path) {
+        int fd = open(path.c_str(), O_RDONLY, S_IRUSR | S_IWUSR);
+        if (fd < 0) {
+            return false;
+        }
+        uint32_t magic = 0;
+        if (::read(fd, &magic, sizeof(magic)) != sizeof(magic)) {
+            close(fd);
+            return false;
+        }
+        close(fd);
+        return magic == FAT_MAGIC
+            || magic == FAT_CIGAM
+            || magic == FAT_MAGIC_64
+            || magic == FAT_CIGAM_64
+            || magic == MH_MAGIC_64
+            || magic == MH_CIGAM_64;
+    }
 };
diff --git a/snake/mach-o/Bin.h b/snake/mach-o/Bin.h
@@ -18,6 +18,7 @@ namespace snake {
         const Arch* arch() const {
             return archs.empty() ? nullptr : &archs.front();
         }
+        static bool isMachO(std::string &path);
     private:
         bool parseArch(char *p);
         std::vector<Arch> archs;

diff --git a/snake/main.cpp b/snake/main.cpp
@@ -11,19 +11,46 @@
 #include "Output.hpp"
 #include "cxxopts.hpp"
 #include "utility.hpp"
+#include <sys/stat.h>
+#include <unistd.h>
+#include <dirent.h>
 
 using namespace snake;
 
+std::string findMachOPath(std::string &path) {
+    if (auto i = path.find_last_of(".app"); i != std::string::npos && i + 1 == path.size()) {
+        struct stat st;
+        if (stat(path.c_str(), &st) == 0 && S_ISDIR(st.st_mode)) {
+            if (auto dir = opendir(path.c_str())) {
+                struct dirent *dp;
+                while ((dp = readdir(dir)) != nullptr) {
+                    if (strcmp(dp->d_name, ".") == 0 || strcmp(dp->d_name, "..") == 0) {
+                        continue;
+                    }
+                    auto sub = path + "/" + dp->d_name;
+                    if (Bin::isMachO(sub)) {
+                        closedir(dir);
+                        return sub;
+                    }
+                }
+                closedir(dir);
+            }
+        }
+    }
+    return path;
+}
+
 int main(int argc, char * argv[]) {
     auto useJson = false;
-    cxxopts::Options options("snake", "🐍 Snake, Yet Another Mach-O Unused ObjC Selector/Class/Protocol Detector\n");
+    cxxopts::Options options("snake(1.5)", "🐍 Snake, Yet Another Mach-O Unused ObjC Selector/Class/Protocol Detector\n");
     options.custom_help("[-dscp] [-l path] path/to/binary ...");
     options.positional_help("");
     options.add_options()
     ("s,selector", "Unused selectors")
     ("c,class", "Unused classes")
     ("p,protocol", "Unused protocoles")
     ("d,duplicate", "Duplicate selectors")
+    ("a,allclass", "All Classes")
     ("l,linkmap", "Linkmap file, which has selector size, library name", cxxopts::value<std::string>())
     ("i,input", "Mach-O binary", cxxopts::value<std::string>())
     ("j,json", "Output json format", cxxopts::value<bool>(useJson))
@@ -54,9 +81,12 @@ int main(int argc, char * argv[]) {
             function = 'p';
         } else if (result.count("d")) {
             function = 'd';
+        }  else if (result.count("a")) {
+            function = 'a';
         } else {
-            throw cxxopts::option_not_present_exception("-dscp");
+            throw cxxopts::option_not_present_exception("-dscpa");
         }
+        machoPath = findMachOPath(machoPath);
         Bin bin;
         bin.read(machoPath);
         Linkmap linkmap;
@@ -87,6 +117,13 @@ int main(int argc, char * argv[]) {
                 std::cout << Output::raw.duplicatSelectors(duplicateSelectors, linkmap).str() << std::endl;
                 break;
             }
+            case 'a': {
+                auto allClasses = arch->ObjCClasses();
+                for (auto &c : allClasses) {
+                    std::cout << c << std::endl;
+                }
+                break;
+            }
             default:
                 break;
         }

diff --git a/snake/objc/ObjCClass.hpp b/snake/objc/ObjCClass.hpp
@@ -24,6 +24,7 @@ namespace snake {
         std::set<std::string> classMethods;
         std::set<std::string> instanceMethodsDup;
         std::set<std::string> classMethodsDup;
+        bool externed {false};
 
         ObjCClass *catWithName(const char *n) {
             for (auto iter = cats.begin(); iter != cats.end(); ++iter) {

diff --git a/snake/output/Output.cpp b/snake/output/Output.cpp
@@ -561,6 +561,19 @@ namespace snake {
         }
         return stream;
     }
+    std::ostringstream Raw::printNames(std::vector<std::string> &classes, Linkmap &linkmap) {
+        std::ostringstream stream;
+        stream << "Total Count: " << classes.size() << std::endl;
+        stream << std::endl;
+        for (auto &c : classes) {
+            stream << c << std::endl;
+        }
+        return stream;
+    }
+    std::ostringstream Json::printNames(std::vector<std::string> &classes, Linkmap &linkmap) {
+        std::ostringstream stream;
+        return stream;
+    }
     std::ostringstream Json::duplicatSelectors(std::vector<std::string> &selectors, Linkmap &linkmap) {
         std::ostringstream stream;
         return stream;

diff --git a/snake/output/Output.hpp b/snake/output/Output.hpp
@@ -26,6 +26,7 @@ namespace snake {
         virtual std::ostringstream unusedProtocols(std::vector<std::string> &protocolsUnused, Linkmap &linkmap) = 0;
         virtual std::ostringstream unusedClasses(std::vector<std::string> &classesUnused, Linkmap &linkmap) = 0;
         virtual std::ostringstream duplicatSelectors(std::vector<std::string> &selectors, Linkmap &linkmap) = 0;
+        virtual std::ostringstream printNames(std::vector<std::string> &classes, Linkmap &linkmap) = 0;
         //(lib->(class->[meth]))
         typedef std::map<const std::string, std::map<std::string, std::vector<std::string>>> OPLibs;
         const OPLibs interact(std::map<std::string, ObjCClass> &selectorsUnused, Linkmap &linkmap);
@@ -39,6 +40,7 @@ namespace snake {
         std::ostringstream unusedProtocols(std::vector<std::string> &protocolsUnused, Linkmap &linkmap);
         std::ostringstream unusedClasses(std::vector<std::string> &classesUnused, Linkmap &linkmap);
         std::ostringstream duplicatSelectors(std::vector<std::string> &selectors, Linkmap &linkmap);
+        std::ostringstream printNames(std::vector<std::string> &classes, Linkmap &linkmap);
     };
 
     class Json: Output {
@@ -47,6 +49,7 @@ namespace snake {
         std::ostringstream unusedProtocols(std::vector<std::string> &protocolsUnused, Linkmap &linkmap);
         std::ostringstream unusedClasses(std::vector<std::string> &classesUnused, Linkmap &linkmap);
         std::ostringstream duplicatSelectors(std::vector<std::string> &selectors, Linkmap &linkmap);
+        std::ostringstream printNames(std::vector<std::string> &classes, Linkmap &linkmap);
     };
 }