Implement multi-keybox support (#1)

This commit adds support for using different `keybox.xml` files for different applications. This allows for greater flexibility, enabling the use of distinct cryptographic identities for specific groups of apps.

The `target.txt` configuration now supports a new syntax. A line with a filename in brackets, like `[demo_keybox.xml]`, sets the active keybox for all subsequent packages in the file.

The implementation refactors the configuration and keybox loading logic:

- PkgConfig now parses the new syntax and maps packages to their designated keybox files.

- KeyBoxUtils is updated to dynamically load and cache multiple keybox files on demand, instead of relying on a single global state.

- Certificate generation and hacking functions now use this new configuration to select the correct keybox for each operation based on the application's UID.

README.md

@@ -40,22 +40,38 @@ This file provides the master cryptographic identity for the simulator. It conta
 </AndroidAttestation>
 ```
 
-### Mode Configuration (`target.txt`)
+### Mode and Keybox Configuration (`target.txt`)
 
-TEESimulator currently operates in two primary modes as it transitions towards full emulation. You can control this behavior on a per-package basis.
+TEESimulator currently operates in two primary modes as it transitions towards full emulation.
+You can control the simulation mode and the specific keybox.xml file used on a per-package basis.
+
+#### Mode Suffixes
 
 *   **`!` → Force Generation Mode:** Creates a complete, software-based virtual key. This is the foundation of the full TEE simulation.
 *   **`?` → Force Leaf Hacking Mode:** A legacy mode where a real TEE key is generated, but its attestation certificate is intercepted and modified.
 *   **No symbol → Automatic Mode:** The module selects the most appropriate mode for the device.
 
+#### Multi-Keybox Configuration
+
+You can specify different keybox files for different groups of applications. This is done by adding a line with the filename in square brackets (e.g., [demo_keybox.xml]).
+
+All applications listed after this line will use the specified keybox file, until a new keybox is declared. Applications listed before any custom keybox declaration will use the default `keybox.xml`.
+
 For example:
 ```
-# target.txt
-# Use full generation/simulation for this app
+# These two apps will use the default /data/adb/tricky_store/keybox.xml
 com.google.android.gms!
-
-# Use the legacy leaf hacking mode
 io.github.vvb2060.keyattestation?
+
+# Switch to a different keybox for the following apps.
+# The file must be located at /data/adb/tricky_store/aosp_keybox.xml
+[aosp_keybox.xml]
+com.google.android.gsf
+
+# Switch again to another keybox.
+# The file must be located at /data/adb/tricky_store/demo_keybox.xml
+[demo_keybox.xml]
+org.matrix.demo
 ```
 
 ### Security Patch Level (`security_patch.txt`)

app/src/main/java/io/github/beakthoven/TrickyStoreOSS/CertificateGen.kt

@@ -176,7 +176,7 @@ object CertificateGen {
     }
     
     fun generateChain(uid: Int, params: KeyGenParameters, keyPair: KeyPair, securityLevel: Int = 1): List<ByteArray>? = runCatching {
-        val keybox = getKeyboxForAlgorithm(params.algorithm) ?: return null
+        val keybox = getKeyboxForAlgorithm(uid, params.algorithm) ?: return null
 
         val issuer = X509CertificateHolder(keybox.certificates[0].encoded).subject
         val leaf = buildCertificate(keyPair, keybox, params, issuer, uid, securityLevel)
@@ -234,7 +234,7 @@ object CertificateGen {
         }
         
         val keyPair = generateKeyPair(params) ?: return null
-        val keybox = getKeyboxForAlgorithm(params.algorithm) ?: return null
+        val keybox = getKeyboxForAlgorithm(uid, params.algorithm) ?: return null
         
         val (signingKeyPair, issuer) = if (hasAttestKey) {
             getAttestationKeyInfo(uid, attestKeyDescriptor!!)?.let { 
@@ -267,9 +267,10 @@ object CertificateGen {
         }
     }
 
-    private fun getKeyboxForAlgorithm(algorithm: Int): KeyBox? {
+    private fun getKeyboxForAlgorithm(uid: Int, algorithm: Int): KeyBox? {
         val algorithmName = mapAlgorithmToName(algorithm) ?: return null
-        return KeyBoxUtils.keyboxes[algorithmName]
+        val keyboxFileName = PkgConfig.getKeyboxFileForUid(uid)
+        return KeyBoxUtils.getKeybox(keyboxFileName, algorithmName)
     }
 
     private fun getAttestationKeyInfo(uid: Int, attestKeyDescriptor: KeyDescriptor): Pair<KeyPair, X500Name>? {
@@ -468,4 +469,4 @@ object CertificateGen {
 
         return DEROctetString(DERSequence(applicationIdArray).encoded)
     }
-}
+}

app/src/main/java/io/github/beakthoven/TrickyStoreOSS/CertificateHack.kt

@@ -5,6 +5,7 @@
 
 package io.github.beakthoven.TrickyStoreOSS
 
+import io.github.beakthoven.TrickyStoreOSS.config.PkgConfig
 import io.github.beakthoven.TrickyStoreOSS.logging.Logger
 import org.bouncycastle.asn1.ASN1Boolean
 import org.bouncycastle.asn1.ASN1Encodable
@@ -49,7 +50,7 @@ object CertificateHack {
         leafAlgorithms.clear()
     }
     
-    fun hackCertificateChain(certificateChain: Array<Certificate>?): Array<Certificate> {
+    fun hackCertificateChain(certificateChain: Array<Certificate>?, uid: Int): Array<Certificate> {
         if (certificateChain == null) {
             throw UnsupportedOperationException("Certificate chain is null!")
         }
@@ -80,8 +81,9 @@ object CertificateHack {
                 }
             }
             
-            val keybox = KeyBoxUtils.keyboxes[leaf.publicKey.algorithm]
-                ?: throw UnsupportedOperationException("Unsupported algorithm: ${leaf.publicKey.algorithm}")
+            val keyboxFileName = PkgConfig.getKeyboxFileForUid(uid)
+            val algorithmName = leaf.publicKey.algorithm
+            val keybox = KeyBoxUtils.getKeybox(keyboxFileName, algorithmName) ?: throw UnsupportedOperationException("Unsupported algorithm '$algorithmName' in keybox '$keyboxFileName'")
             
             val certificates = LinkedList(keybox.certificates)
             val builder = X509v3CertificateBuilder(
@@ -107,7 +109,7 @@ object CertificateHack {
             certificates.addFirst(JcaX509CertificateConverter().getCertificate(builder.build(signer)))
             certificates.toTypedArray()
         } catch (t: Throwable) {
-            Logger.e("Failed to hack certificate chain", t)
+            Logger.e("Failed to hack certificate chain for uid=$uid", t)
             certificateChain
         }
     }
@@ -122,12 +124,13 @@ object CertificateHack {
             val algorithm = leafAlgorithms.remove(key)
                 ?: throw UnsupportedOperationException("No algorithm found for key $key")
             
-            val keybox = KeyBoxUtils.keyboxes[algorithm]
-                ?: throw UnsupportedOperationException("Unsupported algorithm: $algorithm")
+            val keyboxFileName = PkgConfig.getKeyboxFileForUid(uid)
+            val keybox = KeyBoxUtils.getKeybox(keyboxFileName, algorithm)
+                ?: throw UnsupportedOperationException("Unsupported algorithm '$algorithm' in keybox '$keyboxFileName'")
             
             CertificateUtils.run { keybox.certificates.toByteArray() } ?: caList
         } catch (t: Throwable) {
-            Logger.e("Failed to hack CA certificate chain", t)
+            Logger.e("Failed to hack CA certificate chain for uid=$uid", t)
             caList
         }
     }
@@ -166,8 +169,10 @@ object CertificateHack {
                 }
             }
             
-            val keybox = KeyBoxUtils.keyboxes[leaf.publicKey.algorithm]
-                ?: throw UnsupportedOperationException("Unsupported algorithm: ${leaf.publicKey.algorithm}")
+            val keyboxFileName = PkgConfig.getKeyboxFileForUid(uid)
+            val algorithmName = leaf.publicKey.algorithm
+            val keybox = KeyBoxUtils.getKeybox(keyboxFileName, algorithmName)
+                ?: throw UnsupportedOperationException("Unsupported algorithm '$algorithmName' in keybox '$keyboxFileName'")
             
             val builder = X509v3CertificateBuilder(
                 X509CertificateHolder(keybox.certificates[0].encoded).subject,
@@ -191,7 +196,7 @@ object CertificateHack {
 
             JcaX509CertificateConverter().getCertificate(builder.build(signer)).encoded
         } catch (t: Throwable) {
-            Logger.e("Failed to hack user certificate", t)
+            Logger.e("Failed to hack user certificate for uid=$uid", t)
             certificate
         }
     }
@@ -238,4 +243,4 @@ object CertificateHack {
         return Extension(ATTESTATION_OID, false, hackedSequenceOctets)
     }
 
-}
+}

app/src/main/java/io/github/beakthoven/TrickyStoreOSS/XmlParser.kt

@@ -8,10 +8,12 @@ package io.github.beakthoven.TrickyStoreOSS
 import android.security.keystore.KeyProperties
 import io.github.beakthoven.TrickyStoreOSS.CertificateGen.KeyBox
 import io.github.beakthoven.TrickyStoreOSS.CertificateHack.clearLeafAlgorithms
+import io.github.beakthoven.TrickyStoreOSS.config.PkgConfig
 import io.github.beakthoven.TrickyStoreOSS.logging.Logger
 import org.xmlpull.v1.XmlPullParser
 import org.xmlpull.v1.XmlPullParserException
 import org.xmlpull.v1.XmlPullParserFactory
+import java.io.File
 import java.io.IOException
 import java.io.StringReader
 import java.security.cert.Certificate
@@ -159,12 +161,68 @@ class XmlParser(private val xmlContent: String) {
 }
 
 object KeyBoxUtils {
-    val keyboxes = ConcurrentHashMap<String, KeyBox>()
+    private val loadedKeyboxFiles = ConcurrentHashMap<String, ConcurrentHashMap<String, KeyBox>>()
 
-    fun hasKeyboxes(): Boolean = keyboxes.isNotEmpty()
+    /**
+     * The primary public function to get a specific KeyBox for a given algorithm and file.
+     * It will load and cache the file on demand if it hasn't been seen before.
+     *
+     * @param keyboxFileName The simple name of the keybox file (e.g., "keybox.xml").
+     * @param algorithm The algorithm key (e.g., KeyProperties.KEY_ALGORITHM_EC).
+     * @return The requested KeyBox, or null if not found in the specified file.
+     */
+    fun getKeybox(keyboxFileName: String, algorithm: String): KeyBox? {
+        val keyboxesForFile = loadedKeyboxFiles.getOrPut(keyboxFileName) {
+            // If this file is not in our cache, load it now.
+            readFromFile(keyboxFileName)
+        }
+        Logger.i("Retriving keybox $keyboxFileName [$algorithm]")
+        return keyboxesForFile[algorithm]
+    }
 
+    private fun readFromFile(fileName: String): ConcurrentHashMap<String, KeyBox> {
+        val filePath = File(PkgConfig.CONFIG_PATH, fileName)
+        Logger.i("Loading keybox file: $filePath")
+
+        val keyboxes = ConcurrentHashMap<String, KeyBox>()
+
+        if (!filePath.exists()) {
+            Logger.e("Keybox file not found: $filePath")
+            return keyboxes // Return an empty map if file doesn't exist
+        }
+
+        try {
+            val xmlData = filePath.readText()
+            val xmlParser = XmlParser(xmlData.sanitizeXml())
+
+            val numberOfKeyboxesResult = xmlParser.obtainPath("AndroidAttestation.NumberOfKeyboxes")
+            val numberOfKeyboxes = when (numberOfKeyboxesResult) {
+                is XmlParser.ParseResult.Success -> numberOfKeyboxesResult.attributes["text"]?.toIntOrNull() ?: 1
+                is XmlParser.ParseResult.Error -> throw Exception(numberOfKeyboxesResult.message, numberOfKeyboxesResult.cause)
+            }
+
+            repeat(numberOfKeyboxes) { i ->
+                val (algorithmName, keyBox) = processKeybox(xmlParser, i)
+                keyboxes[algorithmName] = keyBox
+            }
+
+            Logger.i("Successfully loaded ${keyboxes.size} keyboxes from $fileName")
+        } catch (t: Throwable) {
+            Logger.e("Error loading XML file ($fileName)", t)
+        }
+
+        return keyboxes
+    }
+
+    fun hasKeyboxes(): Boolean = loadedKeyboxFiles.isNotEmpty() && loadedKeyboxFiles.values.any { it.isNotEmpty() }
+
+    // This function is now deprecated and should be removed. We keep it for now to show the transition.
+    // Its logic is now inside readFromFile.
+    @Deprecated("Use getKeybox(fileName, algorithm) instead for dynamic loading.")
     fun readFromXml(xmlData: String?) {
-        keyboxes.clear()
+        // The old global state is gone. This function's logic is now in readFromFile.
+        // We could make this load into a default keybox for backward compatibility if needed.
+        loadedKeyboxFiles.clear()
         clearLeafAlgorithms()
         
         if (xmlData == null) {
@@ -209,7 +267,7 @@ object KeyBoxUtils {
         return content.trimEnd()
     }
 
-    private fun processKeybox(xmlParser: XmlParser, index: Int) {
+    private fun processKeybox(xmlParser: XmlParser, index: Int): Pair<String, KeyBox> {
         try {
             val algorithmResult = xmlParser.obtainPath("AndroidAttestation.Keybox.Key[$index]")
             val keyboxAlgorithm = when (algorithmResult) {
@@ -264,11 +322,11 @@ object KeyBoxUtils {
                 else -> keyboxAlgorithm
             }
             
-            keyboxes[algorithmName] = KeyBox(pemKeyPair, keyPair, certificateChain)
+            return algorithmName to KeyBox(pemKeyPair, keyPair, certificateChain)
             
         } catch (t: Throwable) {
             Logger.e("Error processing keybox $index", t)
             throw t
         }
     }
-}
+}

app/src/main/java/io/github/beakthoven/TrickyStoreOSS/config/Config.kt

@@ -21,6 +21,18 @@ object PkgConfig {
     private val generatePackages = mutableSetOf<String>()
     private val packageModes = mutableMapOf<String, Mode>()
 
+    private val packageKeyboxes = mutableMapOf<String, String>()
+    private val keyboxRegex = Regex("^\\[([a-zA-Z0-9_.-]+\\.xml)]$")
+    private const val DEFAULT_KEYBOX_FILE = "keybox.xml"
+
+    fun getKeyboxFileForUid(callingUid: Int): String = runCatching {
+        val ps = getPm()?.getPackagesForUid(callingUid) ?: return DEFAULT_KEYBOX_FILE
+        for (pkg in ps) {
+            packageKeyboxes[pkg]?.let { return it }
+        }
+        return DEFAULT_KEYBOX_FILE
+    }.getOrDefault(DEFAULT_KEYBOX_FILE)
+
     enum class Mode {
         AUTO, LEAF_HACK, GENERATE
     }
@@ -29,41 +41,58 @@ object PkgConfig {
         hackPackages.clear()
         generatePackages.clear()
         packageModes.clear()
-        f?.readLines()?.forEach {
-            if (it.isNotBlank() && !it.startsWith("#")) {
-                val n = it.trim()
-                when {
-                    n.endsWith("!") -> {
-                        val pkg = n.removeSuffix("!").trim()
-                        generatePackages.add(pkg)
-                        packageModes[pkg] = Mode.GENERATE
-                    }
-                    n.endsWith("?") -> {
-                        val pkg = n.removeSuffix("?").trim()
-                        hackPackages.add(pkg)
-                        packageModes[pkg] = Mode.LEAF_HACK
-                    }
-                    else -> {
-                        // Auto mode
-                        packageModes[n] = Mode.AUTO
-                    }
+        packageKeyboxes.clear()
+
+        var currentKeyboxFile = DEFAULT_KEYBOX_FILE
+
+        f?.readLines()?.forEach { line ->
+            val n = line.trim()
+            if (n.isBlank() || n.startsWith("#")) {
+                return@forEach // Skip comments and empty lines
+            }
+
+            val matchResult = keyboxRegex.find(n)
+            if (matchResult != null) {
+                currentKeyboxFile = matchResult.groupValues[1]
+                Logger.i("Switched to keybox file: $currentKeyboxFile for subsequent packages")
+                return@forEach
+            }
+
+            when {
+                n.endsWith("!") -> {
+                    val pkg = n.removeSuffix("!").trim()
+                    generatePackages.add(pkg)
+                    packageModes[pkg] = Mode.GENERATE
+                    packageKeyboxes[pkg] = currentKeyboxFile
+                }
+                n.endsWith("?") -> {
+                    val pkg = n.removeSuffix("?").trim()
+                    hackPackages.add(pkg)
+                    packageModes[pkg] = Mode.LEAF_HACK
+                    packageKeyboxes[pkg] = currentKeyboxFile
+                }
+                else -> {
+                    // Auto mode
+                    packageModes[n] = Mode.AUTO
+                    packageKeyboxes[n] = currentKeyboxFile
                 }
             }
         }
-        Logger.i("update hack packages: $hackPackages, generate packages=$generatePackages, packageModes=$packageModes")
+        Logger.i("update hack packages: $hackPackages, generate packages=$generatePackages, packageModes=$packageModes, , packageKeyboxes=$packageKeyboxes")
     }.onFailure {
         Logger.e("failed to update target files", it)
     }
 
+    // This function is now deprecated in favor of a more dynamic approach, but kept for simplicity.
+    // The key logic is now in KeyBoxUtils which will be called from the interceptors.
     private fun updateKeyBox(f: File?) = runCatching {
         KeyBoxUtils.readFromXml(f?.readText())
     }.onFailure {
         Logger.e("failed to update keybox", it)
     }
 
-    private const val CONFIG_PATH = "/data/adb/tricky_store"
+    const val CONFIG_PATH = "/data/adb/tricky_store"
     private const val TARGET_FILE = "target.txt"
-    private const val KEYBOX_FILE = "keybox.xml"
     private const val TEE_STATUS_FILE = "tee_status"
     private const val PATCHLEVEL_FILE = "security_patch.txt"
     private val root = File(CONFIG_PATH)
@@ -100,10 +129,15 @@ object PkgConfig {
                 DELETE, MOVED_FROM -> null
                 else -> return
             }
-            when (path) {
-                TARGET_FILE -> updateTargetPackages(f)
-                KEYBOX_FILE -> updateKeyBox(f)
-                PATCHLEVEL_FILE -> updatePatchLevel(f)
+            when {
+                path == TARGET_FILE -> updateTargetPackages(f)
+                path.endsWith(".xml") -> {
+                    // This is a simplification. A more robust solution would be to reload the specific keybox if it's in use.
+                    // For now, we assume any XML change might affect the active keyboxes, prompting a reload where needed.
+                    // The main logic for loading is now handled dynamically in KeyBoxUtils.
+                    Logger.i("Keybox file $path changed. It will be re-read on next use.")
+                }
+                path == PATCHLEVEL_FILE -> updatePatchLevel(f)
             }
         }
     }
@@ -116,9 +150,9 @@ object PkgConfig {
         } else {
             Logger.e("target.txt file not found, please put it to $scope !")
         }
-        val keybox = File(root, KEYBOX_FILE)
+        val keybox = File(root, DEFAULT_KEYBOX_FILE)
         if (!keybox.exists()) {
-            Logger.e("keybox file not found, please put it to $keybox !")
+            Logger.e("default keybox file not found, please put it to $keybox !")
         } else {
             updateKeyBox(keybox)
         }

app/src/main/java/io/github/beakthoven/TrickyStoreOSS/interceptors/Keystore2Interceptor.kt

@@ -150,7 +150,7 @@ object Keystore2Interceptor : BaseKeystoreInterceptor() {
                 if (response != null) {
                     val chain = CertificateUtils.run { response.getCertificateChain() }
                 if (chain != null) {
-                        val newChain = CertificateHack.hackCertificateChain(chain)
+                        val newChain = CertificateHack.hackCertificateChain(chain, callingUid)
                         response.putCertificateChain(newChain).getOrThrow()
                         Logger.i("Hacked certificate for uid=$callingUid")
                         return createTypedObjectReply(response)
@@ -167,4 +167,4 @@ object Keystore2Interceptor : BaseKeystoreInterceptor() {
         }
         return Skip
     }
-}
+}