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_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.common.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) ])
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) def _log_file_path(self, target_udid: str) -> str: os.makedirs(name=IDB_LOGS_PATH, exist_ok=True) return IDB_LOGS_PATH + "/" + target_udid def check_okay_to_spawn(self) -> None: if os.getuid() == 0: logging.warning( "idb should not be run as root. " "Listing available targets on this host and spawning " "companions will not work" ) async def spawn_companion(self, target_udid: str) -> int: self.check_okay_to_spawn() (process, port) = await do_spawn_companion( path=self.companion_path, udid=target_udid, log_file_path=self._log_file_path(target_udid), device_set_path=None, port=None, cwd=None, tmp_path=None, reparent=True, ) self.pid_saver.save_companion_pid(pid=process.pid) return port def _is_notifier_running(self) -> bool: pid = self.pid_saver.get_notifier_pid() # Taken from https://fburl.com/ibk820b6 if pid <= 0: return False try: # no-op if process exists os.kill(pid, 0) return True except OSError as err: # EPERM clearly means there's a process to deny access to # otherwise proc doesn't exist return err.errno == errno.EPERM except Exception: return False async def spawn_notifier(self, targets_file: str = IDB_LOCAL_TARGETS_FILE) -> None: if self._is_notifier_running(): return self.check_okay_to_spawn() cmd = [self.companion_path, "--notify", targets_file] log_path = self._log_file_path("notifier") with open(log_path, "a") as log_file: process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=log_file ) try: self.pid_saver.save_notifier_pid(pid=process.pid) await self._read_notifier_output(stream=none_throws(process.stdout)) logging.debug(f"started notifier at process id {process.pid}") except Exception as e: raise CompanionSpawnerException( "Failed to spawn the idb notifier. " f"Stderr: {get_last_n_lines(log_path, 30)}" ) from e async def _read_notifier_output(self, stream: asyncio.StreamReader) -> None: while True: line = await stream.readline() if line is None: return update = parse_json_line(line) if update["report_initial_state"]: return