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 distributed under the
11  * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
12  * KIND, either express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15 package android.platform.uiautomator_helpers
16 
17 import android.content.pm.PackageManager
18 import androidx.test.platform.app.InstrumentationRegistry
19 
20 /**
21  * Adopt shell permissions for the target context.
22  *
23  * @param[permissions] the permission to adopt. Adopt all available permission is it's empty.
24  */
25 class ShellPrivilege(vararg permissions: String) : AutoCloseable {
26 
27     private val instrumentation = InstrumentationRegistry.getInstrumentation()
28     private val targetContext = instrumentation.targetContext
29     private val uiAutomation = instrumentation.uiAutomation
30     private var permissionsGranted = false
31 
32     init {
33         permissionsGranted = grantMissingPermissions(*permissions)
34     }
35 
36     /**
37      * @return[Boolean] True is there are any missing permission and we've successfully granted all
38      * of them.
39      */
grantMissingPermissionsnull40     private fun grantMissingPermissions(vararg permissions: String): Boolean {
41         if (permissions.isEmpty()) {
42             uiAutomation.adoptShellPermissionIdentity()
43             return true
44         }
45         val missingPermissions = permissions.filter { !it.isGranted() }.toTypedArray()
46         if (missingPermissions.isEmpty()) return false
47         uiAutomation.adoptShellPermissionIdentity(*missingPermissions)
48         return true
49     }
50 
closenull51     override fun close() {
52         if (permissionsGranted) instrumentation.uiAutomation.dropShellPermissionIdentity()
53         permissionsGranted = false
54     }
55 
isGrantednull56     private fun String.isGranted(): Boolean =
57         targetContext.checkCallingPermission(this) == PackageManager.PERMISSION_GRANTED
58 }
59