async def test_clear(self) -> None:
with tempfile.NamedTemporaryFile() as f:
with open(f.name, "w") as temp:
json.dump(json_data_companions([]), temp)
companion_manager = DirectCompanionManager(logger=mock.MagicMock(),
state_file_path=f.name)
companion = CompanionInfo(udid="asdasda",
host="foohost",
port=123,
is_local=False)
companion_manager.add_companion(companion)
companion_manager.clear()
companions = companion_manager.get_companions()
self.assertEqual(companions, [])
class GrpcClient(IdbClient):
def __init__(
self, target_udid: Optional[str], logger: Optional[logging.Logger] = None
) -> None:
self.logger: logging.Logger = (
logger if logger else logging.getLogger("idb_grpc_client")
)
self.target_udid = target_udid
self.direct_companion_manager = DirectCompanionManager(logger=self.logger)
self.local_targets_manager = LocalTargetsManager(logger=self.logger)
self.companion_info: Optional[CompanionInfo] = None
async def spawn_notifier(self) -> None:
if platform == "darwin":
companion_spawner = CompanionSpawner(
companion_path="/usr/local/bin/idb_companion", logger=self.logger
)
await companion_spawner.spawn_notifier()
@asynccontextmanager
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 spawn_companion(self, target_udid: str) -> Optional[CompanionInfo]:
if (
self.local_targets_manager.is_local_target_available(
target_udid=target_udid
)
or target_udid == "mac"
):
companion_spawner = CompanionSpawner(
companion_path="/usr/local/bin/idb_companion", logger=self.logger
)
self.logger.info(f"will attempt to spawn a companion for {target_udid}")
port = await companion_spawner.spawn_companion(target_udid=target_udid)
if port:
self.logger.info(f"spawned a companion for {target_udid}")
host = "localhost"
companion_info = CompanionInfo(
host=host, port=port, udid=target_udid, is_local=True
)
self.direct_companion_manager.add_companion(companion_info)
return companion_info
return None
@property
def metadata(self) -> Dict[str, str]:
if self.target_udid:
# pyre-fixme[7]: Expected `Dict[str, str]` but got `Dict[str,
# Optional[str]]`.
return {"udid": self.target_udid}
else:
return {}
async def kill(self) -> None:
self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()
@log_and_handle_exceptions
async def list_apps(self) -> List[InstalledAppInfo]:
async with self.get_stub() as stub:
response = await stub.list_apps(ListAppsRequest())
return [
InstalledAppInfo(
bundle_id=app.bundle_id,
name=app.name,
architectures=app.architectures,
install_type=app.install_type,
process_state=AppProcessState(app.process_state),
debuggable=app.debuggable,
)
for app in response.apps
]
@log_and_handle_exceptions
async def accessibility_info(
self, point: Optional[Tuple[int, int]]
) -> AccessibilityInfo:
async with self.get_stub() as stub:
grpc_point = Point(x=point[0], y=point[1]) if point is not None else None
response = await stub.accessibility_info(
AccessibilityInfoRequest(point=grpc_point)
)
return AccessibilityInfo(json=response.json)
@log_and_handle_exceptions
async def add_media(self, file_paths: List[str]) -> None:
async with self.get_stub() as stub, stub.add_media.open() as stream:
if none_throws(self.companion_info).is_local:
for file_path in file_paths:
await stream.send_message(
AddMediaRequest(payload=Payload(file_path=file_path))
)
await stream.end()
await stream.recv_message()
else:
generator = stream_map(
generate_tar(paths=file_paths, place_in_subfolders=True),
lambda chunk: AddMediaRequest(payload=Payload(data=chunk)),
)
await drain_to_stream(
stream=stream, generator=generator, logger=self.logger
)
@log_and_handle_exceptions
async def approve(self, bundle_id: str, permissions: Set[str]) -> None:
async with self.get_stub() as stub:
await stub.approve(
ApproveRequest(
bundle_id=bundle_id,
permissions=[APPROVE_MAP[permission] for permission in permissions],
)
)
@log_and_handle_exceptions
async def clear_keychain(self) -> None:
async with self.get_stub() as stub:
await stub.clear_keychain(ClearKeychainRequest())
@log_and_handle_exceptions
async def contacts_update(self, contacts_path: str) -> None:
async with self.get_stub() as stub:
data = await create_tar([contacts_path])
await stub.contacts_update(
ContactsUpdateRequest(payload=Payload(data=data))
)
@log_and_handle_exceptions
async def screenshot(self) -> bytes:
async with self.get_stub() as stub:
response = await stub.screenshot(ScreenshotRequest())
return response.image_data
@log_and_handle_exceptions
async def set_location(self, latitude: float, longitude: float) -> None:
async with self.get_stub() as stub:
await stub.set_location(
SetLocationRequest(
location=Location(latitude=latitude, longitude=longitude)
)
)
@log_and_handle_exceptions
async def terminate(self, bundle_id: str) -> None:
async with self.get_stub() as stub:
await stub.terminate(TerminateRequest(bundle_id=bundle_id))
@log_and_handle_exceptions
async def describe(self) -> TargetDescription:
async with self.get_stub() as stub:
response = await stub.describe(TargetDescriptionRequest())
return target_to_py(response.target_description)
@log_and_handle_exceptions
async def focus(self) -> None:
async with self.get_stub() as stub:
await stub.focus(FocusRequest())
@log_and_handle_exceptions
async def open_url(self, url: str) -> None:
async with self.get_stub() as stub:
await stub.open_url(OpenUrlRequest(url=url))
@log_and_handle_exceptions
async def uninstall(self, bundle_id: str) -> None:
async with self.get_stub() as stub:
await stub.uninstall(UninstallRequest(bundle_id=bundle_id))
@log_and_handle_exceptions
async def rm(self, bundle_id: str, paths: List[str]) -> None:
async with self.get_stub() as stub:
await stub.rm(RmRequest(bundle_id=bundle_id, paths=paths))
@log_and_handle_exceptions
async def mv(self, bundle_id: str, src_paths: List[str], dest_path: str) -> None:
async with self.get_stub() as stub:
await stub.mv(
MvRequest(bundle_id=bundle_id, src_paths=src_paths, dst_path=dest_path)
)
@log_and_handle_exceptions
async def ls(self, bundle_id: str, path: str) -> List[FileEntryInfo]:
async with self.get_stub() as stub:
response = await stub.ls(LsRequest(bundle_id=bundle_id, path=path))
return [FileEntryInfo(path=file.path) for file in response.files]
@log_and_handle_exceptions
async def mkdir(self, bundle_id: str, path: str) -> None:
async with self.get_stub() as stub:
await stub.mkdir(MkdirRequest(bundle_id=bundle_id, path=path))
@log_and_handle_exceptions
async def crash_delete(self, query: CrashLogQuery) -> List[CrashLogInfo]:
async with self.get_stub() as stub:
response = await stub.crash_delete(_to_crash_log_query_proto(query))
return _to_crash_log_info_list(response)
@log_and_handle_exceptions
async def crash_list(self, query: CrashLogQuery) -> List[CrashLogInfo]:
async with self.get_stub() as stub:
response = await stub.crash_list(_to_crash_log_query_proto(query))
return _to_crash_log_info_list(response)
@log_and_handle_exceptions
async def crash_show(self, name: str) -> CrashLog:
async with self.get_stub() as stub:
response = await stub.crash_show(CrashShowRequest(name=name))
return _to_crash_log(response)
@log_and_handle_exceptions
async def install(self, bundle: Bundle) -> AsyncIterator[InstalledArtifact]:
async for response in self._install_to_destination(
bundle=bundle, destination=InstallRequest.APP
):
yield response
@log_and_handle_exceptions
async def install_xctest(self, xctest: Bundle) -> AsyncIterator[InstalledArtifact]:
async for response in self._install_to_destination(
bundle=xctest, destination=InstallRequest.XCTEST
):
yield response
@log_and_handle_exceptions
async def install_dylib(self, dylib: Bundle) -> AsyncIterator[InstalledArtifact]:
async for response in self._install_to_destination(
bundle=dylib, destination=InstallRequest.DYLIB
):
yield response
@log_and_handle_exceptions
async def install_dsym(self, dsym: Bundle) -> AsyncIterator[InstalledArtifact]:
async for response in self._install_to_destination(
bundle=dsym, destination=InstallRequest.DSYM
):
yield response
@log_and_handle_exceptions
async def install_framework(
self, framework_path: Bundle
) -> AsyncIterator[InstalledArtifact]:
async for response in self._install_to_destination(
bundle=framework_path, destination=InstallRequest.FRAMEWORK
):
yield response
async def _install_to_destination(
self, bundle: Bundle, destination: Destination
) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub, stub.install.open() as stream:
generator = None
if isinstance(bundle, str):
url = urllib.parse.urlparse(bundle)
if url.scheme:
# send url
payload = Payload(url=bundle)
generator = generate_requests([InstallRequest(payload=payload)])
else:
file_path = str(Path(bundle).resolve(strict=True))
if none_throws(self.companion_info).is_local:
# send file_path
generator = generate_requests(
[InstallRequest(payload=Payload(file_path=file_path))]
)
else:
# chunk file from file_path
generator = generate_binary_chunks(
path=file_path, destination=destination, logger=self.logger
)
else:
# chunk file from memory
generator = generate_io_chunks(io=bundle, logger=self.logger)
# stream to companion
await stream.send_message(InstallRequest(destination=destination))
async for message in generator:
await stream.send_message(message)
await stream.end()
async for response in stream:
yield InstalledArtifact(
name=response.name, uuid=response.uuid, progress=response.progress
)
@log_and_handle_exceptions
async def push(self, src_paths: List[str], bundle_id: str, dest_path: str) -> None:
async with self.get_stub() as stub, stub.push.open() as stream:
await stream.send_message(
PushRequest(
inner=PushRequest.Inner(bundle_id=bundle_id, dst_path=dest_path)
)
)
if none_throws(self.companion_info).is_local:
for src_path in src_paths:
await stream.send_message(
PushRequest(payload=Payload(file_path=src_path))
)
await stream.end()
await stream.recv_message()
else:
await drain_to_stream(
stream=stream,
generator=stream_map(
generate_tar(paths=src_paths),
lambda chunk: PushRequest(payload=Payload(data=chunk)),
),
logger=self.logger,
)
@log_and_handle_exceptions
async def pull(self, bundle_id: str, src_path: str, dest_path: str) -> None:
async with self.get_stub() as stub, stub.pull.open() as stream:
request = request = PullRequest(
bundle_id=bundle_id,
src_path=src_path,
# not sending the destination to remote companion
# so it streams the file back
dst_path=dest_path
if none_throws(self.companion_info).is_local
else None,
)
await stream.send_message(request)
await stream.end()
if none_throws(self.companion_info).is_local:
await stream.recv_message()
else:
await drain_untar(generate_bytes(stream), output_path=dest_path)
self.logger.info(f"pulled file to {dest_path}")
@log_and_handle_exceptions
async def list_test_bundle(self, test_bundle_id: str, app_path: str) -> List[str]:
async with self.get_stub() as stub:
response = await stub.xctest_list_tests(
XctestListTestsRequest(bundle_name=test_bundle_id, app_path=app_path)
)
return [name for name in response.names]
@log_and_handle_exceptions
async def list_xctests(self) -> List[InstalledTestInfo]:
async with self.get_stub() as stub:
response = await stub.xctest_list_bundles(XctestListBundlesRequest())
return [
InstalledTestInfo(
bundle_id=bundle.bundle_id,
name=bundle.name,
architectures=bundle.architectures,
)
for bundle in response.bundles
]
@log_and_handle_exceptions
async def send_events(self, events: Iterable[HIDEvent]) -> None:
await self.hid(iterator_to_async_iterator(events))
@log_and_handle_exceptions
async def tap(self, x: int, y: int, duration: Optional[float] = None) -> None:
await self.send_events(tap_to_events(x, y, duration))
@log_and_handle_exceptions
async def button(
self, button_type: HIDButtonType, duration: Optional[float] = None
) -> None:
await self.send_events(button_press_to_events(button_type, duration))
@log_and_handle_exceptions
async def key(self, keycode: int, duration: Optional[float] = None) -> None:
await self.send_events(key_press_to_events(keycode, duration))
@log_and_handle_exceptions
async def text(self, text: str) -> None:
await self.send_events(text_to_events(text))
@log_and_handle_exceptions
async def swipe(
self,
p_start: Tuple[int, int],
p_end: Tuple[int, int],
delta: Optional[int] = None,
) -> None:
await self.send_events(swipe_to_events(p_start, p_end, delta))
@log_and_handle_exceptions
async def key_sequence(self, key_sequence: List[int]) -> None:
events: List[HIDEvent] = []
for key in key_sequence:
events.extend(key_press_to_events(key))
await self.send_events(events)
@log_and_handle_exceptions
async def hid(self, event_iterator: AsyncIterable[HIDEvent]) -> None:
async with self.get_stub() as stub, stub.hid.open() as stream:
grpc_event_iterator = (
event_to_grpc(event) async for event in event_iterator
)
await drain_to_stream(
stream=stream, generator=grpc_event_iterator, logger=self.logger
)
await stream.recv_message()
async def debug_server(self, request: DebugServerRequest) -> DebugServerResponse:
async with self.get_stub() as stub, stub.debugserver.open() as stream:
await stream.send_message(request)
await stream.end()
return await stream.recv_message()
@log_and_handle_exceptions
async def debugserver_start(self, bundle_id: str) -> List[str]:
response = await self.debug_server(
request=DebugServerRequest(
start=DebugServerRequest.Start(bundle_id=bundle_id)
)
)
return response.status.lldb_bootstrap_commands
@log_and_handle_exceptions
async def debugserver_stop(self) -> None:
await self.debug_server(
request=DebugServerRequest(stop=DebugServerRequest.Stop())
)
@log_and_handle_exceptions
async def debugserver_status(self) -> Optional[List[str]]:
response = await self.debug_server(
request=DebugServerRequest(status=DebugServerRequest.Status())
)
commands = response.status.lldb_bootstrap_commands
return commands if commands else None
@log_and_handle_exceptions
async def run_instruments(
self,
stop: asyncio.Event,
trace_basename: str,
template_name: str,
app_bundle_id: str,
app_environment: Optional[Dict[str, str]] = None,
app_arguments: Optional[List[str]] = None,
tool_arguments: Optional[List[str]] = None,
started: Optional[asyncio.Event] = None,
timings: Optional[InstrumentsTimings] = None,
post_process_arguments: Optional[List[str]] = None,
) -> List[str]:
self.logger.info(f"Starting instruments connection")
async with self.get_stub() as stub, stub.instruments_run.open() as stream:
self.logger.info("Sending instruments request")
await stream.send_message(
InstrumentsRunRequest(
start=InstrumentsRunRequest.Start(
template_name=template_name,
app_bundle_id=app_bundle_id,
environment=app_environment,
arguments=app_arguments,
tool_arguments=tool_arguments,
timings=translate_instruments_timings(timings),
)
)
)
self.logger.info("Starting instruments")
await instruments_drain_until_running(stream=stream, logger=self.logger)
if started:
started.set()
self.logger.info("Instruments has started, waiting for stop")
async for response in stop_wrapper(stream=stream, stop=stop):
output = response.log_output
if len(output):
self.logger.info(output.decode())
self.logger.info("Stopping instruments")
await stream.send_message(
InstrumentsRunRequest(
stop=InstrumentsRunRequest.Stop(
post_process_arguments=post_process_arguments
)
)
)
await stream.end()
result = []
with tempfile.TemporaryDirectory() as tmp_trace_dir:
self.logger.info(
f"Writing instruments data from tar to {tmp_trace_dir}"
)
await drain_untar(
instruments_generate_bytes(stream=stream, logger=self.logger),
output_path=tmp_trace_dir,
)
if os.path.exists(
os.path.join(os.path.abspath(tmp_trace_dir), "instrument_data")
):
# tar is an instruments trace (old behavior)
trace_file = f"{trace_basename}.trace"
shutil.copytree(tmp_trace_dir, trace_file)
result.append(trace_file)
self.logger.info(f"Trace written to {trace_file}")
else:
# tar is a folder containing one or more trace files
for file in os.listdir(tmp_trace_dir):
_, file_extension = os.path.splitext(file)
tmp_trace_file = os.path.join(
os.path.abspath(tmp_trace_dir), file
)
trace_file = f"{trace_basename}{file_extension}"
shutil.move(tmp_trace_file, trace_file)
result.append(trace_file)
self.logger.info(f"Trace written to {trace_file}")
return result
@log_and_handle_exceptions
async def launch(
self,
bundle_id: str,
args: Optional[List[str]] = None,
env: Optional[Dict[str, str]] = None,
foreground_if_running: bool = False,
stop: Optional[asyncio.Event] = None,
) -> None:
async with self.get_stub() as stub, stub.launch.open() as stream:
request = LaunchRequest(
start=LaunchRequest.Start(
bundle_id=bundle_id,
env=env,
app_args=args,
foreground_if_running=foreground_if_running,
wait_for=True if stop else False,
)
)
await stream.send_message(request)
if stop:
await asyncio.gather(
drain_launch_stream(stream), end_launch_stream(stream, stop)
)
else:
await stream.end()
await drain_launch_stream(stream)
@log_and_handle_exceptions
async def record_video(self, stop: asyncio.Event, output_file: str) -> None:
self.logger.info(f"Starting connection to backend")
async with self.get_stub() as stub, stub.record.open() as stream:
if none_throws(self.companion_info).is_local:
self.logger.info(
f"Starting video recording to local file {output_file}"
)
await stream.send_message(
RecordRequest(start=RecordRequest.Start(file_path=output_file))
)
else:
self.logger.info(f"Starting video recording with response data")
await stream.send_message(
RecordRequest(start=RecordRequest.Start(file_path=None))
)
await stop.wait()
self.logger.info("Stopping video recording")
await stream.send_message(RecordRequest(stop=RecordRequest.Stop()))
await stream.end()
if none_throws(self.companion_info).is_local:
self.logger.info("Video saved at output path")
await stream.recv_message()
else:
self.logger.info(f"Decompressing gzip to {output_file}")
await drain_gzip_decompress(
generate_video_bytes(stream), output_path=output_file
)
self.logger.info(f"Finished decompression to {output_file}")
@log_and_handle_exceptions
async def run_xctest(
self,
test_bundle_id: str,
app_bundle_id: str,
test_host_app_bundle_id: Optional[str] = None,
is_ui_test: bool = False,
is_logic_test: bool = False,
tests_to_run: Optional[Set[str]] = None,
tests_to_skip: Optional[Set[str]] = None,
env: Optional[Dict[str, str]] = None,
args: Optional[List[str]] = None,
result_bundle_path: Optional[str] = None,
idb_log_buffer: Optional[StringIO] = None,
timeout: Optional[int] = None,
poll_interval_sec: float = TESTS_POLL_INTERVAL,
) -> AsyncIterator[TestRunInfo]:
async with self.get_stub() as stub, stub.xctest_run.open() as stream:
request = make_request(
test_bundle_id=test_bundle_id,
app_bundle_id=app_bundle_id,
test_host_app_bundle_id=test_host_app_bundle_id,
is_ui_test=is_ui_test,
is_logic_test=is_logic_test,
tests_to_run=tests_to_run,
tests_to_skip=tests_to_skip,
env=env,
args=args,
result_bundle_path=result_bundle_path,
timeout=timeout,
)
await stream.send_message(request)
await stream.end()
async for response in stream:
# response.log_output is a container of strings.
# google.protobuf.pyext._message.RepeatedScalarContainer.
for line in [
line
for lines in response.log_output
for line in lines.splitlines(keepends=True)
]:
self.logger.info(line)
if idb_log_buffer:
idb_log_buffer.write(line)
if result_bundle_path:
await write_result_bundle(
response=response,
output_path=result_bundle_path,
logger=self.logger,
)
for result in make_results(response):
yield result
async def _tail_specific_logs(
self,
source: LogRequest.Source,
stop: asyncio.Event,
arguments: Optional[List[str]],
) -> AsyncIterator[str]:
async with self.get_stub() as stub, stub.log.open() as stream:
await stream.send_message(
LogRequest(arguments=arguments, source=source), end=True
)
async for message in cancel_wrapper(stream=stream, stop=stop):
yield message.output.decode()
@log_and_handle_exceptions
async def tail_logs(
self, stop: asyncio.Event, arguments: Optional[List[str]] = None
) -> AsyncIterator[str]:
async for message in self._tail_specific_logs(
source=LogRequest.TARGET, stop=stop, arguments=arguments
):
yield message
@log_and_handle_exceptions
async def tail_companion_logs(self, stop: asyncio.Event) -> AsyncIterator[str]:
async for message in self._tail_specific_logs(
source=LogRequest.COMPANION, stop=stop, arguments=None
):
yield message
@log_and_handle_exceptions
async def boot(self) -> None:
if self.target_udid:
cmd: List[str] = [
"/usr/local/bin/idb_companion",
"--boot",
none_throws(self.target_udid),
]
process = await asyncio.create_subprocess_exec(
*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
)
await process.communicate()
else:
raise IdbException("boot needs --udid to work")
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
@log_and_handle_exceptions
async def list_targets(self) -> List[TargetDescription]:
await self.spawn_notifier()
companions = self.direct_companion_manager.get_companions()
local_targets = self.local_targets_manager.get_local_targets()
connected_targets = await asyncio.gather(
*(
self._companion_to_target(companion=companion)
for companion in companions
)
)
return local_targets + [
target for target in connected_targets if target is not None
]
@log_and_handle_exceptions
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}")
@log_and_handle_exceptions
async def disconnect(self, destination: ConnectionDestination) -> None:
self.direct_companion_manager.remove_companion(destination)
class IdbManagementClient(IdbManagementClientBase):
def __init__(
self,
companion_path: Optional[str] = None,
device_set_path: Optional[str] = None,
logger: Optional[logging.Logger] = None,
prune_dead_companion: bool = True,
companion_command_timeout: int = DEFAULT_COMPANION_COMMAND_TIMEOUT,
) -> None:
os.makedirs(BASE_IDB_FILE_PATH, exist_ok=True)
self.companion_path = companion_path
self.device_set_path = device_set_path
self.logger: logging.Logger = (logger if logger else
logging.getLogger("idb_grpc_client"))
self._prune_dead_companion = prune_dead_companion
self._companion_command_timeout = companion_command_timeout
self.direct_companion_manager = DirectCompanionManager(
logger=self.logger)
self.local_targets_manager = LocalTargetsManager(logger=self.logger)
async def _spawn_notifier(self) -> None:
companion_path = self.companion_path
if companion_path:
companion_spawner = CompanionSpawner(companion_path=companion_path,
logger=self.logger)
await companion_spawner.spawn_notifier()
async def _spawn_companion(self,
target_udid: str) -> Optional[CompanionInfo]:
companion_path = self.companion_path
if companion_path is None:
return None
if (self.local_targets_manager.is_local_target_available(
target_udid=target_udid) or target_udid == "mac"):
companion_spawner = CompanionSpawner(companion_path=companion_path,
logger=self.logger)
self.logger.info(
f"will attempt to spawn a companion for {target_udid}")
port = await companion_spawner.spawn_companion(
target_udid=target_udid)
if port:
self.logger.info(f"spawned a companion for {target_udid}")
host = "localhost"
companion_info = CompanionInfo(host=host,
port=port,
udid=target_udid,
is_local=True)
await self.direct_companion_manager.add_companion(
companion_info)
return companion_info
return None
async def _companion_to_target(
self, companion: CompanionInfo) -> Optional[TargetDescription]:
try:
async with IdbClient.build(
host=companion.host,
port=companion.port,
is_local=False,
logger=self.logger,
) as client:
return await client.describe()
except Exception:
if not self._prune_dead_companion:
self.logger.warning(
f"Failed to describe {companion}, but not removing it")
return None
self.logger.warning(f"Failed to describe {companion}, removing it")
await self.direct_companion_manager.remove_companion(
Address(host=companion.host, port=companion.port))
return None
@asynccontextmanager
async def _start_companion_command(
self, arguments: List[str]
) -> AsyncContextManager[asyncio.subprocess.Process]:
companion_path = self.companion_path
if companion_path is None:
if platform == "darwin":
raise IdbException("Companion path not provided")
else:
raise IdbException(
"Companion interactions do not work on non-macOS platforms"
)
cmd: List[str] = [companion_path]
device_set_path = self.device_set_path
if device_set_path is not None:
cmd.extend(["--device-set-path", device_set_path])
cmd.extend(arguments)
process = await asyncio.create_subprocess_exec(*cmd,
stdout=subprocess.PIPE,
stderr=None)
try:
yield process
finally:
await _terminate_process(process=process, logger=self.logger)
async def _run_companion_command(self, arguments: List[str]) -> str:
timeout = self._companion_command_timeout
async with self._start_companion_command(
arguments=arguments) as process:
try:
(output, _) = await asyncio.wait_for(process.communicate(),
timeout=timeout)
if process.returncode != 0:
raise IdbException(f"Failed to run {arguments}")
self.logger.info(f"Ran {arguments} successfully.")
return output.decode()
except asyncio.TimeoutError:
raise IdbException(
f"Timed out after {timeout} secs on command {' '.join(arguments)}"
)
async def _run_udid_command(self, udid: str, command: str) -> str:
return await self._run_companion_command(
arguments=[f"--{command}", udid])
@asynccontextmanager
async def from_udid(self,
udid: Optional[str]) -> AsyncContextManager[IdbClient]:
await self._spawn_notifier()
try:
companion_info = await self.direct_companion_manager.get_companion_info(
target_udid=udid)
except IdbException as e:
# will try to spawn a companion if on mac.
companion_info = await self._spawn_companion(
target_udid=none_throws(udid))
if companion_info is None:
raise e
async with IdbClient.build(
host=companion_info.host,
port=companion_info.port,
is_local=companion_info.is_local,
logger=self.logger,
) as client:
yield client
@log_call()
async def list_targets(self) -> List[TargetDescription]:
(_, companions) = await asyncio.gather(
self._spawn_notifier(),
self.direct_companion_manager.get_companions())
connected_targets = [
target for target in (await asyncio.gather(
*(self._companion_to_target(companion=companion)
for companion in companions))) if target is not None
]
return merge_connected_targets(
local_targets=self.local_targets_manager.get_local_targets(),
connected_targets=connected_targets,
)
@log_call()
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):
async with IdbClient.build(
host=destination.host,
port=destination.port,
is_local=False,
logger=self.logger,
) as client:
with tempfile.NamedTemporaryFile(mode="w+b") as f:
response = await client.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}")
await self.direct_companion_manager.add_companion(companion)
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}")
@log_call()
async def disconnect(self, destination: ConnectionDestination) -> None:
await self.direct_companion_manager.remove_companion(destination)
@log_call()
async def create(self, device_type: str, os_version: str) -> str:
output = await self._run_companion_command(
arguments=["--create", f"{device_type},{os_version}"])
created = json.loads(output.splitlines()[-1])
return created["udid"]
@log_call()
async def boot(self, udid: str) -> None:
await self._run_udid_command(udid=udid, command="boot")
@asynccontextmanager
async def boot_headless(self, udid: str) -> AsyncContextManager[None]:
async with self._start_companion_command(
["--headless", "1", "--boot", udid]) as process:
# The first line written to stdout is information about the booted sim.
line = await none_throws(process.stdout).readline()
target = json.loads(line.decode())
assert target["udid"] == udid
self.logger.info(f"{udid} is now booted")
yield None
self.logger.info(f"Done with {udid}. Shutting down.")
@log_call()
async def shutdown(self, udid: str) -> None:
await self._run_udid_command(udid=udid, command="shutdown")
@log_call()
async def erase(self, udid: str) -> None:
await self._run_udid_command(udid=udid, command="erase")
@log_call()
async def clone(self, udid: str) -> str:
output = await self._run_udid_command(udid=udid, command="clone")
cloned = json.loads(output.splitlines()[-1])
return cloned["udid"]
@log_call()
async def delete(self, udid: Optional[str]) -> None:
await self._run_udid_command(udid=udid if udid is not None else "all",
command="delete")
@log_call()
async def kill(self) -> None:
await self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()
class IdbManagementClient(IdbManagementClientBase):
def __init__(
self,
companion_path: Optional[str] = None,
device_set_path: Optional[str] = None,
prune_dead_companion: bool = True,
logger: Optional[logging.Logger] = None,
) -> None:
os.makedirs(BASE_IDB_FILE_PATH, exist_ok=True)
self._logger: logging.Logger = (logger if logger else
logging.getLogger("idb_grpc_client"))
self._direct_companion_manager = DirectCompanionManager(
logger=self._logger)
self._local_targets_manager = LocalTargetsManager(logger=self._logger)
self._companion_spawner: Optional[CompanionSpawner] = (
CompanionSpawner(companion_path=companion_path,
logger=self._logger)
if companion_path is not None else None)
self._prune_dead_companion = prune_dead_companion
async def _spawn_notifier(self) -> None:
companion_spawner = self._companion_spawner
if companion_spawner is None:
return
await companion_spawner.spawn_notifier()
async def _spawn_companion(self,
target_udid: str) -> Optional[CompanionInfo]:
companion_spawner = self._companion_spawner
if companion_spawner is None:
return None
local_target_available = await self._local_targets_manager.is_local_target_available(
target_udid=target_udid)
if local_target_available or target_udid == "mac":
self._logger.info(
f"will attempt to spawn a companion for {target_udid}")
port = await companion_spawner.spawn_companion(
target_udid=target_udid)
if port:
self._logger.info(f"spawned a companion for {target_udid}")
host = "localhost"
companion_info = CompanionInfo(
address=TCPAddress(host=host, port=port),
udid=target_udid,
is_local=True,
)
await self._direct_companion_manager.add_companion(
companion_info)
return companion_info
return None
async def _companion_to_target(
self, companion: CompanionInfo) -> Optional[TargetDescription]:
try:
async with IdbClient.build(address=companion.address,
is_local=False,
logger=self._logger) as client:
return await client.describe()
except Exception:
if not self._prune_dead_companion:
self._logger.warning(
f"Failed to describe {companion}, but not removing it")
return None
self._logger.warning(
f"Failed to describe {companion}, removing it")
await self._direct_companion_manager.remove_companion(
companion.address)
return None
@asynccontextmanager
async def from_udid(
self, udid: Optional[str]) -> AsyncGenerator[IdbClient, None]:
await self._spawn_notifier()
try:
companion_info = await self._direct_companion_manager.get_companion_info(
target_udid=udid)
except IdbException as e:
# will try to spawn a companion if on mac.
if udid is None:
raise e
companion_info = await self._spawn_companion(target_udid=udid)
if companion_info is None:
raise e
async with IdbClient.build(
address=companion_info.address,
is_local=companion_info.is_local,
logger=self._logger,
) as client:
yield client
@log_call()
async def list_targets(self) -> List[TargetDescription]:
await self._spawn_notifier()
(companions, local_targets) = await asyncio.gather(
self._direct_companion_manager.get_companions(),
self._local_targets_manager.get_local_targets(),
)
connected_targets = [
target for target in (await asyncio.gather(
*(self._companion_to_target(companion=companion)
for companion in companions))) if target is not None
]
return merge_connected_targets(local_targets=local_targets,
connected_targets=connected_targets)
@log_call()
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, TCPAddress) or isinstance(
destination, DomainSocketAddress):
async with IdbClient.build(address=destination,
is_local=False,
logger=self._logger) as client:
with tempfile.NamedTemporaryFile(mode="w+b") as f:
response = await client.stub.connect(
ConnectRequest(metadata=metadata,
local_file_path=f.name))
companion = CompanionInfo(
address=destination,
udid=response.companion.udid,
is_local=response.companion.is_local,
)
self._logger.debug(f"Connected directly to {companion}")
await self._direct_companion_manager.add_companion(companion)
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}")
@log_call()
async def disconnect(self, destination: ConnectionDestination) -> None:
await self._direct_companion_manager.remove_companion(destination)
@log_call()
async def kill(self) -> None:
await self._direct_companion_manager.clear()
await self._local_targets_manager.clear()
PidSaver(logger=self._logger).kill_saved_pids()
class GrpcClient(IdbClient):
def __init__(self,
target_udid: Optional[str],
logger: Optional[logging.Logger] = None) -> None:
self.logger: logging.Logger = (logger if logger else
logging.getLogger("idb_grpc_client"))
self.target_udid = target_udid
self.direct_companion_manager = DirectCompanionManager(
logger=self.logger)
self.local_targets_manager = LocalTargetsManager(logger=self.logger)
self.companion_info: Optional[CompanionInfo] = None
async def spawn_notifier(self) -> None:
if platform == "darwin" and os.path.exists(
"/usr/local/bin/idb_companion"):
companion_spawner = CompanionSpawner(
companion_path="/usr/local/bin/idb_companion",
logger=self.logger)
await companion_spawner.spawn_notifier()
@asynccontextmanager
async def get_stub(self) -> AsyncContextManager[GrpcStubClient]:
await self.spawn_notifier()
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
async with GrpcStubClient.build(companion_info=none_throws(
self.companion_info),
logger=self.logger) as stub:
yield stub
async def spawn_companion(self,
target_udid: str) -> Optional[CompanionInfo]:
if (self.local_targets_manager.is_local_target_available(
target_udid=target_udid) or target_udid == "mac"):
companion_spawner = CompanionSpawner(
companion_path="/usr/local/bin/idb_companion",
logger=self.logger)
self.logger.info(
f"will attempt to spawn a companion for {target_udid}")
port = await companion_spawner.spawn_companion(
target_udid=target_udid)
if port:
self.logger.info(f"spawned a companion for {target_udid}")
host = "localhost"
companion_info = CompanionInfo(host=host,
port=port,
udid=target_udid,
is_local=True)
self.direct_companion_manager.add_companion(companion_info)
return companion_info
return None
@property
def metadata(self) -> Dict[str, str]:
if self.target_udid:
# pyre-fixme[7]: Expected `Dict[str, str]` but got `Dict[str,
# Optional[str]]`.
return {"udid": self.target_udid}
else:
return {}
async def kill(self) -> None:
self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()
@log_and_handle_exceptions
async def list_apps(self) -> List[InstalledAppInfo]:
async with self.get_stub() as stub:
return await stub.list_apps()
@log_and_handle_exceptions
async def accessibility_info(
self, point: Optional[Tuple[int, int]]) -> AccessibilityInfo:
async with self.get_stub() as stub:
return await stub.accessibility_info(point=point)
@log_and_handle_exceptions
async def add_media(self, file_paths: List[str]) -> None:
async with self.get_stub() as stub:
return await stub.add_media(file_paths=file_paths)
@log_and_handle_exceptions
async def approve(self, bundle_id: str, permissions: Set[str]) -> None:
async with self.get_stub() as stub:
return await stub.approve(bundle_id, permissions)
@log_and_handle_exceptions
async def clear_keychain(self) -> None:
async with self.get_stub() as stub:
await stub.clear_keychain()
@log_and_handle_exceptions
async def contacts_update(self, contacts_path: str) -> None:
async with self.get_stub() as stub:
return await stub.contacts_update(contacts_path=contacts_path)
@log_and_handle_exceptions
async def screenshot(self) -> bytes:
async with self.get_stub() as stub:
return await stub.screenshot()
@log_and_handle_exceptions
async def set_location(self, latitude: float, longitude: float) -> None:
async with self.get_stub() as stub:
await stub.set_location(latitude=latitude, longitude=longitude)
@log_and_handle_exceptions
async def terminate(self, bundle_id: str) -> None:
async with self.get_stub() as stub:
await stub.terminate(bundle_id=bundle_id)
@log_and_handle_exceptions
async def describe(self) -> TargetDescription:
async with self.get_stub() as stub:
return await stub.describe()
@log_and_handle_exceptions
async def focus(self) -> None:
async with self.get_stub() as stub:
return await stub.focus()
@log_and_handle_exceptions
async def open_url(self, url: str) -> None:
async with self.get_stub() as stub:
return await stub.open_url(url=url)
@log_and_handle_exceptions
async def uninstall(self, bundle_id: str) -> None:
async with self.get_stub() as stub:
return await stub.uninstall(bundle_id=bundle_id)
@log_and_handle_exceptions
async def rm(self, bundle_id: str, paths: List[str]) -> None:
async with self.get_stub() as stub:
return await stub.rm(bundle_id=bundle_id, paths=paths)
@log_and_handle_exceptions
async def mv(self, bundle_id: str, src_paths: List[str],
dest_path: str) -> None:
async with self.get_stub() as stub:
return await stub.mv(bundle_id=bundle_id,
src_paths=src_paths,
dest_path=dest_path)
@log_and_handle_exceptions
async def ls(self, bundle_id: str, path: str) -> List[FileEntryInfo]:
async with self.get_stub() as stub:
return await stub.ls(bundle_id=bundle_id, path=path)
@log_and_handle_exceptions
async def mkdir(self, bundle_id: str, path: str) -> None:
async with self.get_stub() as stub:
return await stub.mkdir(bundle_id=bundle_id, path=path)
@log_and_handle_exceptions
async def crash_delete(self, query: CrashLogQuery) -> List[CrashLogInfo]:
async with self.get_stub() as stub:
return await stub.crash_delete(query=query)
@log_and_handle_exceptions
async def crash_list(self, query: CrashLogQuery) -> List[CrashLogInfo]:
async with self.get_stub() as stub:
return await stub.crash_list(query=query)
@log_and_handle_exceptions
async def crash_show(self, name: str) -> CrashLog:
async with self.get_stub() as stub:
return await stub.crash_show(name)
@log_and_handle_exceptions
async def install(self,
bundle: Bundle) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub:
async for response in stub.install(bundle=bundle):
yield response
@log_and_handle_exceptions
async def install_xctest(
self, xctest: Bundle) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub:
async for response in stub.install_xctest(xctest=xctest):
yield response
@log_and_handle_exceptions
async def install_dylib(self,
dylib: Bundle) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub:
async for response in stub.install_dylib(dylib=dylib):
yield response
@log_and_handle_exceptions
async def install_dsym(self,
dsym: Bundle) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub:
async for response in stub.install_dsym(dsym=dsym):
yield response
@log_and_handle_exceptions
async def install_framework(
self, framework_path: Bundle) -> AsyncIterator[InstalledArtifact]:
async with self.get_stub() as stub:
async for response in stub.install_framework(
framework_path=framework_path):
yield response
@log_and_handle_exceptions
async def push(self, src_paths: List[str], bundle_id: str,
dest_path: str) -> None:
async with self.get_stub() as stub:
return await stub.push(src_paths=src_paths,
bundle_id=bundle_id,
dest_path=dest_path)
@log_and_handle_exceptions
async def pull(self, bundle_id: str, src_path: str,
dest_path: str) -> None:
async with self.get_stub() as stub:
return await stub.pull(bundle_id=bundle_id,
src_path=src_path,
dest_path=dest_path)
@log_and_handle_exceptions
async def list_test_bundle(self, test_bundle_id: str,
app_path: str) -> List[str]:
async with self.get_stub() as stub:
return await stub.list_test_bundle(test_bundle_id=test_bundle_id,
app_path=app_path)
@log_and_handle_exceptions
async def list_xctests(self) -> List[InstalledTestInfo]:
async with self.get_stub() as stub:
return await stub.list_xctests()
@log_and_handle_exceptions
async def send_events(self, events: Iterable[HIDEvent]) -> None:
async with self.get_stub() as stub:
return await stub.send_events(events=events)
@log_and_handle_exceptions
async def tap(self,
x: int,
y: int,
duration: Optional[float] = None) -> None:
async with self.get_stub() as stub:
return await stub.tap(x=x, y=y, duration=duration)
@log_and_handle_exceptions
async def button(self,
button_type: HIDButtonType,
duration: Optional[float] = None) -> None:
async with self.get_stub() as stub:
return await stub.button(button_type=button_type,
duration=duration)
@log_and_handle_exceptions
async def key(self,
keycode: int,
duration: Optional[float] = None) -> None:
async with self.get_stub() as stub:
return await stub.key(keycode=keycode, duration=duration)
@log_and_handle_exceptions
async def text(self, text: str) -> None:
async with self.get_stub() as stub:
return await stub.text(text=text)
@log_and_handle_exceptions
async def swipe(
self,
p_start: Tuple[int, int],
p_end: Tuple[int, int],
duration: Optional[float] = None,
delta: Optional[int] = None,
) -> None:
async with self.get_stub() as stub:
return await stub.swipe(p_start=p_start,
p_end=p_end,
duration=duration,
delta=delta)
@log_and_handle_exceptions
async def key_sequence(self, key_sequence: List[int]) -> None:
async with self.get_stub() as stub:
return await stub.key_sequence(key_sequence=key_sequence)
@log_and_handle_exceptions
async def hid(self, event_iterator: AsyncIterable[HIDEvent]) -> None:
async with self.get_stub() as stub:
return await stub.hid(event_iterator=event_iterator)
@log_and_handle_exceptions
async def debugserver_start(self, bundle_id: str) -> List[str]:
async with self.get_stub() as stub:
return await stub.debugserver_start(bundle_id=bundle_id)
@log_and_handle_exceptions
async def debugserver_stop(self) -> None:
async with self.get_stub() as stub:
return await stub.debugserver_stop()
@log_and_handle_exceptions
async def debugserver_status(self) -> Optional[List[str]]:
async with self.get_stub() as stub:
return await stub.debugserver_status()
@log_and_handle_exceptions
async def run_instruments(
self,
stop: asyncio.Event,
trace_basename: str,
template_name: str,
app_bundle_id: str,
app_environment: Optional[Dict[str, str]] = None,
app_arguments: Optional[List[str]] = None,
tool_arguments: Optional[List[str]] = None,
started: Optional[asyncio.Event] = None,
timings: Optional[InstrumentsTimings] = None,
post_process_arguments: Optional[List[str]] = None,
) -> List[str]:
async with self.get_stub() as stub:
return await stub.run_instruments(
stop=stop,
trace_basename=trace_basename,
template_name=template_name,
app_bundle_id=app_bundle_id,
app_environment=app_environment,
app_arguments=app_arguments,
tool_arguments=tool_arguments,
started=started,
timings=timings,
post_process_arguments=post_process_arguments,
)
@log_and_handle_exceptions
async def launch(
self,
bundle_id: str,
args: Optional[List[str]] = None,
env: Optional[Dict[str, str]] = None,
foreground_if_running: bool = False,
stop: Optional[asyncio.Event] = None,
) -> None:
async with self.get_stub() as stub:
return await stub.launch(
bundle_id=bundle_id,
args=args,
env=env,
foreground_if_running=foreground_if_running,
)
@log_and_handle_exceptions
async def record_video(self, stop: asyncio.Event,
output_file: str) -> None:
async with self.get_stub() as stub:
return await stub.record_video(stop=stop, output_file=output_file)
@log_and_handle_exceptions
async def run_xctest(
self,
test_bundle_id: str,
app_bundle_id: str,
test_host_app_bundle_id: Optional[str] = None,
is_ui_test: bool = False,
is_logic_test: bool = False,
tests_to_run: Optional[Set[str]] = None,
tests_to_skip: Optional[Set[str]] = None,
env: Optional[Dict[str, str]] = None,
args: Optional[List[str]] = None,
result_bundle_path: Optional[str] = None,
idb_log_buffer: Optional[StringIO] = None,
timeout: Optional[int] = None,
poll_interval_sec: float = TESTS_POLL_INTERVAL,
) -> AsyncIterator[TestRunInfo]:
async with self.get_stub() as stub:
async for result in stub.run_xctest(
test_bundle_id=test_bundle_id,
app_bundle_id=app_bundle_id,
test_host_app_bundle_id=test_host_app_bundle_id,
is_ui_test=is_ui_test,
is_logic_test=is_logic_test,
tests_to_run=tests_to_run,
tests_to_skip=tests_to_skip,
env=env,
args=args,
result_bundle_path=result_bundle_path,
idb_log_buffer=idb_log_buffer,
timeout=timeout,
poll_interval_sec=poll_interval_sec,
):
yield result
@log_and_handle_exceptions
async def tail_logs(
self,
stop: asyncio.Event,
arguments: Optional[List[str]] = None) -> AsyncIterator[str]:
async with self.get_stub() as stub:
async for message in stub.tail_logs(stop=stop,
arguments=arguments):
yield message
@log_and_handle_exceptions
async def tail_companion_logs(self,
stop: asyncio.Event) -> AsyncIterator[str]:
async with self.get_stub() as stub:
async for message in stub.tail_companion_logs(stop=stop):
yield message
@log_and_handle_exceptions
async def boot(self) -> None:
if self.target_udid:
cmd: List[str] = [
"/usr/local/bin/idb_companion",
"--boot",
none_throws(self.target_udid),
]
process = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE)
await process.communicate()
else:
raise IdbException("boot needs --udid to work")
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
@log_and_handle_exceptions
async def list_targets(self) -> List[TargetDescription]:
await self.spawn_notifier()
companions = self.direct_companion_manager.get_companions()
local_targets = self.local_targets_manager.get_local_targets()
connected_targets = await asyncio.gather(*(self._companion_to_target(
companion=companion) for companion in companions))
return local_targets + [
target for target in connected_targets if target is not None
]
@log_and_handle_exceptions
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}")
@log_and_handle_exceptions
async def disconnect(self, destination: ConnectionDestination) -> None:
self.direct_companion_manager.remove_companion(destination)
class IdbManagementClient(IdbManagementClientBase):
def __init__(self,
companion_path: str,
logger: Optional[logging.Logger] = None) -> None:
self.logger: logging.Logger = (logger if logger else
logging.getLogger("idb_grpc_client"))
self.direct_companion_manager = DirectCompanionManager(
logger=self.logger)
self.local_targets_manager = LocalTargetsManager(logger=self.logger)
self.companion_path = companion_path
async def _spawn_notifier(self) -> None:
if platform == "darwin" and os.path.exists(self.companion_path):
companion_spawner = CompanionSpawner(
companion_path=self.companion_path, logger=self.logger)
await companion_spawner.spawn_notifier()
async def _spawn_companion(self,
target_udid: str) -> Optional[CompanionInfo]:
if (self.local_targets_manager.is_local_target_available(
target_udid=target_udid) or target_udid == "mac"):
companion_spawner = CompanionSpawner(
companion_path=self.companion_path, logger=self.logger)
self.logger.info(
f"will attempt to spawn a companion for {target_udid}")
port = await companion_spawner.spawn_companion(
target_udid=target_udid)
if port:
self.logger.info(f"spawned a companion for {target_udid}")
host = "localhost"
companion_info = CompanionInfo(host=host,
port=port,
udid=target_udid,
is_local=True)
await self.direct_companion_manager.add_companion(
companion_info)
return companion_info
return None
async def _companion_to_target(
self, companion: CompanionInfo) -> Optional[TargetDescription]:
try:
channel = Channel(host=companion.host,
port=companion.port,
loop=asyncio.get_event_loop())
stub = CompanionServiceStub(channel=channel)
response = await stub.describe(TargetDescriptionRequest())
channel.close()
return target_to_py(target=response.target_description,
companion_info=companion)
except Exception:
self.logger.warning(f"Failed to describe {companion}, removing it")
await self.direct_companion_manager.remove_companion(
Address(companion.host, companion.port))
return None
@log_and_handle_exceptions
async def list_targets(self) -> List[TargetDescription]:
(_, companions) = await asyncio.gather(
self._spawn_notifier(),
self.direct_companion_manager.get_companions())
connected_targets = [
target for target in (await asyncio.gather(
*(self._companion_to_target(companion=companion)
for companion in companions))) if target is not None
]
return merge_connected_targets(
local_targets=self.local_targets_manager.get_local_targets(),
connected_targets=connected_targets,
)
@log_and_handle_exceptions
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}")
await 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}")
@log_and_handle_exceptions
async def disconnect(self, destination: ConnectionDestination) -> None:
await self.direct_companion_manager.remove_companion(destination)
@log_and_handle_exceptions
async def boot(self, udid: str) -> None:
cmd: List[str] = [self.companion_path, "--boot", udid]
process = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE)
await process.communicate()
@log_and_handle_exceptions
async def kill(self) -> None:
await self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()