Exclude public

rockbruno · Nov 12, 2018 · be459da · be459da
1 parent 04e067e
commit be459da
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 3 deletions.
diff --git a/swiftshield-Sources/AutomaticSwiftShield.swift b/swiftshield-Sources/AutomaticSwiftShield.swift
@@ -5,6 +5,7 @@ class AutomaticSwiftShield: Protector {
     let projectToBuild: String
     let schemeToBuild: String
     let modulesToIgnore: Set<String>
+    let excludePublic: Bool
 
     var isWorkspace: Bool {
         return projectToBuild.hasSuffix(".xcworkspace")
@@ -14,10 +15,12 @@ class AutomaticSwiftShield: Protector {
          projectToBuild: String,
          schemeToBuild: String,
          modulesToIgnore: Set<String>,
-         protectedClassNameSize: Int) {
+         protectedClassNameSize: Int,
+         excludePublic: Bool = false) {
         self.projectToBuild = projectToBuild
         self.schemeToBuild = schemeToBuild
         self.modulesToIgnore = modulesToIgnore
+        self.excludePublic = excludePublic
         super.init(basePath: basePath, protectedClassNameSize: protectedClassNameSize)
         if self.schemeToBuild.isEmpty || self.projectToBuild.isEmpty {
             Logger.log(.helpText)
@@ -108,6 +111,21 @@ extension AutomaticSwiftShield {
         guard let name = dict.getString(key: sourceKit.nameID)?.trueName, let usr = dict.getString(key: sourceKit.usrID) else {
             return nil
         }
+
+        if excludePublic {
+            let attribtuesArray = dict.getDictionaryValue(key: sourceKit.attributesId)
+            if attribtuesArray.isArray {
+                let count = attribtuesArray.getArrayCount()
+                for i in 0..<count {
+                    let attributeDict = attribtuesArray.getArrayValue(index: i)
+                    let attribute = attributeDict.getUUIDString(key: sourceKit.attributeId)
+                    if attribute.split(separator: ".").last == "public" {
+                        return nil
+                    }
+                }
+            }
+        }
+
         guard let protected = obfuscationData.obfuscationDict[name] else {
             let newName = String.random(length: self.protectedClassNameSize, excluding: obfuscationData.allObfuscatedNames)
             obfuscationData.obfuscationDict[name] = newName

diff --git a/swiftshield-Sources/SourceKit+Variant.swift b/swiftshield-Sources/SourceKit+Variant.swift
@@ -22,7 +22,16 @@ extension sourcekitd_variant_t {
         return String( cString: SKApi.sourcekitd_uid_get_string_ptr( uuid! )! )// ?: "NOUUID"
     }
 
-    func getDictionary(key: sourcekitd_uid_t) -> sourcekitd_variant_t {
+    func getDictionaryValue( key: sourcekitd_uid_t ) -> sourcekitd_variant_t {
         return SKApi.sourcekitd_variant_dictionary_get_value(self, key)
     }
+
+    func getArrayValue( index: Int ) -> sourcekitd_variant_t {
+        return SKApi.sourcekitd_variant_array_get_value(self, index)
+    }
+
+    func getArrayCount() -> Int {
+        let count = sourcekitd_variant_array_get_count(self)
+        return Int(count)
+    }
 }
diff --git a/swiftshield-Sources/SourceKit.swift b/swiftshield-Sources/SourceKit.swift
@@ -73,6 +73,8 @@ final class SourceKit {
     lazy var lineID = SKApi.sourcekitd_uid_get_from_cstr("key.line")!
     lazy var colID = SKApi.sourcekitd_uid_get_from_cstr("key.column")!
     lazy var usrID = SKApi.sourcekitd_uid_get_from_cstr("key.usr")!
+    lazy var attributesId = SKApi.sourcekitd_uid_get_from_cstr("key.attributes")!
+    lazy var attributeId = SKApi.sourcekitd_uid_get_from_cstr("key.attribute")!
 
     func array(argv: [String]) -> sourcekitd_object_t {
         let objects = argv.map { SKApi.sourcekitd_request_string_create($0) }

diff --git a/swiftshield-Sources/main.swift b/swiftshield-Sources/main.swift
@@ -37,7 +37,8 @@ if automatic {
     let schemeToBuild = UserDefaults.standard.string(forKey: "automatic-project-scheme") ?? ""
     let projectToBuild = UserDefaults.standard.string(forKey: "automatic-project-file") ?? ""
     let modulesToIgnore = UserDefaults.standard.string(forKey: "ignore-modules")?.components(separatedBy: ",") ?? []
-    protector = AutomaticSwiftShield(basePath: basePath, projectToBuild: projectToBuild, schemeToBuild: schemeToBuild, modulesToIgnore: Set(modulesToIgnore), protectedClassNameSize: protectedClassNameSize)
+    let excludePublic = CommandLine.arguments.contains("-exclude-public")
+    protector = AutomaticSwiftShield(basePath: basePath, projectToBuild: projectToBuild, schemeToBuild: schemeToBuild, modulesToIgnore: Set(modulesToIgnore), protectedClassNameSize: protectedClassNameSize, excludePublic: excludePublic)
 } else {
     let tag = UserDefaults.standard.string(forKey: "tag") ?? "__s"
     protector = ManualSwiftShield(basePath: basePath, tag: tag, protectedClassNameSize: protectedClassNameSize)