def test_save_notifier_pid(self) -> None:
with tempfile.NamedTemporaryFile() as f:
pid_saver = PidSaver(logger=mock.MagicMock(),
pids_file_path=f.name)
pid = 1
pid_saver.save_notifier_pid(pid)
data = json.load(f)
notifier_pid = data["notifier"]
self.assertEqual(notifier_pid, pid)
def test_save_companion_pid(self) -> None:
with tempfile.NamedTemporaryFile() as f:
pid_saver = PidSaver(logger=mock.MagicMock(),
pids_file_path=f.name)
pid = 1
pid_saver.save_companion_pid(pid)
data = json.load(f)
companion_pids = data["companions"]
self.assertEqual(companion_pids[0], pid)
def test_get_saved_pids(self) -> None:
with tempfile.NamedTemporaryFile() as f:
pid_saver = PidSaver(logger=mock.MagicMock(),
pids_file_path=f.name)
companion_pids = [1, 2]
notifier_pid = 3
with open(f.name, "w") as f:
json.dump(({
"companions": companion_pids,
"notifier": notifier_pid
}), f)
pid_saver._load()
self.assertEqual(pid_saver.companion_pids, companion_pids)
self.assertEqual(pid_saver.notifier_pid, notifier_pid)
def test_clear(self) -> None:
with tempfile.NamedTemporaryFile() as f:
pid_saver = PidSaver(logger=mock.MagicMock(),
pids_file_path=f.name)
pid = 1
pid_saver.save_companion_pid(pid)
pid_saver._clear_saved_pids()
pid_saver._load()
self.assertEqual(pid_saver.companion_pids, [])
self.assertEqual(pid_saver.notifier_pid, 0)
async def test_kill_saved_pids(self) -> None:
with tempfile.NamedTemporaryFile() as f:
pid_saver = PidSaver(logger=mock.MagicMock(),
pids_file_path=f.name)
companion_pid = 1
pid_saver.save_companion_pid(companion_pid)
notifier_pid = 2
pid_saver.save_notifier_pid(notifier_pid)
with mock.patch("idb.client.pid_saver.os.kill") as kill:
pid_saver.kill_saved_pids()
kill.assert_has_calls([
mock.call(1, signal.SIGTERM),
mock.call(2, signal.SIGTERM)
])
async def kill(self) -> None:
self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()
def __init__(self, companion_path: str, logger: logging.Logger) -> None:
self.companion_path = companion_path
self.logger = logger
self.pid_saver = PidSaver(logger=self.logger)
class CompanionSpawner:
def __init__(self, companion_path: str, logger: logging.Logger) -> None:
self.companion_path = companion_path
self.logger = logger
self.pid_saver = PidSaver(logger=self.logger)
async def _read_stream(self, stream: StreamReader) -> int:
port = 0
while True:
line = await stream.readline()
logging.debug(f"read line from companion : {line}")
if line:
update = json.loads(line.decode())
if update:
logging.debug(f"got update from companion {update}")
port = update["grpc_port"]
break
else:
break
return port
def _log_file_path(self, target_udid: str) -> str:
os.makedirs(name=IDB_LOGS_PATH, exist_ok=True)
return IDB_LOGS_PATH + "/" + target_udid
async def spawn_companion(self, target_udid: str) -> int:
if not self.companion_path:
raise CompanionSpawnerException(
f"couldn't instantiate a companion for {target_udid} because\
the companion_path is not available"
)
cmd: List[str] = [
self.companion_path,
"--udid",
target_udid,
"--grpc-port",
"0",
]
with open(self._log_file_path(target_udid), "a") as log_file:
process = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stdin=asyncio.subprocess.PIPE,
stderr=log_file,
)
self.pid_saver.save_companion_pid(pid=process.pid)
logging.debug(f"started companion at process id {process.pid}")
if process.stdout:
# pyre-fixme[6]: Expected `StreamReader` for 1st param but got
# `Optional[StreamReader]`.
port = await self._read_stream(process.stdout)
if not port:
raise CompanionSpawnerException("failed to spawn companion")
return port
raise CompanionSpawnerException("process has no stdout")
def _is_notifier_running(self) -> bool:
return self.pid_saver.get_notifier_pid() > 0
async def spawn_notifier(self) -> None:
if not self._is_notifier_running():
if not self.companion_path:
raise CompanionSpawnerException(
f"couldn't instantiate a notifier because\
the companion_path is not available"
)
cmd: List[str] = [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 _read_notifier_output(self, stream: StreamReader) -> None:
while True:
line = await stream.readline()
if line:
update = json.loads(line.decode())
if update["report_initial_state"]:
return
else:
return
