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.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)
])
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)
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
async def kill(self) -> None:
await self.direct_companion_manager.clear()
self.local_targets_manager.clear()
PidSaver(logger=self.logger).kill_saved_pids()