1 /*
2  * Copyright (C) 2022 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.trust.test.lib
18 
19 import android.app.trust.TrustManager
20 import android.content.ComponentName
21 import android.content.Context
22 import android.trust.BaseTrustAgentService
23 import android.util.Log
24 import androidx.test.core.app.ApplicationProvider.getApplicationContext
25 import com.android.internal.widget.LockPatternUtils
26 import com.google.common.truth.Truth.assertWithMessage
27 import org.junit.rules.TestRule
28 import org.junit.runner.Description
29 import org.junit.runners.model.Statement
30 import kotlin.reflect.KClass
31 
32 /**
33  * Enables a trust agent and causes the system service to bind to it.
34  *
35  * The enabled agent can be accessed during the test via the [agent] property.
36  *
37  * @constructor Creates the rule. Do not use; instead, use [invoke].
38  */
39 class TrustAgentRule<T : BaseTrustAgentService>(
40     private val serviceClass: KClass<T>
41 ) : TestRule {
42     private val context: Context = getApplicationContext()
43     private val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
44     private val lockPatternUtils = LockPatternUtils(context)
45 
46     val agent get() = BaseTrustAgentService.instance(serviceClass) as T
47 
48     override fun apply(base: Statement, description: Description) = object : Statement() {
49         override fun evaluate() {
50             verifyTrustServiceRunning()
51             unlockDeviceWithCredential()
52             enableTrustAgent()
53 
54             try {
55                 verifyAgentIsRunning()
56                 base.evaluate()
57             } finally {
58                 disableTrustAgent()
59             }
60         }
61     }
62 
63     private fun verifyTrustServiceRunning() {
64         assertWithMessage("Trust service is not running").that(trustManager).isNotNull()
65     }
66 
67     private fun unlockDeviceWithCredential() {
68         Log.d(TAG, "Unlocking device with credential")
69         trustManager.reportUnlockAttempt(true, context.userId)
70     }
71 
72     private fun enableTrustAgent() {
73         val componentName = ComponentName(context, serviceClass.java)
74         val userId = context.userId
75         Log.i(TAG, "Enabling trust agent ${componentName.flattenToString()} for user $userId")
76         val agents = mutableListOf(componentName)
77             .plus(lockPatternUtils.getEnabledTrustAgents(userId))
78             .distinct()
79         lockPatternUtils.setEnabledTrustAgents(agents, userId)
80     }
81 
82     private fun verifyAgentIsRunning() {
83         wait("${serviceClass.simpleName} to be running") {
84             BaseTrustAgentService.instance(serviceClass) != null
85         }
86     }
87 
88     private fun disableTrustAgent() {
89         val componentName = ComponentName(context, serviceClass.java)
90         val userId = context.userId
91         Log.i(TAG, "Disabling trust agent ${componentName.flattenToString()} for user $userId")
92         val agents = lockPatternUtils.getEnabledTrustAgents(userId).toMutableList()
93             .distinct()
94             .minus(componentName)
95         lockPatternUtils.setEnabledTrustAgents(agents, userId)
96     }
97 
98     companion object {
99         /**
100          * Creates a new rule for the specified agent class. Example usage:
101          * ```
102          *   @get:Rule val rule = TrustAgentRule<MyTestAgent>()
103          * ```
104          */
105         inline operator fun <reified T : BaseTrustAgentService> invoke() =
106             TrustAgentRule(T::class)
107 
108         private const val TAG = "TrustAgentRule"
109     }
110 }
111