remap-badblocks 0.7__py3-none-any.whl

cli/__main__.py

@@ -0,0 +1,239 @@
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+from typing import Callable
+
+from remap_badblocks.cli.commands import (add, apply, get, remove, update,
+                                          version)
+from remap_badblocks.src.devices.devices_config import DevicesConfig
+from remap_badblocks.src.devices_config_constants import \
+    DEFAULT_DEVICES_CONFIG_PATH
+
+argparse = ArgumentParser(
+    description="Badblocks remapper for block devices.",
+)
+
+argparse.add_argument(
+    "-P",
+    "--db-path",
+    type=Path,
+    default=DEFAULT_DEVICES_CONFIG_PATH,
+    help="Path to the devices configuration database.",
+)
+
+subparser = argparse.add_subparsers(
+    title="Actions",
+    description="Available actions for the badblocks remapper.",
+    dest="action",
+    help="Action to perform",
+    required=True,
+)
+
+add_new_device_parser = subparser.add_parser(
+    "add",
+    help="Add a new device to the database.",
+)
+
+add_new_device_parser.add_argument(
+    "name", type=str, help="The name of the output device."
+)
+add_new_device_parser.add_argument(
+    "--wwn",
+    type=str,
+    help="The WWN of the device to add.",
+)
+add_new_device_parser.add_argument(
+    "--path",
+    type=Path,
+    help="The path to the device to add.",
+)
+add_new_device_parser.add_argument(
+    "--depends-on",
+    type=str,
+    default=[],
+    action="append",
+    help=(
+        "The name of a device that this device depends on. Can be specified multiple times to add multiple"
+        " dependencies. This is used to ensure that devices are applied in the correct order."
+    ),
+)
+add_new_device_parser.add_argument(
+    "--depends-on-id",
+    type=int,
+    default=[],
+    action="append",
+    help=(
+        "The ID of a device that this device depends on. Can be specified multiple times to add multiple"
+        " dependencies. This is used to ensure that devices are applied in the correct order."
+    ),
+)
+add_new_device_parser.add_argument(
+    "--logical-sector-range",
+    type=str,
+    default="-",
+    help=(
+        "Logical sector range to assign to the device. Format: 'start-end', where 'end' can be omitted to mean"
+        " 'till the end of the device'. Both start and end can be specified as sector numbers or as byte offsets."
+        " In case of byte offsets, they must be multiples of the sector size. If not specified, the whole device"
+        " will be used. Note that end is not included in the range."
+    ),
+)
+
+
+get_devices_parser = subparser.add_parser(
+    "get",
+    help="Get the existing device(s) in the database.",
+)
+get_devices_parser.add_argument(
+    "--id",
+    type=int,
+    default=None,
+    help="The ID of the device to get. If neither --id or --name are provided, all devices will be returned.",
+)
+get_devices_parser.add_argument(
+    "--name",
+    type=str,
+    default=None,
+    help="The name of the device to get. If neither --id or --name are provided, all devices will be returned.",
+)
+
+remove_devices_parser = subparser.add_parser(
+    "remove",
+    help="Remove an existing device in the database.",
+)
+remove_devices_parser.add_argument(
+    "--id",
+    type=int,
+    default=None,
+    help="The ID of the device to remove. If neither --id or --name are provided, an error will be raised.",
+)
+remove_devices_parser.add_argument(
+    "--name",
+    type=str,
+    default=None,
+    help="The name of the device to remove. If neither --id or --name are provided, an error will be raised.",
+)
+
+update_mapping_parser = subparser.add_parser(
+    "update",
+    help=(
+        "Update the mapping of a device in the database. More specifically: compute/update the device's badblocks and"
+        " store them in the database, compute the new mapping of a device; if the mapping has changed, the user will"
+        " be prompted to confirm the changes."
+    ),
+)
+update_mapping_parser.add_argument(
+    "--id",
+    type=int,
+    default=None,
+    help="The ID of the device to update.",
+)
+update_mapping_parser.add_argument(
+    "--name",
+    type=str,
+    default=None,
+    help="The name of the device to update.",
+)
+update_mapping_parser.add_argument(
+    "--mode",
+    type=str,
+    choices=["read", "write", "skip"],
+    default="read",
+    help="Mode for computing badblocks",
+)
+update_mapping_parser.add_argument(
+    "--known-badblocks-file",
+    type=Path,
+    default=None,
+    help=(
+        "Path to file with known badblocks (a sector number for each line) that will be merged to those found so far."
+        " Take care the sector size is the same as configured into remap_badblocks"
+    ),
+)
+update_mapping_parser.add_argument(
+    "--block-range",
+    type=str,
+    help=(
+        "Block range to check (e.g., 0-1000, or 1573-), omitted start and end means the whole logical range. Note"
+        " that the end is not included in the range"
+    ),
+    default="-",
+)
+update_mapping_parser.add_argument(
+    "--output",
+    type=Path,
+    default=None,
+    help="Badblocks are stored internally. If provided, they will also be copied to this file.",
+)
+update_mapping_parser.add_argument(
+    "--spare-space",
+    type=str,
+    default=None,
+    help="Number of spare sectors to reserve from the good ranges. Cannot be changed after first time",
+)
+update_mapping_parser.add_argument(
+    "--reset-mapping",
+    action="store_true",
+    help=(
+        "If specified, will ignore existing mapping and build a new mapping from scratch."
+        " DANGER: this might lead to data loss, handle with care"
+    ),
+)
+
+apply_devices_parser = subparser.add_parser(
+    "apply",
+    help="Apply a remapping that's already existing in the database.",
+)
+apply_devices_parser.add_argument(
+    "--id",
+    type=int,
+    default=None,
+    help="The ID of the device to apply. If neither --id or --name are provided, all devices will be applied.",
+)
+apply_devices_parser.add_argument(
+    "--name",
+    type=str,
+    default=None,
+    help="The name of the device to apply. If neither --id or --name are provided, all devices will be applied.",
+)
+apply_devices_parser.add_argument(
+    "--method",
+    type=str,
+    default="device-mapper",
+    choices=["device-mapper"],
+    help="Method to apply the mapping. Only device-mapper is available",
+)
+
+version_parser = subparser.add_parser("version")
+
+
+def main() -> None:
+    args = argparse.parse_args()
+    db_path: Path = args.db_path
+
+    actions: dict[str, Callable[[DevicesConfig, Namespace], None]] = {
+        "add": add,
+        "apply": apply,
+        "get": get,
+        "remove": remove,
+        "update": update,
+    }
+
+    static_actions: dict[str, Callable[[], None]] = {
+        "version": version,
+    }
+
+    if args.action in actions:
+        if db_path.is_dir():
+            raise ValueError("The provided path must not be a directory.")
+        devices_config = DevicesConfig(db_path)
+        actions[args.action](devices_config, args)
+    elif args.action in static_actions:
+        static_actions[args.action]()
+    else:
+        raise ValueError(
+            f"Unknown action: {args.action}. Please use one of {tuple(actions.keys()) + tuple(static_actions.keys())}."
+        )
+
+
+if __name__ == "__main__":
+    main()

