chromite: Fix style issues on_device_fuzz

This CL attempts to address most comments given to
https://crrev/c/3976657 after merge.

Since I can't edit the original, this is a fixup CL.

BUG=b:255365294
TEST=Ran shill_profile_fuzzer on hatch remote device.
TEST=CQ

Change-Id: I03b6752a75dcf0e155d9887d813e61b003735b7a
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/4083415
Commit-Queue: Jordan Abrahams-Whitehead <ajordanr@google.com>
Tested-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Reviewed-by: Jack Neus <jackneus@google.com>
diff --git a/scripts/cros_on_device_fuzz.py b/scripts/cros_on_device_fuzz.py
old mode 100755
new mode 100644
index 2555e32..00711c4
--- a/scripts/cros_on_device_fuzz.py
+++ b/scripts/cros_on_device_fuzz.py
@@ -1,5 +1,4 @@
-#!/usr/bin/env python3
-# Copyright 2022 The ChromiumOS Authors.
+# Copyright 2022 The ChromiumOS Authors
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
@@ -21,7 +20,7 @@
 from pathlib import Path
 import re
 import tempfile
-from typing import Iterable
+from typing import Iterable, Optional
 
 from chromite.lib import commandline
 from chromite.lib import cros_build_lib
@@ -29,31 +28,36 @@
 from chromite.lib import remote_access
 
 
-_DEFAULT_PRIVATE_KEY = Path.home() / ".ssh" / "testing_rsa"
 _DUT_FUZZER_ROOT = Path("/usr/local/tmp/fuzzer_root")
 _IMPLICIT_SYSTEM_PACKAGE = "virtual/implicit-system"
-_PACKAGE_SEPARATOR = re.compile("[ ,]+")
+_PACKAGE_SEPARATOR = re.compile(r"[ ,]+")
+_FUZZER_BOARD = "amd64-generic"
 
 
 def main(argv: Iterable[str]):
     """Dispatch to other functions."""
-    args = parse_args(argv)
-    args.func(args)
+    opts = parse_args(argv)
+    opts.func(opts)
 
 
-def setup_main(args):
+def setup_main(opts):
     """Setup up the DUT for fuzzing."""
     cros_build_lib.AssertInsideChroot()
-    device = _parse_device(args.device_host, _DEFAULT_PRIVATE_KEY)
+    device = _get_device(opts.device, opts.private_key)
+    # Get a nice string representation of the device host for logging
+    host = _format_collection(opts.device)
     # We do this version check to ensure we have a CrOS device
     # immediately.
     try:
         version = device.version
+        if not version:
+            raise RuntimeError("Version is null or empty.")
     except Exception:
-        logging.error("Unable to get version of remote %s", args.device_host)
+        logging.error("Unable to get version of remote %s", host)
         raise
-    logging.info("Connected to %s; CrOS version %s", args.device_host, version)
-    packages = set(re.split(_PACKAGE_SEPARATOR, args.packages))
+
+    logging.info("Connected to %s; CrOS version %s", host, version)
+    packages = set(re.split(_PACKAGE_SEPARATOR, opts.packages))
     # We require virtual/implicit-system so that we can
     # build the on-device sysroot. Without it, we'll
     # miss critical libs and binaries.
@@ -63,7 +67,7 @@
         tarball_name = "sysroot_fuzzer_tarball.tar.xz"
         on_device_fuzz.create_sysroot_tarball(
             packages=packages,
-            board="amd64-generic",
+            board=_FUZZER_BOARD,
             output_path=tmpdir / tarball_name,
         )
         on_device_fuzz.create_dut_sysroot(
@@ -71,59 +75,50 @@
             sysroot_tarball=tmpdir / tarball_name,
             sysroot_device_path=_DUT_FUZZER_ROOT,
         )
-    logging.info("Fuzzer set up complete for %s", args.device_host)
+    logging.info("Fuzzer set up complete for %s", host)
 
 
-def fuzz_main(args):
+def fuzz_main(opts):
     """Run a given fuzzer on the DUT."""
     cros_build_lib.AssertInsideChroot()
     fuzzer_install_path = Path("/usr/libexec/fuzzers")
