def _handle_message(self, message, handle_wrte=True):
"""Handle a message that was read for this stream.
For each message type, this means:
OKAY: Check id's and make sure we are expecting an OKAY. Clear the
self._expecting_okay flag so any pending write()'s know.
CLSE: Set our internal state to closed.
WRTE: Add the data read to our internal read buffer. Note we don't
return the actual data because it may not be this thread that needs it.
Args:
message: Message that was read.
handle_wrte: If True, we can handle WRTE messages, otherwise raise.
Raises:
AdbProtocolError: If we get a WRTE message but handle_wrte is False.
"""
if message.command == 'OKAY':
self._set_or_check_remote_id(message.arg0)
if not self._expecting_okay:
raise usb_exceptions.AdbProtocolError(
'%s received unexpected OKAY: %s', self, message)
self._expecting_okay = False
elif message.command == 'CLSE':
self.closed_state = self.ClosedState.CLOSED
elif not handle_wrte:
raise usb_exceptions.AdbProtocolError(
'%s received WRTE before OKAY/CLSE: %s', self, message)
else:
with self._read_buffer_lock:
self._read_buffer.append(message.data)
self._buffer_size += len(message.data)
def _send_command(self, command, timeout, data=''):
"""Send the given command/data over this transport.
We do a couple sanity checks in here to be sure we can format a valid
AdbMessage for our underlying AdbConnection, and then send it. This method
can be used to send any message type, and doesn't do any state tracking or
acknowledgement checking.
Args:
command: The command to send, should be one of 'OKAY', 'WRTE', or 'CLSE'
timeout: timeouts.PolledTimeout to use for this operation
data: If provided, data to send with the AdbMessage.
"""
if len(data) > self.adb_connection.maxdata:
raise usb_exceptions.AdbProtocolError(
'Message data too long (%s>%s): %s', len(data),
self.adb_connection.maxdata, data)
if not self.remote_id:
# If we get here, we probably missed the OKAY response to our OPEN. We
# should have failed earlier, but in case someone does something tricky
# with multiple threads, we sanity check this here.
raise usb_exceptions.AdbProtocolError('%s send before OKAY: %s',
self, data)
self.adb_connection.transport.write_message(
adb_message.AdbMessage(command, self.local_id, self.remote_id,
data), timeout)
def _handle_message_for_stream(self, stream_transport, message, timeout):
"""Handle an incoming message, check if it's for the given stream.
If the message is not for the stream, then add it to the appropriate
message queue.
Args:
stream_transport: AdbStreamTransport currently waiting on a message.
message: Message to check and handle.
timeout: Timeout to use for the operation, should be an instance of
timeouts.PolledTimeout.
Returns:
The message read if it was for this stream, None otherwise.
Raises:
AdbProtocolError: If we receive an unexepcted message type.
"""
if message.command not in ('OKAY', 'CLSE', 'WRTE'):
raise usb_exceptions.AdbProtocolError(
'%s received unexpected message: %s', self, message)
if message.arg1 == stream_transport.local_id:
# Ack writes immediately.
if message.command == 'WRTE':
# Make sure we don't get a WRTE before an OKAY/CLSE message.
if not stream_transport.remote_id:
raise usb_exceptions.AdbProtocolError(
'%s received WRTE before OKAY/CLSE: %s',
stream_transport, message)
self.transport.write_message(
adb_message.AdbMessage('OKAY', stream_transport.local_id,
stream_transport.remote_id),
timeout)
elif message.command == 'CLSE':
self.close_stream_transport(stream_transport, timeout)
return message
else:
# Message was not for this stream, add it to the right stream's queue.
with self._stream_transport_map_lock:
dest_transport = self._stream_transport_map.get(message.arg1)
if dest_transport:
if message.command == 'CLSE':
self.close_stream_transport(dest_transport, timeout)
dest_transport.enqueue_message(message, timeout)
else:
_LOG.warning('Received message for unknown local-id: %s',
message)
def read_message(self, timeout=None):
"""Read a message from this transport and return it.
Reads a message of RECV_MSG_TYPE and returns it. Note that this method
abstracts the data length and data read so that the caller simply gets the
data along with the header in the returned message.
Args:
timeout: timeouts.PolledTimeout to use for the operation.
Returns:
An instance of self.RECV_MSG_TYPE that was read from self.stream.
Raises:
AdbProtocolError: If an invalid response is received.
AdbRemoteError: If a FAIL response is received.
"""
raw_data = self.stream.read(
struct.calcsize(self.RECV_MSG_TYPE.struct_format), timeout)
try:
raw_message = struct.unpack(self.RECV_MSG_TYPE.struct_format,
raw_data)
except struct.error:
raise usb_exceptions.AdbProtocolError(
'%s expected format "%s", got data %s' %
(self, self.RECV_MSG_TYPE.struct_format, raw_data))
if raw_message[0] not in self.WIRE_TO_CMD:
raise usb_exceptions.AdbProtocolError(
'Unrecognized command id: %s' % raw_message)
# Swap out the wire command with the string equivalent.
raw_message = (self.WIRE_TO_CMD[raw_message[0]], ) + raw_message[1:]
if self.RECV_MSG_TYPE.has_data and raw_message[-1]:
# For messages that have data, the length of the data is the last field
# in the struct. We do another read and swap out that length for the
# actual data read before we create the namedtuple to return.
data_len = raw_message[-1]
raw_message = raw_message[:-1] + (self.stream.read(
data_len, timeout), )
if raw_message[0] not in self.VALID_RESPONSES:
raise usb_exceptions.AdbProtocolError(
'%s not a valid response for %s' % (raw_message[0], self))
if raw_message[0] == 'FAIL':
raise usb_exceptions.AdbRemoteError('Remote ADB failure: %s' %
raw_message)
return self.RECV_MSG_TYPE(*raw_message)
def _set_or_check_remote_id(self, remote_id):
"""Set or check the remote id."""
if not self.remote_id:
assert self.closed_state == self.ClosedState.PENDING, 'Bad ClosedState!'
self.remote_id = remote_id
self.closed_state = self.ClosedState.OPEN
elif self.remote_id != remote_id:
raise usb_exceptions.AdbProtocolError('%s remote-id change to %s',
self, remote_id)
def __init__(self, command, arg0=0, arg1=0, data=''):
if command not in self.CMD_TO_WIRE:
raise usb_exceptions.AdbProtocolError(
'Unrecognized ADB command: %s' % command)
self._command = self.CMD_TO_WIRE[command]
self.arg0 = arg0
self.arg1 = arg1
self.data = data
self.magic = self._command ^ 0xFFFFFFFF
def read_message(self, timeout):
"""Read an AdbMessage from this transport.
Args:
timeout: Timeout for the entire read operation, in the form of a
timeouts.PolledTimeout instance. Note that for packets with a data
payload, two USB reads are performed.
Returns:
The ADB message read from the device.
Raises:
UsbReadFailedError: There's an error during read, including timeout.
AdbProtocolError: A message is incorrectly formatted.
AdbTimeoutError: timeout is already expired, or expires before we read the
entire message, specifically between reading header and data packets.
"""
with self._reader_lock:
raw_header = self._transport.read(
struct.calcsize(AdbMessage.HEADER_STRUCT_FORMAT),
timeout.remaining_ms)
if not raw_header:
raise usb_exceptions.AdbProtocolError('Adb connection lost')
try:
raw_message = RawAdbMessage(*struct.unpack(
AdbMessage.HEADER_STRUCT_FORMAT, raw_header))
except struct.error as exception:
raise usb_exceptions.AdbProtocolError(
'Unable to unpack ADB command (%s): %s (%s)' %
(AdbMessage.HEADER_STRUCT_FORMAT, raw_header, exception))
if raw_message.data_length > 0:
if timeout.has_expired():
_LOG.warning(
'Timed out between AdbMessage header and data, reading '
'data anyway with 10ms timeout')
timeout = timeouts.PolledTimeout.from_millis(10)
data = self._transport.read(raw_message.data_length,
timeout.remaining_ms)
else:
data = ''
return raw_message.to_adb_message(data)
def Write(self, data, timeout):
"""Write data to this stream, using the given timeouts.PolledTimeout."""
if not self.remote_id:
raise usb_exceptions.AdbStreamClosedError(
'Cannot Write() to half-opened %s', self)
if self.closed_state != self.ClosedState.OPEN:
raise usb_exceptions.AdbStreamClosedError(
'Cannot Write() to closed %s', self)
elif self._expecting_okay:
raise usb_exceptions.AdbProtocolError(
'Previous WRTE failed, %s in unknown state', self)
# Make sure we only have one WRTE in flight at a time, because ADB doesn't
# identify which WRTE it is ACK'ing when it sends the OKAY message back.
with self._write_lock:
self._expecting_okay = True
self._SendCommand('WRTE', timeout, data)
self._ReadMessagesUntilTrue(lambda: not self._expecting_okay, timeout)
def __init__(self, transport, maxdata, remote_banner):
"""Create an ADB connection to a device.
Args:
transport: AdbTransportAdapter to use for reading/writing AdbMessages
maxdata: Max data size the remote endpoint will accept.
remote_banner: Banner received from the remote endpoint.
"""
try:
self.systemtype, self.serial, self.banner = remote_banner.split(':', 2)
except ValueError:
raise usb_exceptions.AdbProtocolError('Received malformed banner %s',
remote_banner)
self.transport = transport
self.maxdata = maxdata
self._last_id_used = 0
self._reader_lock = threading.Lock()
self._open_lock = threading.Lock()
# Maps local_id: AdbStreamTransport object for the relevant stream.
self._stream_transport_map = {}
self._stream_transport_map_lock = threading.RLock()
def connect(cls,
transport,
rsa_keys=None,
timeout_ms=1000,
auth_timeout_ms=100):
"""Establish a new connection to a device, connected via transport.
Args:
transport: A transport to use for reads/writes from/to the device,
usually an instance of UsbHandle, but really it can be anything with
read() and write() methods.
rsa_keys: List of AuthSigner subclass instances to be used for
authentication. The device can either accept one of these via the sign
method, or we will send the result of get_public_key from the first one
if the device doesn't accept any of them.
timeout_ms: Timeout to wait for the device to respond to our CNXN
request. Actual timeout may take longer if the transport object passed
has a longer default timeout than timeout_ms, or if auth_timeout_ms is
longer than timeout_ms and public key auth is used. This argument may
be a PolledTimeout object.
auth_timeout_ms: Timeout to wait for when sending a new public key. This
is only relevant when we send a new public key. The device shows a
dialog and this timeout is how long to wait for that dialog. If used
in automation, this should be low to catch such a case as a failure
quickly; while in interactive settings it should be high to allow
users to accept the dialog. We default to automation here, so it's low
by default. This argument may be a PolledTimeout object.
Returns:
An instance of AdbConnection that is connected to the device.
Raises:
usb_exceptions.DeviceAuthError: When the device expects authentication,
but we weren't given any valid keys.
usb_exceptions.AdbProtocolError: When the device does authentication in an
unexpected way, or fails to respond appropriately to our CNXN request.
"""
timeout = timeouts.PolledTimeout.from_millis(timeout_ms)
if ADB_MESSAGE_LOG:
adb_transport = adb_message.DebugAdbTransportAdapter(transport)
else:
adb_transport = adb_message.AdbTransportAdapter(transport)
adb_transport.write_message(
adb_message.AdbMessage(command='CNXN',
arg0=ADB_VERSION,
arg1=MAX_ADB_DATA,
data='host::%s\0' % ADB_BANNER), timeout)
msg = adb_transport.read_until(('AUTH', 'CNXN'), timeout)
if msg.command == 'CNXN':
return cls(adb_transport, msg.arg1, msg.data)
# We got an AUTH response, so we have to try to authenticate.
if not rsa_keys:
raise usb_exceptions.DeviceAuthError(
'Device authentication required, no keys available.')
# Loop through our keys, signing the last 'banner' or token.
for rsa_key in rsa_keys:
if msg.arg0 != cls.AUTH_TOKEN:
raise usb_exceptions.AdbProtocolError('Bad AUTH response: %s',
msg)
signed_token = rsa_key.sign(msg.data)
adb_transport.write_message(
adb_message.AdbMessage(command='AUTH',
arg0=cls.AUTH_SIGNATURE,
arg1=0,
data=signed_token), timeout)
msg = adb_transport.read_until(('AUTH', 'CNXN'), timeout)
if msg.command == 'CNXN':
return cls(adb_transport, msg.arg1, msg.data)
# None of the keys worked, so send a public key.
adb_transport.write_message(
adb_message.AdbMessage(command='AUTH',
arg0=cls.AUTH_RSAPUBLICKEY,
arg1=0,
data=rsa_keys[0].get_public_key() + '\0'),
timeout)
try:
msg = adb_transport.read_until(
('CNXN', ),
timeouts.PolledTimeout.from_millis(auth_timeout_ms))
except usb_exceptions.UsbReadFailedError as exception:
if exception.is_timeout():
exceptions.reraise(usb_exceptions.DeviceAuthError,
'Accept auth key on device, then retry.')
raise
# The read didn't time-out, so we got a CNXN response.
return cls(adb_transport, msg.arg1, msg.data)
def assert_command_is(self, command):
"""Assert that a message's command matches the given command."""
if self.command != command:
raise usb_exceptions.AdbProtocolError(
'Expected %s command, received %s' % (command, self))
def assert_command_is(self, command): # pylint: disable=invalid-name
"""Assert that a message's command matches the given command."""
if self.command != command:
raise usb_exceptions.AdbProtocolError(
'Expected %s command, received %s', command, self)