cli/commands/__init__.py

@@ -0,0 +1,8 @@
+from .add import add
+from .apply import apply
+from .get import get
+from .remove import remove
+from .update import update
+from .version import version
+
+__all__ = ["add", "apply", "get", "remove", "update", "version"]

cli/commands/add.py

@@ -0,0 +1,91 @@
+import os
+from argparse import Namespace
+from pathlib import Path
+from typing import Optional, TypedDict
+
+from remap_badblocks.src.devices.devices_config import DevicesConfig
+from remap_badblocks.src.utils._get_device_info import (
+    get_disk_block_size, get_disk_number_of_blocks, resolve_device_name)
+from remap_badblocks.src.utils._parse_inputs import \
+    parse_memory_range_to_sectors
+
+
+class AddArgs(TypedDict):
+    wwn: Optional[str]
+    path: Optional[Path]
+    name: str
+    depends_on_id: set[int]
+    depends_on: set[str]
+    logical_sector_range: str
+
+
+def parse_args(args: Namespace) -> AddArgs:
+    return {
+        "wwn": args.wwn,
+        "path": args.path,
+        "name": args.name,
+        "depends_on_id": set(args.depends_on_id),
+        "depends_on": set(args.depends_on),
+        "logical_sector_range": args.logical_sector_range.strip(),
+    }
+
+
+def get_path_from_args(args: AddArgs) -> Path:
+    wwn = args["wwn"]
+    path = args["path"]
+
+    if wwn is not None and path is not None:
+        raise ValueError("Only one of --wwn or --path can be provided, not both.")
+
+    if path is not None:
+        if isinstance(path, str):
+            path = Path(path)
+        return path
+    elif wwn is not None:
+        if not wwn.startswith("/dev/disk/by-id/"):
+            wwn = os.path.join("/dev/disk/by-id/", wwn)
+        return Path(wwn)
+    else:
+        raise ValueError("Either --wwn or --path must be provided.")
+
+
+def get_device_dependencies_from_args(dc: DevicesConfig, args: AddArgs) -> set[int]:
+    return args["depends_on_id"] | {
+        dc.get_device(name=dep).id for dep in args["depends_on"] if dep
+    }
+
+
+def get_logical_block_range_from_args(
+    args: AddArgs, sector_size: int, total_sectors: int
+) -> tuple[int, int]:
+    logical_start_block, logical_end_block = parse_memory_range_to_sectors(
+        args["logical_sector_range"], sector_size
+    )
+
+    if logical_start_block is None:
+        logical_start_block = 0
+    if logical_end_block is None:
+        logical_end_block = total_sectors
+
+    return logical_start_block, logical_end_block
+
+
+def add(dc: DevicesConfig, _args: Namespace) -> None:
+    """
+    Add a new device to the database.
+    """
+    args = parse_args(_args)
+    path: Path = get_path_from_args(args)
+
+    sector_size = get_disk_block_size(resolve_device_name(path))
+    total_sectors = get_disk_number_of_blocks(resolve_device_name(path))
+
+    dc.add_device(
+        path=path,
+        name=args["name"],
+        sector_size=sector_size,
+        depends_on=get_device_dependencies_from_args(dc, args),
+        logical_range=get_logical_block_range_from_args(
+            args, sector_size, total_sectors
+        ),
+    )

