1 /*
<lambda>null2  * 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.content.Context
20 import android.util.Log
21 import android.view.KeyEvent
22 import android.view.WindowManagerGlobal
23 import androidx.test.core.app.ApplicationProvider.getApplicationContext
24 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
25 import androidx.test.uiautomator.UiDevice
26 import com.android.internal.widget.LockPatternUtils
27 import com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_NOT_REQUIRED
28 import com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN
29 import com.android.internal.widget.LockscreenCredential
30 import com.google.common.truth.Truth.assertWithMessage
31 import org.junit.rules.TestRule
32 import org.junit.runner.Description
33 import org.junit.runners.model.Statement
34 
35 /**
36  * Sets a screen lock on the device for the duration of the test.
37  *
38  * @param requireStrongAuth Whether a strong auth is required at the beginning.
39  * If true, trust agents will not be available until the user verifies their credentials.
40  */
41 class ScreenLockRule(val requireStrongAuth: Boolean = false) : TestRule {
42     private val context: Context = getApplicationContext()
43     private val userId = context.userId
44     private val uiDevice = UiDevice.getInstance(getInstrumentation())
45     private val windowManager = checkNotNull(WindowManagerGlobal.getWindowManagerService())
46     private val lockPatternUtils = LockPatternUtils(context)
47     private var instantLockSavedValue = false
48     private var strongAuthSavedValue: Int = 0
49 
50     override fun apply(base: Statement, description: Description) = object : Statement() {
51         override fun evaluate() {
52             verifyNoScreenLockAlreadySet()
53             dismissKeyguard()
54             setScreenLock()
55             setLockOnPowerButton()
56             configureStrongAuthState()
57 
58             try {
59                 base.evaluate()
60             } finally {
61                 restoreStrongAuthState()
62                 removeScreenLock()
63                 revertLockOnPowerButton()
64                 dismissKeyguard()
65             }
66         }
67     }
68 
69     private fun configureStrongAuthState() {
70         strongAuthSavedValue = lockPatternUtils.getStrongAuthForUser(userId)
71         if (requireStrongAuth) {
72             Log.d(TAG, "Triggering strong auth due to simulated lockdown")
73             lockPatternUtils.requireStrongAuth(STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN, userId)
74             wait("strong auth required after lockdown") {
75                 lockPatternUtils.getStrongAuthForUser(userId) ==
76                         STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN
77             }
78         }
79     }
80 
81     private fun restoreStrongAuthState() {
82         lockPatternUtils.requireStrongAuth(strongAuthSavedValue, userId)
83     }
84 
85     private fun verifyNoScreenLockAlreadySet() {
86         assertWithMessage("Screen Lock must not already be set on device")
87                 .that(lockPatternUtils.isSecure(context.userId))
88                 .isFalse()
89     }
90 
91     fun dismissKeyguard() {
92         wait("keyguard dismissed") { count ->
93             if (!uiDevice.isScreenOn) {
94                 Log.i(TAG, "Waking device, +500ms")
95                 uiDevice.wakeUp()
96             }
97 
98             // Bouncer may be shown due to a race; back dismisses it
99             if (count >= 10) {
100                 Log.i(TAG, "Pressing back to dismiss Bouncer")
101                 uiDevice.pressKeyCode(KeyEvent.KEYCODE_BACK)
102             }
103 
104             windowManager.dismissKeyguard(null, null)
105 
106             !windowManager.isKeyguardLocked
107         }
108     }
109 
110     fun successfulScreenLockAttempt() {
111         lockPatternUtils.verifyCredential(LockscreenCredential.createPin(PIN), context.userId, 0)
112         lockPatternUtils.userPresent(context.userId)
113         wait("strong auth not required") {
114             lockPatternUtils.getStrongAuthForUser(context.userId) == STRONG_AUTH_NOT_REQUIRED
115         }
116     }
117 
118     fun failedScreenLockAttempt() {
119         lockPatternUtils.verifyCredential(
120             LockscreenCredential.createPin(WRONG_PIN),
121             context.userId,
122             0
123         )
124     }
125 
126     private fun setScreenLock() {
127         lockPatternUtils.setLockCredential(
128                 LockscreenCredential.createPin(PIN),
129                 LockscreenCredential.createNone(),
130                 context.userId
131         )
132         wait("screen lock set") { lockPatternUtils.isSecure(context.userId) }
133         Log.i(TAG, "Device PIN set to $PIN")
134     }
135 
136     private fun setLockOnPowerButton() {
137         instantLockSavedValue = lockPatternUtils.getPowerButtonInstantlyLocks(context.userId)
138         lockPatternUtils.setPowerButtonInstantlyLocks(true, context.userId)
139     }
140 
141     private fun removeScreenLock() {
142         var lockCredentialUnset = lockPatternUtils.setLockCredential(
143                 LockscreenCredential.createNone(),
144                 LockscreenCredential.createPin(PIN),
145                 context.userId)
146         Log.i(TAG, "Removing screen lock")
147         assertWithMessage("Lock screen credential should be unset")
148                 .that(lockCredentialUnset)
149                 .isTrue()
150 
151         lockPatternUtils.setLockScreenDisabled(true, context.userId)
152         wait("screen lock un-set") {
153             lockPatternUtils.isLockScreenDisabled(context.userId)
154         }
155         wait("screen lock insecure") { !lockPatternUtils.isSecure(context.userId) }
156     }
157 
158     private fun revertLockOnPowerButton() {
159         lockPatternUtils.setPowerButtonInstantlyLocks(instantLockSavedValue, context.userId)
160     }
161 
162     companion object {
163         private const val TAG = "ScreenLockRule"
164         private const val PIN = "0000"
165         private const val WRONG_PIN = "0001"
166     }
167 }
168