1#!/usr/bin/env python3
2#
3#   Copyright 2016 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import encodings
18import logging
19import shlex
20import shutil
21
22from mobly.controllers.android_device_lib.adb import AdbError
23from mobly.controllers.android_device_lib.adb import AdbProxy
24
25ROOT_USER_ID = '0'
26SHELL_USER_ID = '2000'
27UTF_8 = encodings.utf_8.getregentry().name
28
29
30class BlueberryAdbProxy(AdbProxy):
31    """Proxy class for ADB.
32
33    For syntactic reasons, the '-' in adb commands need to be replaced with
34    '_'. Can directly execute adb commands on an object:
35    >> adb = BlueberryAdbProxy(<serial>)
36    >> adb.start_server()
37    >> adb.devices() # will return the console output of "adb devices".
38    """
39
40    def __init__(self, serial="", ssh_connection=None):
41        """Construct an instance of AdbProxy.
42
43        Args:
44            serial: str serial number of Android device from `adb devices`
45            ssh_connection: SshConnection instance if the Android device is
46                            connected to a remote host that we can reach via SSH.
47        """
48        super().__init__(serial)
49        self._server_local_port = None
50        adb_path = shutil.which('adb')
51        adb_cmd = [shlex.quote(adb_path)]
52        if serial:
53            adb_cmd.append("-s %s" % serial)
54        if ssh_connection is not None:
55            # Kill all existing adb processes on the remote host (if any)
56            # Note that if there are none, then pkill exits with non-zero status
57            ssh_connection.run("pkill adb", ignore_status=True)
58            # Copy over the adb binary to a temp dir
59            temp_dir = ssh_connection.run("mktemp -d").stdout.strip()
60            ssh_connection.send_file(adb_path, temp_dir)
61            # Start up a new adb server running as root from the copied binary.
62            remote_adb_cmd = "%s/adb %s root" % (temp_dir, "-s %s" % serial if serial else "")
63            ssh_connection.run(remote_adb_cmd)
64            # Proxy a local port to the adb server port
65            local_port = ssh_connection.create_ssh_tunnel(5037)
66            self._server_local_port = local_port
67
68        if self._server_local_port:
69            adb_cmd.append("-P %d" % local_port)
70        self.adb_str = " ".join(adb_cmd)
71        self._ssh_connection = ssh_connection
72
73    def get_user_id(self):
74        """Returns the adb user. Either 2000 (shell) or 0 (root)."""
75        return self.shell('id -u').decode(UTF_8).rstrip()
76
77    def is_root(self, user_id=None):
78        """Checks if the user is root.
79
80        Args:
81            user_id: if supplied, the id to check against.
82        Returns:
83            True if the user is root. False otherwise.
84        """
85        if not user_id:
86            user_id = self.get_user_id()
87        return user_id == ROOT_USER_ID
88
89    def ensure_root(self):
90        """Ensures the user is root after making this call.
91
92        Note that this will still fail if the device is a user build, as root
93        is not accessible from a user build.
94
95        Returns:
96            False if the device is a user build. True otherwise.
97        """
98        self.ensure_user(ROOT_USER_ID)
99        return self.is_root()
100
101    def ensure_user(self, user_id=SHELL_USER_ID):
102        """Ensures the user is set to the given user.
103
104        Args:
105            user_id: The id of the user.
106        """
107        if self.is_root(user_id):
108            self.root()
109        else:
110            self.unroot()
111        self.wait_for_device()
112        return self.get_user_id() == user_id
113
114    def tcp_forward(self, host_port, device_port):
115        """Starts tcp forwarding from localhost to this android device.
116
117        Args:
118            host_port: Port number to use on localhost
119            device_port: Port number to use on the android device.
120
121        Returns:
122            Forwarded port on host as int or command output string on error
123        """
124        if self._ssh_connection:
125            # We have to hop through a remote host first.
126            #  1) Find some free port on the remote host's localhost
127            #  2) Setup forwarding between that remote port and the requested
128            #     device port
129            remote_port = self._ssh_connection.find_free_port()
130            host_port = self._ssh_connection.create_ssh_tunnel(remote_port, local_port=host_port)
131        try:
132            output = self.forward(["tcp:%d" % host_port, "tcp:%d" % device_port])
133        except AdbError as error:
134            return error
135        # If hinted_port is 0, the output will be the selected port.
136        # Otherwise, there will be no output upon successfully
137        # forwarding the hinted port.
138        if not output:
139            return host_port
140        try:
141            output_int = int(output)
142        except ValueError:
143            return output
144        return output_int
145
146    def remove_tcp_forward(self, host_port):
147        """Stop tcp forwarding a port from localhost to this android device.
148
149        Args:
150            host_port: Port number to use on localhost
151        """
152        if self._ssh_connection:
153            remote_port = self._ssh_connection.close_ssh_tunnel(host_port)
154            if remote_port is None:
155                logging.warning("Cannot close unknown forwarded tcp port: %d", host_port)
156                return
157            # The actual port we need to disable via adb is on the remote host.
158            host_port = remote_port
159        self.forward(["--remove", "tcp:%d" % host_port])
160
161    def path_exists(self, path):
162        """Check if a file path exists on an Android device
163
164        :param path: file path, could be a directory
165        :return: True if file path exists
166        """
167        try:
168            ret = self.shell("ls {}".format(path))
169            if ret is not None and len(ret) > 0:
170                return True
171            else:
172                return False
173        except AdbError as e:
174            logging.debug("path {} does not exist, error={}".format(path, e))
175            return False
176