async def get_stub(self) -> AsyncContextManager[CompanionServiceStub]: await self.spawn_notifier() channel: Optional[Channel] = None try: try: self.companion_info = self.direct_companion_manager.get_companion_info( target_udid=self.target_udid ) except IdbException as e: # will try to spawn a companion if on mac. companion_info = await self.spawn_companion( target_udid=none_throws(self.target_udid) ) if companion_info: self.companion_info = companion_info else: raise e self.logger.info(f"using companion {self.companion_info}") channel = Channel( # pyre-fixme[16]: `Optional` has no attribute `host`. self.companion_info.host, # pyre-fixme[16]: `Optional` has no attribute `port`. self.companion_info.port, loop=asyncio.get_event_loop(), ) yield CompanionServiceStub(channel=channel) finally: if channel: channel.close()
async def connect( self, destination: ConnectionDestination, metadata: Optional[Dict[str, str]] = None, ) -> CompanionInfo: self.logger.debug(f"Connecting directly to {destination} with meta {metadata}") if isinstance(destination, Address): channel = Channel( destination.host, destination.port, loop=asyncio.get_event_loop() ) stub = CompanionServiceStub(channel=channel) with tempfile.NamedTemporaryFile(mode="w+b") as f: response = await stub.connect( ConnectRequest( destination=destination_to_grpc(destination), metadata=metadata, local_file_path=f.name, ) ) companion = CompanionInfo( udid=response.companion.udid, host=destination.host, port=destination.port, is_local=response.companion.is_local, ) self.logger.debug(f"Connected directly to {companion}") self.direct_companion_manager.add_companion(companion) channel.close() return companion else: companion = await self.spawn_companion(target_udid=destination) if companion: return companion else: raise IdbException(f"can't find target for udid {destination}")
async def get_stub(self) -> CompanionServiceStub: channel: Optional[Channel] = None try: try: self.companion_info = self.direct_companion_manager.get_companion_info( target_udid=self.target_udid ) except IdbException as e: # will try to spawn a companion if on mac. if platform.system() == "Darwin" and self.target_udid: self.logger.info( f"will attempt to spawn a companion for {self.target_udid}" ) port = await self.companion_spawner.spawn_companion( target_udid=self.target_udid ) host = "localhost" self.companion_info = CompanionInfo( host=host, port=port, udid=self.target_udid, is_local=True ) self.direct_companion_manager.add_companion(self.companion_info) else: raise e self.logger.info(f"using companion {self.companion_info}") channel = Channel( self.companion_info.host, self.companion_info.port, loop=asyncio.get_event_loop(), ) yield CompanionServiceStub(channel=channel) finally: if channel: channel.close()
async def build( cls, address: Address, logger: logging.Logger, is_local: Optional[bool] = None, exchange_metadata: bool = True, extra_metadata: Optional[Dict[str, str]] = None, use_tls: bool = False, ) -> AsyncGenerator["Client", None]: metadata_to_companion = ( { **{ key: value for (key, value) in plugin.resolve_metadata(logger=logger).items() if isinstance(value, str) }, **(extra_metadata or {}), } if exchange_metadata else {} ) ssl_context = plugin.channel_ssl_context() if use_tls else None if use_tls: assert ssl_context is not None async with ( Channel( host=address.host, port=address.port, loop=asyncio.get_event_loop(), ssl=ssl_context, ) if isinstance(address, TCPAddress) else Channel(path=address.path, loop=asyncio.get_event_loop()) ) as channel: stub = CompanionServiceStub(channel=channel) with tempfile.NamedTemporaryFile(mode="w+b") as f: try: response = await stub.connect( ConnectRequest( metadata=metadata_to_companion, local_file_path=f.name ) ) except Exception as ex: raise IdbException( f"Failed to connect to companion at address {address}: {ex}" ) companion = companion_to_py( companion=response.companion, address=address, is_local=is_local ) if exchange_metadata: metadata_from_companion = { key: value for (key, value) in companion.metadata.items() if isinstance(value, str) } plugin.append_companion_metadata( logger=logger, metadata=metadata_from_companion ) yield Client(stub=stub, companion=companion, logger=logger)
async def _companion_to_target( self, companion: CompanionInfo) -> TargetDescription: channel = Channel(companion.host, companion.port, loop=asyncio.get_event_loop()) stub = CompanionServiceStub(channel=channel) response = await stub.describe(TargetDescriptionRequest()) channel.close() return target_to_py(response.target_description)
async def provide_client(self) -> CompanionClient: await self.daemon_spawner.start_daemon_if_needed( force_kill=self.force_kill_daemon) if not self.channel or not self.stub: self.channel = Channel(self.host, self.port, loop=asyncio.get_event_loop()) self.stub = CompanionServiceStub(channel=self.channel) return CompanionClient(stub=self.stub, is_local=True, udid=self.target_udid, logger=self.logger)
async def build( cls, host: str, port: int, is_local: bool, logger: logging.Logger) -> AsyncContextManager[IdbClientBase]: channel = Channel(host=host, port=port, loop=asyncio.get_event_loop()) try: yield IdbClient( stub=CompanionServiceStub(channel=channel), is_local=is_local, logger=logger, ) finally: channel.close()
async def _companion_to_target( self, companion: CompanionInfo) -> Optional[TargetDescription]: try: channel = Channel(companion.host, companion.port, loop=asyncio.get_event_loop()) stub = CompanionServiceStub(channel=channel) response = await stub.describe(TargetDescriptionRequest()) channel.close() return target_to_py(response.target_description) except Exception: self.logger.warning(f"Failed to describe {companion}, removing it") self.direct_companion_manager.remove_companion( Address(companion.host, companion.port)) return None
async def build( cls, companion_info: CompanionInfo, logger: logging.Logger) -> AsyncContextManager["GrpcStubClient"]: channel = Channel( host=companion_info.host, port=companion_info.port, loop=asyncio.get_event_loop(), ) try: yield GrpcStubClient( stub=CompanionServiceStub(channel=channel), companion_info=companion_info, logger=logger, ) finally: channel.close()
async def build( cls, address: Address, is_local: bool, logger: logging.Logger) -> AsyncContextManager["IdbClient"]: channel = (Channel(host=address.host, port=address.port, loop=asyncio.get_event_loop()) if isinstance( address, TCPAddress) else Channel(path=address.path, loop=asyncio.get_event_loop())) try: yield IdbClient( stub=CompanionServiceStub(channel=channel), address=address, is_local=is_local, logger=logger, ) finally: channel.close()
def func_wrapper(client, *args, **kwargs): # pyre-ignore try: client.companion_info: CompanionInfo = client.direct_companion_manager.get_companion_info( target_udid=client.target_udid) client.logger.info(f"using companion {client.companion_info}") client.channel = Channel( client.companion_info.host, client.companion_info.port, loop=asyncio.get_event_loop(), ) client.stub: Optional[CompanionServiceStub] = CompanionServiceStub( channel=client.channel) return func(client, *args, **kwargs) except GRPCError as e: raise IdbException(e.message) from e # noqa B306 except (ProtocolError, StreamTerminatedError) as e: raise IdbException(e.args) from e
def __init__( self, port: int, host: str, target_udid: Optional[str], logger: Optional[logging.Logger] = None, force_kill_daemon: bool = False, ) -> None: self.port: int = port self.host: str = host self.logger: logging.Logger = (logger if logger else logging.getLogger("idb_grpc_client")) self.force_kill_daemon = force_kill_daemon self.target_udid = target_udid self.daemon_spawner = DaemonSpawner(host=self.host, port=self.port) self.daemon_channel: Optional[Channel] = None self.daemon_stub: Optional[CompanionServiceStub] = None for (call_name, f) in ipc_loader.client_calls( daemon_provider=self.provide_client): setattr(self, call_name, f) # this is temporary while we are killing the daemon # the cli needs access to the new direct_companion_manager to route direct # commands. # this overrides the stub to talk directly to the companion self.direct_companion_manager = DirectCompanionManager( logger=self.logger) try: self.companion_info: CompanionInfo = self.direct_companion_manager.get_companion_info( target_udid=self.target_udid) self.logger.info(f"using companion {self.companion_info}") self.channel = Channel( self.companion_info.host, self.companion_info.port, loop=asyncio.get_event_loop(), ) self.stub: CompanionServiceStub = CompanionServiceStub( channel=self.channel) except IdbException as e: self.logger.info(e)
def get_stub_for_address(self, host: str, port: int) -> CompanionServiceStub: self._logger.debug( f"creating grpc stub for companion at {host}:{port}") channel = Channel(host, port, loop=asyncio.get_event_loop()) return CompanionServiceStub(channel)