From 133713379f958224b1ff3afe9e7a2658c2aff6cb Mon Sep 17 00:00:00 2001 From: Billy Lau Date: Mon, 28 Sep 2026 07:19:45 -0500 Subject: [PATCH] automate_observation: add --serial to target specific devices Add a repeatable -s/--serial flag so that callers (and the upcoming web UI helper) can choose which connected devices to observe instead of always iterating over every device listed by `adb devices`. - Selected devices are processed in the order given; duplicates are ignored. - A requested serial that is not connected is logged, reported as FAILED in the final summary with a dedicated message, and causes exit code 1. Remaining requested devices are still observed. - The "More than 1 device connected!" warning is only emitted when --serial is not used. - Without --serial, behaviour, log text and exit codes are unchanged. - This also applies when no device is connected at all (or `adb devices` fails): with --serial, the "No devices connected!" early return is skipped so the requested serials are reported as FAILED. Selection logic lives in a new pure helper, select_target_devices(), covered by unit tests alongside end-to-end main() tests. Change-Id: I9f3fc21eb1b85b9fda4fa0a5c5b4da95322f3091 --- uraniborg/docs/automate_observation.md | 15 ++ .../scripts/python/automate_observation.py | 80 ++++++- .../python/tests/test_automate_observation.py | 219 ++++++++++++++++++ 3 files changed, 305 insertions(+), 9 deletions(-) diff --git a/uraniborg/docs/automate_observation.md b/uraniborg/docs/automate_observation.md index 3d3dd42..2c62914 100644 --- a/uraniborg/docs/automate_observation.md +++ b/uraniborg/docs/automate_observation.md @@ -27,6 +27,21 @@ If you see the final SUCCESS message, feel free to ignore earlier ERROR messages Those resulted from some files that cannot be extracted from device, but does not affect the core files required for further analysis. +### Selecting Devices + +By default, the script observes **every** device listed by `adb devices`, one +after another. To observe only specific devices, pass their serial numbers with +`-s`/`--serial`. The flag may be repeated, and devices are processed in the +order given: + +```bash +python3 automate_observation.py --serial ABCDEF012345 --serial 9876543210FEDCBA +``` + +A requested serial that is not connected is not silently skipped: it is logged +as an error, reported as `FAILED` in the final summary, and makes the script +exit with code `1`. The remaining requested devices are still observed. + ## Performing Inclusion Proof Checks You can automatically verify extracted package APK splits against Android Binary diff --git a/uraniborg/scripts/python/automate_observation.py b/uraniborg/scripts/python/automate_observation.py index 78a338c..0de14f9 100644 --- a/uraniborg/scripts/python/automate_observation.py +++ b/uraniborg/scripts/python/automate_observation.py @@ -68,6 +68,14 @@ def parse_arguments() -> argparse.Namespace: "\"results\" directory can be found.") parser.add_argument("-D", "--debug", required=False, action="count", help="If specified, debugging mode is turned on.") + parser.add_argument("-s", "--serial", required=False, action="append", + default=None, metavar="SERIAL", + help="Serial number of a connected device to observe " + "(as listed by `adb devices`). May be repeated to " + "observe several devices, which are processed in " + "the order given. If omitted, every connected " + "device is observed. A requested serial that is " + "not connected is reported as FAILED.") parser.add_argument("--pull-all-apks", required=False, action="count", help="If specified, the script will attempt to download " @@ -196,6 +204,46 @@ def adb_installed(logger: logging.Logger) -> bool: return False +def select_target_devices(connected_devices: list, + requested_serials: Optional[list[str]], + logger: logging.Logger) -> tuple[list, list[str]]: + """Narrows connected devices down to the ones requested via --serial. + + Args: + connected_devices: DeviceInfo-like objects as returned by + AdbWrapper.devices(). + requested_serials: Serial numbers passed via --serial, or None if the flag + was not used. Duplicates are ignored. + logger: A logger object to log debug or error messages. + + Returns: + A tuple (target_devices, missing_serials). When requested_serials is None, + target_devices is connected_devices unchanged and missing_serials is empty. + Otherwise target_devices holds the connected devices matching the request, + in the order the serials were requested, and missing_serials holds the + requested serials that are not connected, also in request order. + """ + if requested_serials is None: + return list(connected_devices), [] + + by_serial = {d.serial_number: d for d in connected_devices} + target_devices = [] + missing_serials = [] + seen = set() + for serial in requested_serials: + if serial in seen: + logger.debug("Ignoring duplicate --serial %s", serial) + continue + seen.add(serial) + if serial in by_serial: + target_devices.append(by_serial[serial]) + else: + logger.error("Requested device with serial number %s is not connected.", + serial) + missing_serials.append(serial) + return target_devices, missing_serials + + def clear_logcat(adb_wrapper): adb_wrapper.logcat_clear() @@ -900,23 +948,27 @@ def main(): if not AdbWrapper.start_server(logger): return - connected_devices = AdbWrapper.devices(logger) - if not connected_devices or len(connected_devices) < 1: + connected_devices = AdbWrapper.devices(logger) or [] + # With --serial, fall through: every requested serial is then reported as + # missing, FAILED in the summary, and the run exits 1. + if not connected_devices and args.serial is None: logger.error("No devices connected!") return logger.debug("There are %d connected device(s)", len(connected_devices)) - # Here we'll handle the cases where there can be 1 or more devices connected - if len(connected_devices) > 1: - # TODO(billylau): Branch off based on user-input - do all, or selectively + target_devices, missing_serials = select_target_devices( + connected_devices, args.serial, logger) + # Without --serial every connected device is observed; warn in case that + # was not intended. With --serial the user has already chosen explicitly. + if args.serial is None and len(connected_devices) > 1: logger.warning("More than 1 device connected!") results = {} prefetched = False - has_errors = False + has_errors = bool(missing_serials) verification_failed_devices = set() collection_error_devices = set() - for target_device in connected_devices: + for target_device in target_devices: try: if target_device.unauthorized: logger.error("Please authorize device with serial number %s for ADB via " @@ -1023,8 +1075,18 @@ def main(): if target_device.serial_number in results: collection_error_devices.add(target_device.serial_number) - for target_device in connected_devices: - device = target_device.serial_number + # Summarise in the order devices were requested (or discovered, without + # --serial), including requested serials that were never connected. + missing = set(missing_serials) + summary_serials = list(dict.fromkeys( + args.serial if args.serial is not None + else [d.serial_number for d in connected_devices])) + for device in summary_serials: + if device in missing: + logger.error( + "FAILED: Requested device %s is not connected (exiting 1).", + device) + continue if device not in results: logger.error( "FAILED: Hubble data collection failed on connected device %s " diff --git a/uraniborg/scripts/python/tests/test_automate_observation.py b/uraniborg/scripts/python/tests/test_automate_observation.py index 5d8f6a8..3e3541a 100644 --- a/uraniborg/scripts/python/tests/test_automate_observation.py +++ b/uraniborg/scripts/python/tests/test_automate_observation.py @@ -17,6 +17,7 @@ """Unit tests for automate_observation.py CLI parsing and multi-device prefetch latch.""" import argparse +import contextlib import io import json import os @@ -1654,5 +1655,223 @@ def make_adb_wrapper(serial: str, _logger): ) +# --- --serial ----------------------------------------------------------------- + + +_MULTI_DEVICE_WARNING = "More than 1 device connected!" + + +@pytest.fixture +def serial_main_mocks(): + """Patches main()'s collaborators so every device collects successfully. + + extract_results_and_apks returns /tmp/out/ for whichever device is + being processed, so tests can assert exactly which devices ran. + """ + patch_targets = [ + "supported_platform", "verify_hubble", "adb_installed", "AdbWrapper", + "is_hubble_installed", "is_xiaomi_phone", "install_hubble", + "clear_logcat", "launch_hubble", "wait_for_results", + "extract_results_and_apks", "extract_selinux_policies", + "set_up_logging", + ] + with contextlib.ExitStack() as stack: + mocks = { + name: stack.enter_context(mock.patch(f"automate_observation.{name}")) + for name in patch_targets + } + mocks["supported_platform"].return_value = True + mocks["verify_hubble"].return_value = True + mocks["adb_installed"].return_value = True + mocks["AdbWrapper"].start_server.return_value = True + mocks["is_hubble_installed"].return_value = False + mocks["is_xiaomi_phone"].return_value = False + mocks["install_hubble"].return_value = True + mocks["launch_hubble"].return_value = True + mocks["wait_for_results"].return_value = "/sdcard/hubble/results" + + processed = [] + + def _adb_wrapper_ctor(serial, logger): + processed.append(serial) + return mock.Mock() + + mocks["AdbWrapper"].side_effect = _adb_wrapper_ctor + mocks["extract_results_and_apks"].side_effect = ( + lambda *a, **kw: f"/tmp/out/{processed[-1]}") + mocks["logger"] = mock.Mock() + mocks["set_up_logging"].return_value = mocks["logger"] + mocks["processed"] = processed + yield mocks + + +def _set_argv(monkeypatch: pytest.MonkeyPatch, *extra: str): + monkeypatch.setattr( + sys, "argv", + ["automate_observation.py", "-H", "/path/to/hubble.apk", + "-o", "/tmp/out", *extra]) + + +def test_parse_arguments_serial_default_and_repeatable( + monkeypatch: pytest.MonkeyPatch, +): + """--serial defaults to None and accumulates in order across -s/--serial.""" + monkeypatch.setattr(sys, "argv", ["automate_observation.py"]) + assert automate_observation.parse_arguments().serial is None + + monkeypatch.setattr( + sys, "argv", + ["automate_observation.py", "-s", "B", "--serial", "A", "--serial=C"]) + assert automate_observation.parse_arguments().serial == ["B", "A", "C"] + + +def test_select_target_devices_without_serial_returns_all(): + logger = mock.Mock() + devices = [_make_mock_device("A"), _make_mock_device("B")] + targets, missing = automate_observation.select_target_devices( + devices, None, logger) + assert targets == devices + assert targets is not devices # a copy, so callers cannot mutate the input + assert missing == [] + logger.error.assert_not_called() + + +def test_select_target_devices_filters_orders_dedupes_and_reports_missing(): + logger = mock.Mock() + a, b, c = (_make_mock_device(s) for s in ("A", "B", "C")) + targets, missing = automate_observation.select_target_devices( + [a, b, c], ["C", "GONE1", "A", "C", "GONE2", "GONE1"], logger) + assert targets == [c, a] + assert missing == ["GONE1", "GONE2"] + assert logger.error.call_count == 2 + logger.error.assert_any_call( + "Requested device with serial number %s is not connected.", "GONE1") + logger.error.assert_any_call( + "Requested device with serial number %s is not connected.", "GONE2") + + +def test_select_target_devices_empty_request_selects_nothing(): + targets, missing = automate_observation.select_target_devices( + [_make_mock_device("A")], [], mock.Mock()) + assert targets == [] + assert missing == [] + + +def test_main_serial_limits_and_orders_devices( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + """Only requested devices run, in request order, without the multi-device warning.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1"), + _make_mock_device("DEV2"), + _make_mock_device("DEV3"), + ] + _set_argv(monkeypatch, "--serial", "DEV3", "-s", "DEV1") + + automate_observation.main() # no SystemExit: every selected device succeeded + + assert m["processed"] == ["DEV3", "DEV1"] + assert m["install_hubble"].call_count == 2 + for call in m["logger"].warning.call_args_list: + assert call.args != (_MULTI_DEVICE_WARNING,) + + success_msg = ("SUCCESS! Hubble was successfully deployed and executed on " + "connected device %s.") + success_calls = [c.args for c in m["logger"].info.call_args_list + if c.args and c.args[0] == success_msg] + assert success_calls == [(success_msg, "DEV3"), (success_msg, "DEV1")] + + +def test_main_serial_not_connected_is_reported_failed_and_exits_1( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + """A requested serial that is not connected fails loudly; others still run.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1"), + _make_mock_device("DEV2"), + ] + _set_argv(monkeypatch, "--serial", "GONE", "--serial", "DEV1") + + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + assert exc_info.value.code == 1 + + assert m["processed"] == ["DEV1"] + m["logger"].error.assert_any_call( + "Requested device with serial number %s is not connected.", "GONE") + m["logger"].error.assert_any_call( + "FAILED: Requested device %s is not connected (exiting 1).", "GONE") + # The generic collection-failure line is reserved for connected devices. + for call in m["logger"].error.call_args_list: + assert call.args != ( + "FAILED: Hubble data collection failed on connected device %s " + "(exiting 1).", + "GONE", + ) + m["logger"].info.assert_any_call( + "SUCCESS! Hubble was successfully deployed and executed on " + "connected device %s.", + "DEV1", + ) + # DEV2 was not requested, so it must not appear in the summary at all. + for method in (m["logger"].info, m["logger"].warning, m["logger"].error): + for call in method.call_args_list: + assert "DEV2" not in call.args + + +def test_main_without_serial_still_runs_all_and_warns( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + """Default behaviour is unchanged: all devices run and the warning is kept.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1"), + _make_mock_device("DEV2"), + ] + _set_argv(monkeypatch) + + automate_observation.main() + + assert m["processed"] == ["DEV1", "DEV2"] + m["logger"].warning.assert_any_call(_MULTI_DEVICE_WARNING) + + +@pytest.mark.parametrize("listing", [[], None], ids=["empty", "adb_failed"]) +def test_main_serial_with_no_connected_devices_fails_and_exits_1( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, listing, +): + """--serial must not be swallowed by the "No devices connected!" early return.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = listing + _set_argv(monkeypatch, "--serial", "GONE") + + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + assert exc_info.value.code == 1 + + assert m["processed"] == [] + m["logger"].error.assert_any_call( + "FAILED: Requested device %s is not connected (exiting 1).", "GONE") + for call in m["logger"].error.call_args_list: + assert call.args != ("No devices connected!",) + + +@pytest.mark.parametrize("listing", [[], None], ids=["empty", "adb_failed"]) +def test_main_without_serial_no_devices_keeps_early_return( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, listing, +): + """Without --serial, the pre-existing early return (exit 0) is unchanged.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = listing + _set_argv(monkeypatch) + + automate_observation.main() # no SystemExit + + assert m["processed"] == [] + m["logger"].error.assert_any_call("No devices connected!") + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))