From 5a0d13b34d86ed70f379e7a00ad086c1df7e0323 Mon Sep 17 00:00:00 2001
From: Wang Han <416810799@qq.com>
Date: Sun, 14 Jul 2024 16:11:02 +0800
Subject: [PATCH] Allow specify injection targets

---
 app/src/main/cpp/main.cpp                | 60 +++++++++++++++++++++---
 build.gradle.kts                         |  4 +-
 gradle/wrapper/gradle-wrapper.properties |  2 +-
 module/customize.sh                      | 15 +++++-
 module/module.prop                       |  2 +-
 5 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/app/src/main/cpp/main.cpp b/app/src/main/cpp/main.cpp
index d73cc99d..aa9c3691 100644
--- a/app/src/main/cpp/main.cpp
+++ b/app/src/main/cpp/main.cpp
@@ -2,14 +2,15 @@
 #include <sys/system_properties.h>
 #include <string>
 #include <vector>
+#include <sstream>
 #include <unistd.h>
 #include "zygisk.hpp"
 
 #define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG, "PIF", __VA_ARGS__)
 
 #define CLASSES_DEX "/data/adb/modules/keystoreinjection/classes.dex"
-
-#define KEYBOX_FILE_PATH "/data/adb/keybox.xml"
+#define APPLIST_FILE_PATH "/data/adb/keystoreinjection/targetlist"
+#define KEYBOX_FILE_PATH "/data/adb/keystoreinjection/keybox.xml"
 
 ssize_t xread(int fd, void *buffer, size_t count) {
     ssize_t total = 0;
@@ -37,6 +38,18 @@ ssize_t xwrite(int fd, void *buffer, size_t count) {
     return total;
 }
 
+std::vector<std::string> split(const std::string &strTotal) {
+    std::vector<std::string> vecResult;
+    std::istringstream iss(strTotal);
+    std::string token;
+
+    while (std::getline(iss, token, '\n')) {
+        vecResult.push_back("/" + token);
+    }
+
+    return std::move(vecResult);
+}
+
 class KeystoreInjection : public zygisk::ModuleBase {
 public:
     void onLoad(zygisk::Api *api, JNIEnv *env) override {
@@ -55,13 +68,41 @@ class KeystoreInjection : public zygisk::ModuleBase {
 
         std::string dir(rawDir);
         env->ReleaseStringUTFChars(args->app_data_dir, rawDir);
-        if (!dir.ends_with("/io.github.vvb2060.keyattestation")) {
+
+        int fd = api->connectCompanion();
+
+        // Check target app and return earlier if not related
+        long applistSize = 0;
+        xread(fd, &applistSize, sizeof(long));
+
+        if (applistSize < 1) {
+            close(fd);
             return;
         }
 
-        long dexSize = 0, xmlSize = 0;
+        std::vector<uint8_t> applistVector;
+        applistVector.resize(applistSize);
+        xread(fd, applistVector.data(), applistSize);
+
+        // Generate split app list
+        std::string applist(applistVector.begin(), applistVector.end());
+        std::vector<std::string> splitlist= split(applist);
+
+        // Find target app
+        bool found = false;
+        for (const std::string& app : splitlist) {
+            if (dir.ends_with(app)) {
+                found = true;
+                break;
+            }
+        }
 
-        int fd = api->connectCompanion();
+        if (!found) {
+            close(fd);
+            return;
+        }
+
+        long dexSize = 0, xmlSize = 0;
 
         xread(fd, &dexSize, sizeof(long));
         xread(fd, &xmlSize, sizeof(long));
@@ -153,15 +194,20 @@ static std::vector<uint8_t> readFile(const char *path) {
 }
 
 static void companion(int fd) {
+    std::vector<uint8_t> applistVector, dexVector, xmlVector;
 
-    std::vector<uint8_t> dexVector, xmlVector;
-
+    applistVector = readFile(APPLIST_FILE_PATH);
     dexVector = readFile(CLASSES_DEX);
     xmlVector = readFile(KEYBOX_FILE_PATH);
 
+    long applistSize = applistVector.size();
     long dexSize = dexVector.size();
     long xmlSize = xmlVector.size();
 
+    xwrite(fd, &applistSize, sizeof(long));
+    // Write applist earlier, so we can avoid reading dex for unrelated apps
+    xwrite(fd, applistVector.data(), applistSize);
+
     xwrite(fd, &dexSize, sizeof(long));
     xwrite(fd, &xmlSize, sizeof(long));
 
diff --git a/build.gradle.kts b/build.gradle.kts
index 1d39658e..8d98714c 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -1,6 +1,6 @@
 buildscript {
-    val agp_version by extra("8.5.0")
+    val agp_version by extra("8.5.1")
 }
 plugins {
-    id("com.android.application") version "8.5.0" apply false
+    id("com.android.application") version "8.5.1" apply false
 }
\ No newline at end of file
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index c0bf2d54..731c8978 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,7 +1,7 @@
 #Thu Apr 18 12:59:17 CEST 2024
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-8.8-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.9-bin.zip
 networkTimeout=10000
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/module/customize.sh b/module/customize.sh
index 0b42dda7..907a5c24 100644
--- a/module/customize.sh
+++ b/module/customize.sh
@@ -1,4 +1,17 @@
+LIST="/data/adb/keystoreinjection/targetlist"
+
 # Error on < Android 14.
-if [ "$API" -lt 30]; then
+if [ "$API" -lt 30 ]; then
     abort "- !!! You can't use this module on Android < 11"
 fi
+
+mkdir -p /data/adb/keystoreinjection
+if [ ! -e "$LIST" ]; then
+    echo "io.github.vvb2060.keyattestation" > "$LIST"
+fi
+
+ui_print "***********************************************************************"
+ui_print "- Please move keybox to /data/adb/keystoreinjection/keybox.xml"
+ui_print "- Please define target apps in /data/adb/keystoreinjection/targetlist"
+ui_print "-     Format: one app per line"
+ui_print "***********************************************************************"
diff --git a/module/module.prop b/module/module.prop
index 1b909731..a372a028 100644
--- a/module/module.prop
+++ b/module/module.prop
@@ -3,4 +3,4 @@ name=Keystore Injection
 version=v0.1.0
 versionCode=10
 author=aviraxp
-description=Generate a valid certificate chain by injecting custom Keystore provider.
+description=Generate a valid certificate chain from scratch by injecting custom Keystore provider.