-    device = _parse_device(args.device_host, _DEFAULT_PRIVATE_KEY)
-    libfuzzer_options = {}
+    device = _get_device(opts.device, opts.private_key)
     on_device_fuzz.run_fuzzer_executable(
         device,
         sysroot_device_path=_DUT_FUZZER_ROOT,
-        sysroot_fuzzer_path=fuzzer_install_path / args.fuzzer_name,
-        libfuzzer_options=libfuzzer_options,
+        sysroot_fuzzer_path=fuzzer_install_path / opts.fuzzer_name,
+        libfuzzer_options={},
     )
 
 
-def _parse_device(
-    device_host: str, private_key: Path
+def _get_device(
+    collection: commandline.Device, private_key: Optional[Path]
 ) -> remote_access.ChromiumOSDevice:
-    """Get ChromiumOSDevice host from device_host format string.
+    """Get ChromiumOSDevice host from commandline.Device collection.
 
-    Gets a ChromiumOSDevice, executing as root.
+    The device will always execute commands as root.
 
     Args:
-        device_host: 'hostname:port' or 'hostname' string.
+        collection: Device collection object to connect to.
         private_key: Path to private key to log into the device for.
 
     Returns:
         A remote_access.ChromiumOSDevice connection.
-
-    Raises:
-        ValueError if device_host is not formatted correctly.
     """
-    host = device_host.split(":")
-    if len(host) == 2:
-        device = remote_access.ChromiumOSDevice(
-            host[0],
-            port=host[1],
-            username="root",
-            private_key=private_key,
-            connect=True,
-        )
-    elif len(host) == 1:
-        device = remote_access.ChromiumOSDevice(
-            host[0],
-            username="root",
-            private_key=private_key,
-            connect=True,
-        )
-    else:
-        raise ValueError(f"Badly formatted device host: {device_host}")
-    return device
+    return remote_access.ChromiumOSDevice(
+        collection.hostname,
+        port=collection.port,
+        username="root",
+        private_key=private_key,
+        base_dir=_DUT_FUZZER_ROOT,
+        connect=True,
+    )
+
+
+def _format_collection(collection: commandline.Device) -> str:
+    if collection.port is None:
+        return collection.hostname
+    return f"{collection.hostname}:{collection.port}"
 
 
 def parse_args(raw_args: Iterable[str]):
@@ -132,26 +127,29 @@
     parser.add_argument(
         "--private-key",
         type=Path,
-        default=_DEFAULT_PRIVATE_KEY,
-        help=f"RSA key for device. (default: {_DEFAULT_PRIVATE_KEY})",
+        default=None,
+        help="RSA key for device. (default: lib.remote_access's key path)",
     )
-    subparsers = parser.add_subparsers()
 
+    # Device parsing aids.
+    device_parser = commandline.DeviceParser(commandline.DEVICE_SCHEME_SSH)
+    device_host_help = "Device host, in the format '[ssh://]hostname[:port]'"
+
+    subparsers = parser.add_subparsers()
     # setup subcommand
-    setup_parser = subparsers.add_parser(
+    subparser = subparsers.add_parser(
         "setup", help="Installs fuzzers on device sysroot."
     )
-    setup_parser.set_defaults(func=setup_main)
-    device_host_help = "Device host, in the format 'hostname:port'"
-    setup_parser.add_argument(
+    subparser.set_defaults(func=setup_main)
+    subparser.add_argument(
         "--device-sysroot-path",
         type=Path,
         help="Absolute location on device for the"
         f" fuzzer sysroot. (default: {_DUT_FUZZER_ROOT})",
         default=_DUT_FUZZER_ROOT,
     )
-    setup_parser.add_argument("device_host", help=device_host_help)
-    setup_parser.add_argument(
+    subparser.add_argument("device", type=device_parser, help=device_host_help)
+    subparser.add_argument(
         "packages",
         help="Portage packages to install fuzzers for."
         " Space or comma separated."
@@ -159,12 +157,12 @@
     )
 
     # fuzz subcommand
-    fuzz_parser = subparsers.add_parser(
+    subparser = subparsers.add_parser(
         "fuzz", help="Runs a fuzzer in device sysroot."
     )
-    fuzz_parser.set_defaults(func=fuzz_main)
-    fuzz_parser.add_argument("device_host", help=device_host_help)
-    fuzz_parser.add_argument(
+    subparser.set_defaults(func=fuzz_main)
+    subparser.add_argument("device", type=device_parser, help=device_host_help)
+    subparser.add_argument(
         "fuzzer_name", help="Name of the fuzzer executable to run."
     )