def close_stream_transport(self, stream_transport, timeout): """Remove the given stream transport's id from our map of id's. If the stream id is actually removed, we send a CLSE message to let the remote end know (this happens when we are ack'ing a CLSE message we received). The ADB protocol doesn't say this is a requirement, but ADB does it, so we do too. Args: stream_transport: The stream transport to close. timeout: Timeout on the operation. Returns: True if the id was removed and message sent, False if it was already missing from the stream map (already closed). """ with self._stream_transport_map_lock: if stream_transport.local_id in self._stream_transport_map: del self._stream_transport_map[stream_transport.local_id] # If we never got a remote_id, there's no CLSE message to send. if stream_transport.remote_id: self.transport.write_message( adb_message.AdbMessage('CLSE', stream_transport.local_id, stream_transport.remote_id), timeout) return True return False
def open_stream(self, destination, timeout_ms=None): """Opens a new stream to a destination service on the device. Not the same as the posix 'open' or any other Open methods, this corresponds to the OPEN message described in the ADB protocol documentation mentioned above. It creates a stream (uniquely identified by remote/local ids) that connects to a particular service endpoint. Args: destination: The service:command string, see ADB documentation. timeout_ms: Timeout in milliseconds for the Open to succeed (or as a PolledTimeout object). Raises: AdbProtocolError: Wrong local_id sent to us, or we didn't get a ready response. Returns: An AdbStream object that can be used to read/write data to the specified service endpoint, or None if the requested service couldn't be opened. """ timeout = timeouts.PolledTimeout.from_millis(timeout_ms) stream_transport = self._make_stream_transport() self.transport.write_message( adb_message.AdbMessage(command='OPEN', arg0=stream_transport.local_id, arg1=0, data=destination + '\0'), timeout) if not stream_transport.ensure_opened(timeout): return None return AdbStream(destination, stream_transport)
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 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)