cli/commands/apply.py

@@ -0,0 +1,81 @@
+from argparse import Namespace
+from pathlib import Path
+from typing import Iterable, Literal, Optional, TypedDict
+
+from remap_badblocks.src.devices.device_config import DeviceConfig
+from remap_badblocks.src.devices.devices_config import DevicesConfig
+from remap_badblocks.src.mapping import Mapping
+from remap_badblocks.src.remappers._check_applied_devices import \
+    filter_applied_devices
+from remap_badblocks.src.remappers._generate_dm_table import generate_dm_table
+from remap_badblocks.src.utils._run_command import run_command_realtime
+from remap_badblocks.src.utils._sort_devices import \
+    sort_devices_by_dependencies
+
+
+class ApplyArgs(TypedDict):
+    device_id: Optional[int]
+    device_name: Optional[str]
+    method: Literal["device-mapper"]
+
+
+def parse_args(args: Namespace) -> ApplyArgs:
+    return {
+        "device_id": args.id,
+        "device_name": args.name,
+        "method": args.method,
+    }
+
+
+def get_user_wanted_device_list(
+    dc: DevicesConfig, args: ApplyArgs
+) -> Iterable[DeviceConfig]:
+    if args["device_id"] is not None:
+        return [dc.get_device(id=args["device_id"])]
+    elif args["device_name"] is not None:
+        return [dc.get_device(name=args["device_name"])]
+    else:
+        return dc.get_devices()
+
+
+def plan_apply_devices(
+    dc: DevicesConfig, devices: Iterable[DeviceConfig]
+) -> Iterable[DeviceConfig]:
+    already_applied_devices = set(filter_applied_devices(dc.get_devices()))
+    return sort_devices_by_dependencies(
+        list(devices), already_applied_devices=already_applied_devices
+    )
+
+
+def apply_device_mapper(device: DeviceConfig):
+    mapping: Optional[Mapping] = device.mapping
+    sector_size: int = device.sector_size
+    device_path: Path = device.path
+
+    assert mapping is not None, "Call ‘remap_badblocks update‘ before applying"
+
+    device_dmtable_identifier = device_path
+
+    dmtable = "\n".join(
+        generate_dm_table(device_dmtable_identifier, mapping, sector_size)
+    )
+
+    for _ in run_command_realtime(["dmsetup", "create", device.name], stdin=dmtable):
+        pass
+
+
+def apply(dc: DevicesConfig, _args: Namespace):
+    args = parse_args(_args)
+
+    devices: Iterable[DeviceConfig] = get_user_wanted_device_list(dc, args)
+    devices = plan_apply_devices(dc, devices)
+
+    # TODO: remove already applied devices
+
+    method = args["method"]
+    assert (
+        method == "device-mapper"
+    ), f"Only ‘device-mapper‘ method is available, found ‘{method}‘"
+
+    for device in devices:
+        apply_device_mapper(device)

