예제 #1
0
 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)
예제 #2
0
 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)
예제 #3
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)
             ])
예제 #4
0
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