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