cli/commands/get.py

@@ -0,0 +1,28 @@
+from argparse import Namespace
+from typing import Iterable, Optional
+
+from remap_badblocks.src.devices.device_config import DeviceConfig
+from remap_badblocks.src.devices.devices_config import DevicesConfig
+
+
+def parse_id_from_args(args: Namespace) -> Optional[int]:
+    return args.id
+
+
+def get_devices_to_print(
+    dc: DevicesConfig, id: Optional[int]
+) -> Iterable[DeviceConfig]:
+    if id is not None:
+        return (dc.get_device(id=id),)
+    else:
+        return dc.get_devices()
+
+
+def get(dc: DevicesConfig, args: Namespace) -> None:
+    """
+    Get the configuration of a device from the database.
+    """
+    id: Optional[int] = parse_id_from_args(args)
+
+    for device in get_devices_to_print(dc, id):
+        print(str(device))

cli/commands/remove.py

@@ -0,0 +1,45 @@
+from argparse import Namespace
+from typing import Optional, TypedDict
+
+from remap_badblocks.src.devices.devices_config import DevicesConfig
+
+
+class RemoveArgs(TypedDict):
+    id: Optional[int]
+    name: Optional[str]
+
+
+def parse_args(args: Namespace) -> RemoveArgs:
+    return {
+        "id": args.id,
+        "name": args.name,
+    }
+
+
+def get_device_id_from_args(dc: DevicesConfig, args: RemoveArgs) -> int:
+    id: Optional[int] = args["id"]
+    name: Optional[str] = args["name"]
+
+    if id is not None and name is not None:
+        raise ValueError(
+            "Only one of --id or --name can be provided to remove a device from the database."
+        )
+    elif id is None and name is None:
+        raise ValueError(
+            "Either --id or --name must be provided to remove a device from the database."
+        )
+    elif id is not None:
+        return id
+    else:
+        assert name is not None
+        return dc.get_device(name=name).id
+
+
+def remove(dc: DevicesConfig, _args: Namespace) -> None:
+    """
+    Remove the configuration of a device from the database.
+    """
+    args = parse_args(_args)
+    id = get_device_id_from_args(dc, args)
+
+    dc.remove_device(id)