diff --git a/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py b/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py index 09d47293d..493eea04f 100644 --- a/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py +++ b/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py @@ -476,7 +476,7 @@ def base(): @base.command() @click.argument("file") - @click.option("--partition", type=str) + @click.option("--target", type=str) @click.option("--os-image-checksum", help="SHA256 checksum of OS image (direct value)") @click.option( "--os-image-checksum-file", @@ -488,7 +488,7 @@ def base(): @debug_console_option def flash( file, - partition, + target, os_image_checksum, os_image_checksum_file, console_debug, @@ -504,7 +504,7 @@ def flash( self.set_console_debug(console_debug) self.flash( file, - partition=partition, + partition=target, force_exporter_http=force_exporter_http, force_flash_bundle=force_flash_bundle, ) diff --git a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py index e75d0b0a7..2d7fbef62 100644 --- a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py +++ b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from io import BytesIO from pathlib import Path +from typing import cast from urllib.parse import urlparse from uuid import UUID @@ -514,10 +515,10 @@ class FlasherClientInterface(metaclass=ABCMeta): @abstractmethod def flash( self, - path: PathBuf, + path: PathBuf | dict[str, PathBuf], *, - partition: str | None = None, - operator: Operator | None = None, + target: str | None = None, + operator: Operator | dict[str, Operator] | None = None, compression: Compression | None = None, ): """Flash image to DUT""" @@ -528,7 +529,7 @@ def dump( self, path: PathBuf, *, - partition: str | None = None, + target: str | None = None, operator: Operator | None = None, compression: Compression | None = None, ): @@ -542,45 +543,118 @@ def base(): pass @base.command() - @click.argument("file") - @click.option("--partition", type=str) + @click.argument("file", nargs=-1, required=False) + @click.option( + "--target", + "-t", + "target_specs", + multiple=True, + help="name:file", + ) @click.option("--compression", type=click.Choice(Compression, case_sensitive=False)) - def flash(file, partition, compression): - """Flash image to DUT from file""" - self.flash(file, partition=partition, compression=compression) + def flash(file, target_specs, compression): + if target_specs: + mapping: dict[str, str] = {} + for spec in target_specs: + if ":" not in spec: + raise click.ClickException(f"Invalid target spec '{spec}', expected name:file") + name, img = spec.split(":", 1) + mapping[name] = img + self.flash(cast(dict[str, PathBuf], mapping), compression=compression) + return + + if not file: + raise click.ClickException("FILE argument is required unless --target/-t is used") + + self.flash(file[0], target=None, compression=compression) @base.command() @click.argument("file") - @click.option("--partition", type=str) + @click.option("--target", type=str) @click.option("--compression", type=click.Choice(Compression, case_sensitive=False)) - def dump(file, partition, compression): + def dump(file, target, compression): """Dump image from DUT to file""" - self.dump(file, partition=partition, compression=compression) + self.dump(file, target=target, compression=compression) return base class FlasherClient(FlasherClientInterface, DriverClient): - def flash( + def _should_upload_file(self, storage, filename: str, src_path: PathBuf, src_operator: Operator) -> bool: + """Check if file should be uploaded by comparing existence and hash.""" + if not storage.exists(filename): + return True + + try: + import hashlib + + m = hashlib.sha256() + with src_operator.open(src_path, "rb") as f: + while True: + data = f.read(size=65536) + if len(data) == 0: + break + m.update(data) + src_hash = m.hexdigest() + + storage_hash = storage.hash(filename) + + if storage_hash == src_hash: + return False + else: + return True + except Exception: + return True + + def _flash_single( self, - path: PathBuf, + image: PathBuf, *, - partition: str | None = None, - operator: Operator | None = None, - compression: Compression | None = None, + target: str | None, + operator: Operator | None, + compression: Compression | None, ): """Flash image to DUT""" if operator is None: - path, operator, _ = operator_for_path(path) + image, operator, _ = operator_for_path(image) - with OpendalAdapter(client=self, operator=operator, path=path, mode="rb", compression=compression) as handle: - return self.call("flash", handle, partition) + with OpendalAdapter(client=self, operator=operator, path=image, mode="rb", compression=compression) as handle: + return self.call("flash", handle, target) + + def flash( + self, + path: PathBuf | dict[str, PathBuf], + *, + target: str | None = None, + operator: Operator | dict[str, Operator] | None = None, + compression: Compression | None = None, + ): + if isinstance(path, dict): + if target is not None: + raise ArgumentError("'target' parameter is not valid when flashing multiple images") + + results: dict[str, object] = {} + + oper_map = operator if isinstance(operator, dict) else {} + + for part, img in path.items(): + op_val = oper_map.get(part) if isinstance(operator, dict) else operator + results[part] = self._flash_single( + img, target=part, operator=cast(Operator | None, op_val), compression=compression + ) + + return results + + if isinstance(operator, dict): + raise ArgumentError("operator mapping provided for single image flash") + + return self._flash_single(path, target=target, operator=operator, compression=compression) def dump( self, path: PathBuf, *, - partition: str | None = None, + target: str | None = None, operator: Operator | None = None, compression: Compression | None = None, ): @@ -589,7 +663,7 @@ def dump( path, operator, _ = operator_for_path(path) with OpendalAdapter(client=self, operator=operator, path=path, mode="wb", compression=compression) as handle: - return self.call("dump", handle, partition) + return self.call("dump", handle, target) class StorageMuxClient(DriverClient): @@ -648,7 +722,7 @@ def off(): """Disconnect storage""" self.off() - @base.command() + @base.command @click.argument("file") def write_local_file(file): self.write_local_file(file) @@ -661,13 +735,13 @@ def flash( self, path: PathBuf, *, - partition: str | None = None, + target: str | None = None, operator: Operator | None = None, compression: Compression | None = None, ): """Flash image to DUT""" - if partition is not None: - raise ArgumentError(f"partition is not supported for StorageMuxFlasherClient, {partition} provided") + if target is not None: + raise ArgumentError(f"target is not supported for StorageMuxFlasherClient, {target} provided") self.host() @@ -684,13 +758,13 @@ def dump( self, path: PathBuf, *, - partition: str | None = None, + target: str | None = None, operator: Operator | None = None, compression: Compression | None = None, ): """Dump image from DUT""" - if partition is not None: - raise ArgumentError(f"partition is not supported for StorageMuxFlasherClient, {partition} provided") + if target is not None: + raise ArgumentError(f"target is not supported for StorageMuxFlasherClient, {target} provided") self.call("host") diff --git a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py index 5986eb8ab..e327276ce 100644 --- a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py +++ b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py @@ -209,7 +209,7 @@ def client(cls) -> str: return "jumpstarter_driver_opendal.client.FlasherClient" @abstractmethod - def flash(self, source, partition: str | None = None): ... + def flash(self, source, target: str | None = None): ... @abstractmethod def dump(self, target, partition: str | None = None): ... diff --git a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py index b4808c592..6e5a33c21 100644 --- a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py +++ b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py @@ -139,13 +139,13 @@ def test_driver_opendal_presign(tmp_path): ) -@pytest.mark.parametrize("partition", [None, "uboot"]) -def test_driver_flasher(tmp_path, partition): +@pytest.mark.parametrize("target", [None, "uboot"]) +def test_driver_flasher(tmp_path, target): with serve(MockFlasher()) as flasher: (tmp_path / "disk.img").write_bytes(b"hello") - flasher.flash(tmp_path / "disk.img", partition=partition) - flasher.dump(tmp_path / "dump.img", partition=partition) + flasher.flash(tmp_path / "disk.img", target=target) + flasher.dump(tmp_path / "dump.img", target=target) assert (tmp_path / "dump.img").read_bytes() == b"hello" diff --git a/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py b/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py index babc9c432..47d4ae90b 100644 --- a/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py +++ b/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py @@ -5,8 +5,7 @@ import click from jumpstarter_driver_composite.client import CompositeClient -from jumpstarter_driver_opendal.client import operator_for_path -from jumpstarter_driver_opendal.common import PathBuf +from jumpstarter_driver_opendal.client import FlasherClient, operator_for_path from jumpstarter_driver_power.client import PowerClient from opendal import Operator @@ -14,7 +13,7 @@ @dataclass(kw_only=True) -class RideSXClient(CompositeClient): +class RideSXClient(FlasherClient, CompositeClient): """Client for RideSX""" def __post_init__(self): @@ -27,23 +26,20 @@ def _upload_file_if_needed(self, file_path: str, operator: Operator | None = Non if operator is None: path_buf, operator, operator_scheme = operator_for_path(file_path) else: - path_buf = PathBuf(file_path) + path_buf = Path(file_path) operator_scheme = "unknown" - if isinstance(path_buf, Path): - filename = path_buf.name - else: - filename = Path(str(path_buf)).name + filename = Path(path_buf).name - if self.storage.exists(filename): - self.logger.info(f"File {filename} already exists in storage, skipping upload") - else: + if self._should_upload_file(self.storage, filename, path_buf, operator): if operator_scheme == "http": self.logger.info(f"Downloading {file_path} to storage as {filename}") else: self.logger.info(f"Uploading {file_path} to storage as {filename}") self.storage.write_from_path(filename, path_buf, operator=operator) + else: + self.logger.info(f"File {filename} already exists in storage with matching hash, skipping upload") return filename @@ -78,62 +74,41 @@ def flash_images(self, partitions: Dict[str, str], operators: Optional[Dict[str, return flash_result - def flash(self, partitions: Dict[str, str], operators: Optional[Dict[str, Operator]] = None): - """Flash partitions to the device""" + def flash( + self, + path: str | Dict[str, str], + *, + partition: str | None = None, + operator: Operator | Dict[str, Operator] | None = None, + compression=None, + ): + if isinstance(path, dict): + partitions = path + operators = operator if isinstance(operator, dict) else None + else: + if partition is None: + raise ValueError("'partition' must be provided") + partitions = {partition: path} + operators = {partition: operator} if isinstance(operator, Operator) else None + self.logger.info("Starting RideSX flash operation") self.boot_to_fastboot() - self.logger.info("waiting for fastboot mode to be ready...") - time.sleep(3) - result = self.flash_images(partitions, operators) self.logger.info("flash operation completed successfully") return result def cli(self): + generic_cli = FlasherClient.cli(self) + @click.group() def storage(): - """Storage operations""" pass - @storage.command() - @click.option("--partition", "-p", multiple=True, help="Partition to flash in format partition:file") - def flash(partition): - """Flash partitions to device. - - Examples: - j storage flash --partition aboot:/path/to/aboot.img --partition rootfs:/path/to/rootfs.img - j storage flash -p boot:boot.img -p system:system.img -p userdata:userdata.img - """ - if not partition: - click.echo("Error: At least one --partition must be provided") - click.echo("Usage: j storage flash --partition : [--partition : ...]") - raise click.Abort() - - partitions = {} - - for part_spec in partition: - if ":" not in part_spec: - click.echo(f"Error: Invalid partition format '{part_spec}'. Expected format: partition:file") - raise click.Abort() - - partition_name, file_path = part_spec.split(":", 1) - if not partition_name or not file_path: - click.echo( - f"Error: Invalid partition format '{part_spec}'. Both partition and file must be specified" - ) - raise click.Abort() - - partitions[partition_name] = file_path - - try: - self.flash(partitions) - click.echo("Flash operation completed successfully") - except Exception as e: - click.echo(f"Flash operation failed: {e}") - raise click.Abort() from e + for name, cmd in generic_cli.commands.items(): + storage.add_command(cmd, name=name) return storage diff --git a/packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/driver_test.py b/packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/driver_test.py index 04c3cd76f..bf6077eaa 100644 --- a/packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/driver_test.py +++ b/packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/driver_test.py @@ -50,7 +50,7 @@ def test_driver_uboot_console(uboot_image): } ) ) as root: - root.qemu.flasher.flash(uboot_image, partition="bios") + root.qemu.flasher.flash(uboot_image, target="bios") uboot = root.uboot