async def run_with_client(self, args: Namespace, client: Client) -> None:
artifact: Optional[InstalledArtifact] = None
compression = (Compression[args.compression]
if args.compression is not None else None)
async for info in client.install(
bundle=args.bundle_path,
make_debuggable=args.make_debuggable,
compression=compression,
override_modification_time=args.override_mtime,
):
artifact = info
progress = info.progress
if progress is None:
continue
self.logger.info(f"Progress: {progress}")
if args.json:
print(
json.dumps({
"installedAppBundleId": none_throws(artifact).name,
"uuid": none_throws(artifact).uuid,
}))
else:
print(
f"Installed: {none_throws(artifact).name} {none_throws(artifact).uuid}"
)
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}")
async def daemon(
client: CompanionClient, stream: Stream[InstallResponse, InstallRequest]
) -> None:
destination_message = none_throws(await stream.recv_message())
payload_message = none_throws(await stream.recv_message())
file_path = payload_message.payload.file_path
url = payload_message.payload.url
data = payload_message.payload.data
destination = destination_message.destination
async with client.stub.install.open() as forward_stream:
await forward_stream.send_message(destination_message)
if client.is_local or len(url):
await forward_stream.send_message(payload_message)
await forward_stream.end()
response = none_throws(await forward_stream.recv_message())
elif file_path:
response = await drain_to_stream(
stream=forward_stream,
generator=_generate_binary_chunks(
path=file_path, destination=destination, logger=client.logger
),
logger=client.logger,
)
elif data:
await forward_stream.send_message(payload_message)
response = await drain_to_stream(
stream=forward_stream, generator=stream, logger=client.logger
)
else:
raise Exception(f"Unrecognised payload message")
await stream.send_message(response)
async def pipe(
self,
input_stream: StreamReader,
output_stream: StreamWriter,
stop: asyncio.Event,
) -> None:
"""
Pipe stdin and stdout to remote dap server
"""
read_future: Optional[asyncio.Future[StreamReader]] = None
write_future: Optional[asyncio.Future[StreamWriter]] = None
stop_future = asyncio.ensure_future(stop.wait())
while True:
if read_future is None:
read_future = asyncio.ensure_future(self._stream.recv_message())
if write_future is None:
write_future = asyncio.ensure_future(
read_next_dap_protocol_message(input_stream)
)
done, pending = await asyncio.wait(
[read_future, write_future, stop_future],
return_when=asyncio.FIRST_COMPLETED,
)
if stop_future in done:
self.logger.debug("Received stop command! Closing stream...")
read_future.cancel()
self.logger.debug("Read future cancelled!")
write_future.cancel()
self.logger.debug("Write future cancelled!")
break
if write_future in done:
data = none_throws(write_future.result())
write_future = None
await self._stream.send_message(
DapRequest(pipe=DapRequest.Pipe(data=data))
)
if read_future in done:
self.logger.debug("Received a message from companion.")
result = none_throws(read_future.result())
read_future = None
if result is None:
# Reached the end of the stream
break
output_stream.write(result.stdout.data)
async def notifierProcess(self) -> asyncio.subprocess.Process:
cmd = [self.notifier_path, "--notify", "1"]
with open(self._log_file_path(), "a") as log_file:
process = await asyncio.create_subprocess_exec(
*cmd, stdout=asyncio.subprocess.PIPE, stderr=log_file)
await self._read_stream(none_throws(process.stdout))
return process
async def _install_to_destination(
self, bundle: Bundle, destination: Destination
) -> 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))
response = await drain_to_stream(
stream=stream, generator=generator, logger=self.logger
)
return InstalledArtifact(name=response.name, uuid=response.uuid)
async def run_with_client(self, args: Namespace, client: Client) -> None:
# Setup
root_command = none_throws(self.root_command)
prompt = "" if args.no_prompt else "idb> "
# Prompt loop
while True:
sys.stdout.flush()
sys.stderr.flush()
new_args = shlex.split(input(prompt))
# Special shell commands
if new_args == ["exit"]:
return
# Run the specified command
try:
args = self.parser.parse_args(new_args)
command = root_command.resolve_command_from_args(args)
if not isinstance(command, ClientCommand):
print("shell commands must be client commands",
file=sys.stderr)
continue
await command.run_with_client(args, client)
print("SUCCESS=1")
except IdbException as e:
print(e.args[0], file=sys.stderr)
print("SUCCESS=0")
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_domain_sock_server(
self, config: CompanionServerConfig,
path: str) -> asyncio.subprocess.Process:
(process, log_file_path) = await self._spawn_server(
config=config,
port_env_variables={},
bind_arguments=["--grpc-domain-sock", path],
)
stdout = none_throws(process.stdout)
try:
extracted_path = await _extract_domain_sock_from_spawned_companion(
stream=stdout, log_file_path=log_file_path)
except Exception as e:
raise CompanionSpawnerException(
f"Failed to spawn companion, couldn't read domain socket path "
f"stderr: {get_last_n_lines(log_file_path, 30)}") from e
if not extracted_path:
raise CompanionSpawnerException(
f"Failed to spawn companion, no extracted domain socket path "
f"stderr: {get_last_n_lines(log_file_path, 30)}")
if extracted_path != path:
raise CompanionSpawnerException(
"Failed to spawn companion, extracted domain socket path "
f"is not correct (expected {path} got {extracted_path})"
f"stderr: {get_last_n_lines(log_file_path, 30)}")
return process
async def daemon(client: CompanionClient,
stream: Stream[RecordResponse, RecordRequest]) -> None:
client.logger.info(f"Starting connection to backend")
request = await stream.recv_message()
output_file = none_throws(request).start.file_path
async with client.stub.record.open() as forward_stream:
if client.is_local:
client.logger.info(
f"Starting video recording to local file {output_file}")
await forward_stream.send_message(
RecordRequest(start=Start(file_path=output_file)))
else:
client.logger.info(f"Starting video recording with response data")
await forward_stream.send_message(
RecordRequest(start=Start(file_path=None)))
client.logger.info("Request sent")
await stream.recv_message()
client.logger.info("Stopping video recording")
await forward_stream.send_message(RecordRequest(stop=Stop()))
await forward_stream.stream.end()
if client.is_local:
client.logger.info("Responding with file path")
response = await forward_stream.recv_message()
await stream.send_message(response)
else:
client.logger.info(f"Decompressing gzip to {output_file}")
await drain_gzip_decompress(_generate_bytes(forward_stream),
output_path=output_file)
client.logger.info(f"Finished decompression to {output_file}")
await stream.send_message(
RecordResponse(payload=Payload(file_path=output_file)))
async def get_stub_for_udid(self, udid: Optional[str]) -> CompanionClient:
is_companion_available = (
self.is_companion_available_for_target_udid(udid)
if udid else False)
if udid and udid in self._stub_map and udid in self._udid_companion_map:
return CompanionClient(
stub=self._stub_map[udid],
is_local=self._udid_companion_map[udid].is_local,
udid=udid,
logger=self._logger,
is_companion_available=is_companion_available,
)
else:
async with self.create_companion_for_target_with_udid(
target_udid=udid) as companion:
stub = self.get_stub_for_address(companion.host,
none_throws(companion.port))
self._stub_map[companion.udid] = stub
return CompanionClient(
stub=stub,
is_local=companion.is_local,
udid=udid,
logger=self._logger,
is_companion_available=is_companion_available,
)
async def spawn_tcp_server(
self,
config: CompanionServerConfig,
port: Optional[int],
tls_cert_path: Optional[str] = None,
) -> Tuple[asyncio.subprocess.Process, int]:
bind_arguments = [
"--grpc-port",
str(port) if port is not None else "0"
]
if tls_cert_path is not None:
bind_arguments.extend(["--tls-cert-path", tls_cert_path])
(process, log_file_path) = await self._spawn_server(
config=config,
bind_arguments=bind_arguments,
)
stdout = none_throws(process.stdout)
try:
extracted_port = await _extract_port_from_spawned_companion(stdout)
except Exception as e:
raise CompanionSpawnerException(
f"Failed to spawn companion, couldn't read port output "
f"stderr: {get_last_n_lines(log_file_path, 30)}") from e
if extracted_port == 0:
raise CompanionSpawnerException(
f"Failed to spawn companion, port zero is invalid "
f"stderr: {get_last_n_lines(log_file_path, 30)}")
if port is not None and extracted_port != port:
raise CompanionSpawnerException(
"Failed to spawn companion, invalid port "
f"(expected {port} got {extracted_port})"
f"stderr: {get_last_n_lines(log_file_path, 30)}")
return (process, extracted_port)
async def generate_gzip(path: str) -> AsyncIterator[bytes]:
async with _create_gzip_compress_command(path=path) as process:
reader = none_throws(process.stdout)
while not reader.at_eof():
data = await reader.read(READ_CHUNK_SIZE)
if not data:
return
yield data
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 drain_gzip_decompress(stream: AsyncIterator[bytes], output_path: str) -> None:
async with _create_gzip_decompress_command(extract_path=output_path) as process:
writer = none_throws(process.stdin)
async for data in stream:
writer.write(data)
await writer.drain()
writer.write_eof()
await writer.drain()
async def do_spawn_companion(
path: str,
udid: str,
log_file_path: str,
device_set_path: Optional[str],
port: Optional[int],
cwd: Optional[str],
tmp_path: Optional[str],
reparent: bool,
tls_cert_path: Optional[str] = None,
) -> Tuple[asyncio.subprocess.Process, int]:
arguments: List[str] = [
path,
"--udid",
udid,
"--grpc-port",
str(port) if port is not None else "0",
]
if tls_cert_path is not None:
arguments.extend(["--tls-cert-path", tls_cert_path])
if device_set_path is not None:
arguments.extend(["--device-set-path", device_set_path])
env = dict(os.environ)
if tmp_path:
env["TMPDIR"] = tmp_path
with open(log_file_path, "a") as log_file:
process = await asyncio.create_subprocess_exec(
*arguments,
stdout=asyncio.subprocess.PIPE,
stdin=asyncio.subprocess.PIPE if reparent else None,
stderr=log_file,
cwd=cwd,
env=env,
preexec_fn=os.setpgrp if reparent else None,
)
logging.debug(f"started companion at process id {process.pid}")
stdout = none_throws(process.stdout)
try:
extracted_port = await _extract_port_from_spawned_companion(stdout)
except Exception as e:
raise CompanionSpawnerException(
f"Failed to spawn companion, couldn't read port "
f"stderr: {get_last_n_lines(log_file_path, 30)}"
) from e
if extracted_port == 0:
raise CompanionSpawnerException(
f"Failed to spawn companion, port is zero"
f"stderr: {get_last_n_lines(log_file_path, 30)}"
)
if port is not None and extracted_port != port:
raise CompanionSpawnerException(
"Failed to spawn companion, port is not correct "
f"(expected {port} got {extracted_port})"
f"stderr: {get_last_n_lines(log_file_path, 30)}"
)
return (process, extracted_port)
async def boot(self) -> None:
if self.target_udid:
cmd: List[str] = ["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 run_with_client(self, args: Namespace, client: Client) -> None:
artifact: Optional[InstalledArtifact] = None
async for info in client.install(args.bundle_path):
artifact = info
progress = info.progress
if progress is None:
continue
self.logger.info(f"Progress: {progress}")
if args.json:
print(
json.dumps({
"installedAppBundleId": none_throws(artifact).name,
"uuid": none_throws(artifact).uuid,
}))
else:
print(
f"Installed: {none_throws(artifact).name} {none_throws(artifact).uuid}"
)
async def tail_targets(
self,
only: Optional[OnlyFilter] = None
) -> AsyncGenerator[List[TargetDescription], None]:
arguments = ["--notify", "stdout"] + _only_arg_from_filter(only=only)
async with self._start_companion_command(
arguments=arguments) as process:
async for line in none_throws(process.stdout):
yield target_descriptions_from_json(data=line.decode().strip())
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}")
async def start(self) -> None:
logging.debug(f"Started tailing notifier")
if self.process:
logging.warning(
f"Trying to start companion tailer when already running")
return
self.process = await self.notifierProcess()
if self.process:
self._reading_forever_fut = asyncio.ensure_future(
self._read_stream(stream=none_throws(self.process.stdout)))
async def _spawn_daemon(self) -> None:
cmd: List[str] = [sys.argv[0], "daemon"]
if platform.system() == "Darwin":
logging.debug("Mac Detected. passing notifier path to daemon")
cmd.extend(["--notifier-path", "idb_companion"])
with open(self._log_file_path(), "w") as log_file:
process = await asyncio.create_subprocess_exec(
*cmd, stdout=asyncio.subprocess.PIPE, stderr=log_file)
logging.debug(f"daemon process id {process.pid}")
save_pid(pid=process.pid)
await self._read_daemon_output(none_throws(process.stdout))
async def drain_to_stream(stream: Stream[_TSend, _TRecv],
generator: AsyncIterator[_TSend],
logger: Logger) -> _TRecv:
while True:
async for message in generator:
await stream.send_message(message)
await stream.end()
logger.debug(
f"Streamed all chunks to companion, waiting for completion")
response = none_throws(await stream.recv_message())
logger.debug(f"Companion completed")
return response
async def instruments_drain_until_running(stream: Stream[
InstrumentsRunRequest, InstrumentsRunResponse], logger: Logger) -> None:
while True:
response = none_throws(await stream.recv_message())
log_output = response.log_output
if len(log_output):
logger.info(log_output.decode().strip())
continue
state = response.state
if state == InstrumentsRunResponse.RUNNING_INSTRUMENTS:
logger.info("State changed to RUNNING_INSTRUMENTS")
return
async def daemon(context: DaemonContext, client: CompanionClient,
request: BootRequest) -> BootResponse:
if client.is_companion_available:
await none_throws(context.boot_manager
).boot(udid=none_throws(client.udid))
return BootResponse()
else:
# TODO T41660845
raise GRPCError(
status=Status(Status.UNIMPLEMENTED),
message="boot with chained daemons hasn't been implemented yet",
)
async def drain_untar(generator: AsyncIterator[bytes], output_path: str) -> None:
try:
os.mkdir(output_path)
except FileExistsError:
pass
async with _create_untar_command(output_path=output_path) as process:
writer = none_throws(process.stdin)
async for data in generator:
writer.write(data)
await writer.drain()
writer.write_eof()
await writer.drain()
await process.wait()
async def spawn_notifier(self) -> None:
if self._is_notifier_running():
return
self.check_okay_to_spawn()
cmd = [self.companion_path, "--notify", IDB_LOCAL_TARGETS_FILE]
with open(self._log_file_path("notifier"), "a") as log_file:
process = await asyncio.create_subprocess_exec(
*cmd, stdout=asyncio.subprocess.PIPE, stderr=log_file)
self.pid_saver.save_notifier_pid(pid=process.pid)
await asyncio.ensure_future(
self._read_notifier_output(stream=none_throws(process.stdout)))
logging.debug(f"started notifier at process id {process.pid}")
async def __aexit__(
self,
exc_type: Optional[Type[Exception]],
exception: Optional[Exception],
traceback: Optional[TracebackType],
) -> bool:
name = none_throws(self.name)
duration = int((time.time() - none_throws(self.start)) * 1000)
if exception:
logger.debug(f"{name} failed")
await plugin.failed_invocation(
name=name,
duration=duration,
exception=exception,
metadata=self.metadata,
)
else:
logger.debug(f"{name} succeeded")
await plugin.after_invocation(
name=name, duration=duration, metadata=self.metadata
)
return False
async def instruments_drain_until_stop(
stream: Stream[InstrumentsRunRequest, InstrumentsRunResponse],
stop: asyncio.Future,
logger: Logger,
) -> None:
while True:
read = asyncio.ensure_future(stream.recv_message())
done, pending = await asyncio.wait([stop, read],
return_when=asyncio.FIRST_COMPLETED)
if stop in done:
return
response = none_throws(read.result())
output = response.log_output
if len(output):
logger.info(output.decode())
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)
