async def send_and_receive(self, message, generate_identifier=True, timeout=5): """Send a message and wait for a response.""" if self._state not in [ ProtocolState.CONNECTED, ProtocolState.READY, ]: raise exceptions.InvalidStateError(self._state.name) # Some messages will respond with the same identifier as used in the # corresponding request. Others will not and one example is the crypto # message (for pairing). They will never include an identifier, but it # it is in turn only possible to have one of those message outstanding # at one time (i.e. it's not possible to mix up the responses). In # those cases, a "fake" identifier is used that includes the message # type instead. if generate_identifier: identifier = str(uuid.uuid4()).upper() message.identifier = identifier else: identifier = "type_" + str(message.type) self.connection.send(message) return await self._receive(identifier, timeout)
def takeover(self, protocol: Protocol) -> None: """Temporary override priority list with a specific protocol.""" if self._takeover_protocol: raise exceptions.InvalidStateError( f"{self._takeover_protocol[0]} has already done takeover" ) self._takeover_protocol = [protocol]
async def connect(self) -> None: """Connect to device.""" if self.remote_control.data_channel is None: raise exceptions.InvalidStateError( "remote control channel not connected") self.data_channel = self.remote_control.data_channel self.data_channel.listener = self
def port(self) -> int: """Remote connection port number.""" if self.transport is None: raise exceptions.InvalidStateError("not connected") sock = self.transport.get_extra_info("socket") _, dstport = sock.getpeername() return dstport
async def send(self, message): """Send a message and expect no response.""" if self._state not in [ ProtocolState.CONNECTED, ProtocolState.READY, ]: raise exceptions.InvalidStateError(self._state.name) self.connection.send(message)
def add_protocol(self, setup_data: SetupData): """Add a new protocol to the relay.""" # Connecting commits current configuration, thus adding new protocols is not # allowed anymore after that if self._protocol_handlers: raise exceptions.InvalidStateError( "cannot add protocol after connect was called") _LOGGER.debug("Adding handler for protocol %s", setup_data.protocol) self._protocols_to_setup.put(setup_data)
async def start(self, skip_initial_messages: bool = False) -> None: """Connect to device and listen to incoming messages.""" if self._state != ProtocolState.NOT_CONNECTED: raise exceptions.InvalidStateError(self._state.name) self._state = ProtocolState.CONNECTING try: await self.connection.connect() self._state = ProtocolState.CONNECTED # In case credentials have been given externally (i.e. not by pairing # with a device), then use that client id if self.service.credentials: self.srp.pairing_id = parse_credentials( self.service.credentials).client_id # The first message must always be DEVICE_INFORMATION, otherwise the # device will not respond with anything self.device_info = await self.send_and_receive( messages.device_information("pyatv", self.srp.pairing_id.decode())) # Distribute the device information to all listeners (as the # send_and_receive will stop that propagation). self.dispatch(protobuf.DEVICE_INFO_MESSAGE, self.device_info) # This is a hack to support re-use of a protocol object in # proxy (will be removed/refactored later) if skip_initial_messages: return await self._enable_encryption() # This should be the first message sent after encryption has # been enabled await self.send(messages.set_connection_state()) # Subscribe to updates at this stage await self.send_and_receive(messages.client_updates_config()) await self.send_and_receive(messages.get_keyboard_session()) except Exception: # Something went wrong, let's do cleanup self.stop() raise else: # We're now ready self._state = ProtocolState.READY
async def connect(self) -> None: """Initiate connection to device.""" # No protocols to setup + no protocols previously set up => no service if self._protocols_to_setup.empty() and not self._protocol_handlers: raise exceptions.NoServiceError("no service to connect to") # Protocols set up already => we have already connected if self._protocol_handlers: raise exceptions.InvalidStateError("already connected") devinfo: Dict[str, Any] = {} # Set up protocols, ignoring duplicates while not self._protocols_to_setup.empty(): setup_data = self._protocols_to_setup.get() if setup_data.protocol in self._protocol_handlers: _LOGGER.debug("Protocol %s already set up, ignoring", setup_data.protocol) continue _LOGGER.debug("Connecting to protocol: %s", setup_data.protocol) if await setup_data.connect(): _LOGGER.debug("Connected to protocol: %s", setup_data.protocol) self._protocol_handlers[setup_data.protocol] = setup_data for iface, instance in setup_data.interfaces.items(): self._interfaces[iface].register(instance, setup_data.protocol) self._features.add_mapping(setup_data.protocol, setup_data.features) dict_merge(devinfo, setup_data.device_info()) self._device_info = interface.DeviceInfo(devinfo) # Forward power events in case an interface exists for it try: power = cast(interface.Power, self._interfaces[interface.Power].main_instance) power.listener = self._interfaces[interface.Power] except exceptions.NotSupportedError: _LOGGER.debug("Power management not supported by any protocols")
def send(self, frame_type: FrameType, data: bytes) -> None: """Send message without waiting for a response.""" if self.transport is None: raise exceptions.InvalidStateError("not connected") payload_length = len(data) + (AUTH_TAG_LENGTH if self._chacha else 0) header = bytes([frame_type.value]) + payload_length.to_bytes( 3, byteorder="big") log_binary( _LOGGER, ">> Send data", FrameType=bytes([frame_type.value]), Data=data, ) if self._chacha: data = self._chacha.encrypt(data, aad=header) log_binary(_LOGGER, ">> Send", Header=header, Encrypted=data) self.transport.write(header + data)
def _send_first_in_queue(self) -> None: if self.transport is None: raise exceptions.InvalidStateError("not connected") if not self._queue: return _, frame_type, data, _ = self._queue[0] log_binary(_LOGGER, ">> Send data", FrameType=bytes([frame_type.value]), Data=data) payload_length = len(data) + (AUTH_TAG_LENGTH if self._chacha else 0) header = bytes([frame_type.value]) + payload_length.to_bytes( 3, byteorder="big") if self._chacha: data = self._chacha.encrypt(data, aad=header) log_binary(_LOGGER, ">> Send", Header=header, Encrypted=data) self.transport.write(header + data)
def acquire(self) -> None: """Acquire playback manager for playback.""" if self._is_acquired: raise exceptions.InvalidStateError("already streaming to device") self._is_acquired = True