automate_observation: add --serial to target specific devices

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

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 "

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/<serial> 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__]))