async def test_list_apps_direct_companion_uxd(self) -> None:
self.direct_client_mock.list_apps = AsyncMock(return_value=[])
await cli_main(cmd_input=["--companion", "/foo/sock", "list-apps"])
self.direct_client_mock.list_apps.assert_called_once()
self.direct_client_mock.build.assert_called_once_with(
address=DomainSocketAddress(path="/foo/sock"), is_local=False, logger=ANY
)
def json_to_companion_info(data: List[Dict[str, Any]]) -> List[CompanionInfo]:
return [
CompanionInfo(
udid=item["udid"],
address=(TCPAddress(host=item["host"], port=item["port"]) if "host"
in item else DomainSocketAddress(path=item["path"])),
is_local=item["is_local"],
pid=item.get("pid"),
) for item in data
]
async def for_companion(
cls, companion: Companion, udid: str, logger: logging.Logger
) -> AsyncGenerator["IdbClient", None]:
with tempfile.NamedTemporaryFile() as temp:
# Remove the tempfile so we can bind to it first.
os.remove(temp.name)
async with companion.unix_domain_server(
udid=udid, path=temp.name
) as resolved_path, IdbClient.build(
address=DomainSocketAddress(path=resolved_path),
is_local=True,
logger=logger,
) as client:
yield client
async def test_add_then_remove_companion_by_uxd_address(self) -> None:
async for manager in self._managers():
companion = CompanionInfo(
udid="asdasda",
address=DomainSocketAddress(path="/tmp/foo.sock"),
is_local=False,
)
replaced = await manager.add_companion(companion)
self.assertIsNone(replaced)
companions = await manager.get_companions()
self.assertEqual(companions, [companion])
removed = await manager.remove_companion(companion.address)
companions = await manager.get_companions()
self.assertEqual(companions, [])
self.assertEqual(removed, [companion])
async def test_replace_companion(self) -> None:
async for manager in self._managers():
companion_first = CompanionInfo(
udid="a",
address=TCPAddress(host="ahost", port=123),
is_local=False,
pid=None,
)
replaced = await manager.add_companion(companion_first)
self.assertIsNone(replaced)
companions = await manager.get_companions()
self.assertEqual(companions, [companion_first])
companion_second = CompanionInfo(
udid="a",
address=DomainSocketAddress(path="/some/path"),
is_local=False,
pid=123,
)
replaced = await manager.add_companion(companion_second)
self.assertEqual(replaced, companion_first)
companions = await manager.get_companions()
self.assertEqual(companions, [companion_second])
async def test_add_then_clear(self) -> None:
async for manager in self._managers():
first = CompanionInfo(
udid="asdasda",
address=TCPAddress(host="foohost", port=123),
is_local=False,
pid=None,
)
second = CompanionInfo(
udid="fooo",
address=DomainSocketAddress(path="/bar/bar"),
is_local=False,
pid=324,
)
await manager.add_companion(first)
await manager.add_companion(second)
companions = await manager.get_companions()
self.assertEqual(companions, [first, second])
cleared = await manager.clear()
companions = await manager.get_companions()
self.assertEqual(companions, [])
self.assertEqual(cleared, [first, second])
async def _spawn_companion_server(self, udid: str) -> CompanionInfo:
companion = self._companion
if companion is None:
raise IdbException(
f"Cannot spawn companion for {udid}, no companion executable")
target_type = await _local_target_type(companion=companion, udid=udid)
path = os.path.join(BASE_IDB_FILE_PATH, f"{udid}_companion.sock")
address = DomainSocketAddress(path=path)
self._logger.info(
f"Checking whether domain sock {path} is bound for {udid}")
is_bound = await _check_domain_socket_is_bound(path=path)
if is_bound:
self._logger.info(
f"Domain socket {path} is bound for {udid}, connecting to it.")
companion_info = await self.connect(destination=address)
else:
self._logger.info(
f"No existing companion at {path}, spawning one...")
process = await companion.spawn_domain_sock_server(
config=CompanionServerConfig(
udid=udid,
only=target_type,
log_file_path=None,
cwd=None,
tmp_path=None,
reparent=True,
),
path=path,
)
self._logger.info(f"Companion at {path} spawned for {udid}")
companion_info = CompanionInfo(
address=address,
udid=udid,
is_local=True,
pid=process.pid,
)
await self._companion_set.add_companion(companion_info)
return companion_info
def _parse_address(value: str) -> Address:
values = value.rsplit(":", 1)
if len(values) == 1:
return DomainSocketAddress(path=value)
(host, port) = values
return TCPAddress(host=host, port=int(port))