From 7125f684641abb7aa7677a2e15769537d1534f16 Mon Sep 17 00:00:00 2001 From: annon Date: Tue, 14 Jan 2025 01:14:16 -0500 Subject: [PATCH] Adds shell v2 support, which can return stdout/stderr as separate information, as well as the return code from the shell execution --- README.rst | 3 ++ adb_shell/adb_device.py | 81 +++++++++++++++++++++++++++++--- adb_shell/adb_device_async.py | 80 +++++++++++++++++++++++++++++--- adb_shell/constants.py | 18 ++++++++ tests/patchers.py | 10 ++-- tests/test_adb_device.py | 84 ++++++++++++++++++++++++++++++++++ tests/test_adb_device_async.py | 77 +++++++++++++++++++++++++++++++ 7 files changed, 338 insertions(+), 15 deletions(-) diff --git a/README.rst b/README.rst index 1f53d40..7b88c49 100644 --- a/README.rst +++ b/README.rst @@ -75,6 +75,9 @@ Example Usage response1 = device1.shell('echo TEST1') response2 = device2.shell('echo TEST2') + # If the device supports the newer "shell protocol", you can also obtain individual streams + response3 = device1.shell_v2('echo TEST1') + Generate ADB Key Files ********************** diff --git a/adb_shell/adb_device.py b/adb_shell/adb_device.py index 9c3b7b7..a19a863 100644 --- a/adb_shell/adb_device.py +++ b/adb_shell/adb_device.py @@ -229,7 +229,8 @@ def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_info): # 4. If ``cmd`` is not ``b'AUTH'``, then authentication is not necesary and so we are done if cmd != constants.AUTH: - return True, maxdata + # TODO confirm that a rooted device will return device features here + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False # 5. If no ``rsa_keys`` are provided, raise an exception if not rsa_keys: @@ -248,12 +249,12 @@ def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_info): msg = AdbMessage(constants.AUTH, constants.AUTH_SIGNATURE, 0, signed_token) self._send(msg, adb_info) - # 6.3. Read the response from the device + # 6.3. Read the response from the device, which will include the device's feature list cmd, arg0, maxdata, banner2 = self._read_expected_packet_from_device([constants.CNXN, constants.AUTH], adb_info) # 6.4. If ``cmd`` is ``b'CNXN'``, we are done if cmd == constants.CNXN: - return True, maxdata + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False # 7. None of the keys worked, so send ``rsa_keys[0]``'s public key; if the response does not time out, we must have connected successfully pubkey = rsa_keys[0].GetPublicKey() @@ -267,8 +268,9 @@ def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_info): self._send(msg, adb_info) adb_info.transport_timeout_s = auth_timeout_s - _, _, maxdata, _ = self._read_expected_packet_from_device([constants.CNXN], adb_info) - return True, maxdata + # TODO confirm device features are sent during type 3 authentication + _, _, maxdata, banner2 = self._read_expected_packet_from_device([constants.CNXN], adb_info) + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False def read(self, expected_cmds, adb_info, allow_zeros=False): """Read packets from the device until we get an expected packet type. @@ -672,7 +674,7 @@ def connect(self, rsa_keys=None, transport_timeout_s=None, auth_timeout_s=consta self._available = False # Use the IO manager to connect - self._available, self._maxdata = self._io_manager.connect(self._banner, rsa_keys, auth_timeout_s, auth_callback, adb_info) + self._available, self._maxdata, self._shell_v2_supported = self._io_manager.connect(self._banner, rsa_keys, auth_timeout_s, auth_callback, adb_info) return self._available @@ -839,6 +841,73 @@ def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFA return self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, decode) + def shell_v2(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True): + """ + Send an ADB shell command to the device, using the shell protocol. + + The shell protocol is able to separate stdout/stderr streams, as well as provide a return code from the called process. + + Parameters + ---------- + command : str + The shell command that will be sent + transport_timeout_s : float, None + Timeout in seconds for sending and receiving data, or ``None``; see :meth:`BaseTransport.bulk_read() ` + and :meth:`BaseTransport.bulk_write() ` + read_timeout_s : float + The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`_AdbIOManager.read` + timeout_s : float, None + The total time in seconds to wait for the ADB command to finish + decode : bool + Whether to decode the output to utf8 before returning + + Returns + ------- + dict + A dict representing the result of the ADB shell command. Dictionary keys are 'stdout', 'stderr', and 'return_code'. If ``decode`` is True, + stdout/stderr will be returned as strings - otherwise as bytes. + + """ + if not self.available: + raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established. (Did you call `AdbDevice.connect()`?)") + + if not self._shell_v2_supported: + raise exceptions.AdbCommandFailureException("ADB shell_v2 service is not supported on this device.") + + # For shell_v2 calls, we force decode to be False, so that we can separate out the streams using the raw bytes + shell_v2_result = self._service(b'shell,v2', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, False) + + data = BytesIO(shell_v2_result) + stdout = b'' + stderr = b'' + return_code: int = None + while True: + header = data.read(constants.SHELL_V2_HEADER_LENGTH) + if len(header) < constants.SHELL_V2_HEADER_LENGTH: + if len(header) != 0: + _LOGGER.warning(f'ADB shell output contained unprocessed data: {header}') + break + output_fd, output_len = struct.unpack(constants.V2_RESPONSE_FORMAT, header) + output = data.read(output_len) + + if output_fd == constants.KID_STDOUT: + stdout += output + elif output_fd == constants.KID_STDERR: + stderr += output + elif output_fd == constants.KID_EXIT: + if return_code is not None: + _LOGGER.warning('Multiple return codes found for a single shell execution') + # Return code is a single byte + return_code = output[0] + else: + _LOGGER.warning(b'shell v2 output was not processed correctly: %s%s' % (header, output)) + + return { + 'stdout': stdout.decode('utf8', _DECODE_ERRORS) if decode else stdout, + 'stderr': stderr.decode('utf8', _DECODE_ERRORS) if decode else stderr, + 'return_code': return_code + } + def streaming_shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True): """Send an ADB shell command to the device, yielding each line of output. diff --git a/adb_shell/adb_device_async.py b/adb_shell/adb_device_async.py index 84ce523..205525f 100644 --- a/adb_shell/adb_device_async.py +++ b/adb_shell/adb_device_async.py @@ -266,7 +266,8 @@ async def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_inf # 4. If ``cmd`` is not ``b'AUTH'``, then authentication is not necesary and so we are done if cmd != constants.AUTH: - return True, maxdata + # TODO confirm that a rooted device will return device features here + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False # 5. If no ``rsa_keys`` are provided, raise an exception if not rsa_keys: @@ -285,12 +286,12 @@ async def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_inf msg = AdbMessage(constants.AUTH, constants.AUTH_SIGNATURE, 0, signed_token) await self._send(msg, adb_info) - # 6.3. Read the response from the device + # 6.3. Read the response from the device, which will include the device's feature list cmd, arg0, maxdata, banner2 = await self._read_expected_packet_from_device([constants.CNXN, constants.AUTH], adb_info) # 6.4. If ``cmd`` is ``b'CNXN'``, we are done if cmd == constants.CNXN: - return True, maxdata + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False # 7. None of the keys worked, so send ``rsa_keys[0]``'s public key; if the response does not time out, we must have connected successfully pubkey = rsa_keys[0].GetPublicKey() @@ -304,8 +305,9 @@ async def connect(self, banner, rsa_keys, auth_timeout_s, auth_callback, adb_inf await self._send(msg, adb_info) adb_info.transport_timeout_s = auth_timeout_s - _, _, maxdata, _ = await self._read_expected_packet_from_device([constants.CNXN], adb_info) - return True, maxdata + # TODO confirm device features are sent during type 3 authentication + _, _, maxdata, banner2 = await self._read_expected_packet_from_device([constants.CNXN], adb_info) + return True, maxdata, True if constants.SHELL_V2_FEATURE in banner2 else False async def read(self, expected_cmds, adb_info, allow_zeros=False): """Read packets from the device until we get an expected packet type. @@ -711,7 +713,7 @@ async def connect(self, rsa_keys=None, transport_timeout_s=None, auth_timeout_s= self._available = False # Use the IO manager to connect - self._available, self._maxdata = await self._io_manager.connect(self._banner, rsa_keys, auth_timeout_s, auth_callback, adb_info) + self._available, self._maxdata, self._shell_v2_supported = await self._io_manager.connect(self._banner, rsa_keys, auth_timeout_s, auth_callback, adb_info) return self._available @@ -880,6 +882,72 @@ async def shell(self, command, transport_timeout_s=None, read_timeout_s=constant return await self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, decode) + async def shell_v2(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True): + """ + Send an ADB shell command to the device, using the shell protocol. + + The shell protocol is able to separate stdout/stderr streams, as well as provide a return code from the called process. + + Parameters + ---------- + command : str + The shell command that will be sent + transport_timeout_s : float, None + Timeout in seconds for sending and receiving data, or ``None``; see :meth:`BaseTransportAsync.bulk_read() ` + and :meth:`BaseTransportAsync.bulk_write() ` + read_timeout_s : float + The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`_AdbIOManagerAsync.read` + timeout_s : float, None + The total time in seconds to wait for the ADB command to finish + decode : bool + Whether to decode the output to utf8 before returning + + + Returns + ------- + dict + A dict representing the result of the ADB shell command. Dictionary keys are 'stdout', 'stderr', and 'return_code'. If ``decode`` is True, + stdout/stderr will be returned as strings - otherwise as bytes. + + """ + if not self.available: + raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established. (Did you call `AdbDeviceAsync.connect()`?)") + + if not self._shell_v2_supported: + raise exceptions.AdbCommandFailureException("ADB shell_v2 service is not supported on this device.") # For shell_v2 calls, we force decode to be False, so that we can separate out the streams using the raw bytes + shell_v2_result = self._service(b'shell,v2', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, False) + + data = BytesIO(await shell_v2_result) + stdout = b'' + stderr = b'' + return_code: int = None + while True: + header = data.read(constants.SHELL_V2_HEADER_LENGTH) + if len(header) < constants.SHELL_V2_HEADER_LENGTH: + if len(header) != 0: + _LOGGER.warning(f'ADB shell output contained unprocessed data: {header}') + break + output_fd, output_len = struct.unpack(constants.V2_RESPONSE_FORMAT, header) + output = data.read(output_len) + + if output_fd == constants.KID_STDOUT: + stdout += output + elif output_fd == constants.KID_STDERR: + stderr += output + elif output_fd == constants.KID_EXIT: + if return_code is not None: + _LOGGER.warning('Multiple return codes found for a single shell execution') + # Return code is a single byte + return_code = output[0] + else: + _LOGGER.warning(b'shell v2 output was not processed correctly: %s%s' % (header, output)) + + return { + 'stdout': stdout.decode('utf8', 'backslashreplace') if decode else stdout, + 'stderr': stderr.decode('utf8', 'backslashreplace') if decode else stderr, + 'return_code': return_code + } + async def streaming_shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True): """Send an ADB shell command to the device, yielding each line of output. diff --git a/adb_shell/constants.py b/adb_shell/constants.py index 5b2c702..35998dc 100644 --- a/adb_shell/constants.py +++ b/adb_shell/constants.py @@ -36,6 +36,15 @@ #: From adb.h PROTOCOL = 0x01 +#: From shell_service.h +KID_STDOUT = 1 + +#: From shell_service.h +KID_STDERR = 2 + +#: From shell_service.h +KID_EXIT = 3 + #: ADB protocol version. VERSION = 0x01000000 @@ -61,6 +70,9 @@ #: AUTH constant for ``arg0`` AUTH_RSAPUBLICKEY = 3 +#: For shell_v2 calls, there is a five byte header for each data segment +SHELL_V2_HEADER_LENGTH = 5 + AUTH = b'AUTH' CLSE = b'CLSE' CNXN = b'CNXN' @@ -112,6 +124,9 @@ #: The format for FileSync "stat" messages FILESYNC_STAT_FORMAT = b'<4I' +# The format for a "shell_v2" response +V2_RESPONSE_FORMAT = b'