def get_destination(args: Namespace) -> Union[Address, str]:
if is_udid(args.companion):
return args.companion
elif args.port and args.grpc_port and args.companion:
return Address(host=args.companion,
port=args.port,
grpc_port=args.grpc_port)
elif args.port and args.companion:
return Address(host=args.companion, grpc_port=args.port)
else:
raise ConnectCommandException(
"provide either a UDID or the host and port of the companion")
def get_destination(args: Namespace) -> Union[Address, str]:
target_udid = args.companion if "-" in args.companion else None
companion_host = args.companion if not target_udid else None
if target_udid:
return target_udid
elif args.port and args.grpc_port and companion_host:
return Address(host=companion_host, port=args.port, grpc_port=args.grpc_port)
elif args.port and companion_host:
return Address(host=companion_host, grpc_port=args.port)
else:
raise ConnectCommandException(
"provide either a UDID or the host and port of the companion"
)
async def test_add_multiple(self) -> None:
async for manager in self._managers():
companion_a = CompanionInfo(udid="a",
host="ahost",
port=123,
is_local=False)
replaced = await manager.add_companion(companion_a)
self.assertIsNone(replaced)
companions = await manager.get_companions()
self.assertEqual(companions, [companion_a])
companion_b = CompanionInfo(udid="b",
host="bhost",
port=123,
is_local=False)
replaced = await manager.add_companion(companion_b)
self.assertIsNone(replaced)
companions = await manager.get_companions()
self.assertEqual(companions, [companion_a, companion_b])
companion_c = CompanionInfo(udid="c",
host="chost",
port=123,
is_local=False)
replaced = await manager.add_companion(companion_c)
self.assertIsNone(replaced)
companions = await manager.get_companions()
self.assertEqual(companions,
[companion_a, companion_b, companion_c])
removed = await manager.remove_companion(
Address(host=companion_b.host, port=companion_b.port))
companions = await manager.get_companions()
self.assertEqual(companions, [companion_a, companion_c])
self.assertEqual(removed, [companion_b])
removed = await manager.remove_companion("a")
companions = await manager.get_companions()
self.assertEqual(companions, [companion_c])
async def create_companion_for_target_with_udid(
self,
target_udid: Optional[str],
metadata: Optional[Dict[str, str]] = None,
timeout: Optional[float] = None,
) -> AsyncContextManager[CompanionInfo]:
self._logger.debug(f"getting companion for {target_udid}")
companion = self._get_companion_for_target(target_udid=target_udid)
if companion:
yield companion
elif target_udid is None:
raise Exception("Please provide a UDID for your target")
elif self.companion_spawner and target_udid in self._udid_target_map:
self._logger.debug(f"spawning a companion for {target_udid}")
port = await self.companion_spawner.spawn_companion(
target_udid=target_udid)
# overriding host here with localhost as spawning
# the companion only works locally
host = "localhost"
self._logger.info(f"companion started at {host}:{port}")
async with self.create_companion_for_target_with_address(
address=Address(host=host, port=port, grpc_port=port),
metadata=metadata,
timeout=timeout,
) as companion:
yield companion
else:
raise Exception(f"no companion available for {target_udid}")
def destination_to_py(
destination: GrpcConnectionDestination) -> ConnectionDestination:
if destination.HasField("address"):
return Address(host=destination.address.host,
port=destination.address.port)
else:
return destination.target_udid
async def test_disconnect_with_host_and_port(self) -> None:
self.management_client_mock().disconnect = AsyncMock()
host = "someHost"
port = 1234
await cli_main(cmd_input=["disconnect", host, str(port)])
self.management_client_mock().disconnect.assert_called_once_with(
destination=Address(host=host, port=port))
async def test_remove_companion_by_address(self) -> None:
companion_manager = CompanionManager(logger=mock.MagicMock())
add_companion(companion_manager, TEST_COMPANION)
self.assertEqual(len(companion_manager._udid_companion_map), 1)
self.assertEqual(len(companion_manager._udid_target_map), 1)
result = companion_manager.remove_companion(
Address(host=TEST_COMPANION.host, port=TEST_COMPANION.port))
self.assertEqual(result, TEST_COMPANION)
self.assertEqual(len(companion_manager._udid_companion_map), 0)
self.assertEqual(len(companion_manager._udid_target_map), 0)
async def build(
cls, host: str, port: int, is_local: bool, logger: logging.Logger
) -> AsyncContextManager["IdbClient"]:
channel = Channel(host=host, port=port, loop=asyncio.get_event_loop())
try:
yield IdbClient(
stub=CompanionServiceStub(channel=channel),
address=Address(host=host, port=port),
is_local=is_local,
logger=logger,
)
finally:
channel.close()
async def test_remove_companion_with_udid(self) -> None:
with tempfile.NamedTemporaryFile() as f:
companion_manager = DirectCompanionManager(logger=mock.MagicMock(),
state_file_path=f.name)
companion = CompanionInfo(udid="asdasda",
host="foohost",
port=123,
is_local=False)
with open(f.name, "w") as f:
json.dump(json_data_companions([companion]), f)
companion_manager.remove_companion(
Address(host=companion.host, port=companion.port))
companions = companion_manager._load()
self.assertEqual(companions, [])
async def test_add_then_remove_companion_by_address(self) -> None:
async for manager in self._managers():
companion = CompanionInfo(udid="asdasda",
host="foohost",
port=123,
is_local=False)
await manager.add_companion(companion)
companions = await manager.get_companions()
self.assertEqual(companions, [companion])
removed = await manager.remove_companion(
Address(host=companion.host, port=companion.port))
companions = await manager.get_companions()
self.assertEqual(companions, [])
self.assertEqual(removed, [companion])
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 _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:
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
