remap-badblocks 0.7__py3-none-any.whl

cli/commands/update.py

@@ -0,0 +1,208 @@
+from argparse import Namespace
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Iterable, Literal, Optional, TypedDict
+
+from remap_badblocks.src.badblocks._compute_good_ranges import (
+    compute_good_ranges, reserve_space_from_good_ranges)
+from remap_badblocks.src.badblocks._find_badblocks import (
+    get_all_badblocks, read_known_badblocks)
+from remap_badblocks.src.badblocks._mapping_generation import generate_mapping
+from remap_badblocks.src.badblocks._remap_badblocks import (
+    count_used_spare_sectors, iter_all_spare_sectors, iter_free_spare_sectors,
+    remap_badblocks, simplify_mapping)
+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, MappingElement
+from remap_badblocks.src.utils._parse_inputs import (
+    parse_bytes_to_sectors, parse_memory_number_to_bytes,
+    parse_memory_range_to_sectors)
+from remap_badblocks.src.utils._run_command import pipe_lines_to_file
+
+
+class UpdateArgs(TypedDict):
+    device_id: Optional[int]
+    device_name: Optional[str]
+    mode: Literal["read", "write", "skip"]
+    output: Optional[Path]
+    block_range: str
+    spare_space: Optional[str]
+    reset_mapping: bool
+    external_known_badblocks_file: Optional[Path]
+
+
+def parse_args(args: Namespace) -> UpdateArgs:
+    return {
+        "device_id": args.id,
+        "device_name": args.name,
+        "mode": args.mode,
+        "output": args.output,
+        "block_range": args.block_range,
+        "spare_space": (None if args.spare_space is None else args.spare_space.strip()),
+        "reset_mapping": args.reset_mapping,
+        "external_known_badblocks_file": args.known_badblocks_file,
+    }
+
+
+def get_device_from_args(args: UpdateArgs, dc: DevicesConfig) -> DeviceConfig:
+    if args["device_id"] is None:
+        if args["device_name"] is not None:
+            return dc.get_device(name=args["device_name"])
+        else:
+            raise ValueError(
+                "Either --id or --name must be provided to update a device in the database."
+            )
+
+    return dc.get_device(id=args["device_id"])
+
+
+def get_update_block_range_from_args(
+    args: UpdateArgs, device: DeviceConfig
+) -> tuple[int, int]:
+    start_block, end_block = parse_memory_range_to_sectors(
+        args["block_range"],
+        sector_size=device.sector_size,
+    )
+
+    if (start_block is not None and start_block < device.logical_range[0]) or (
+        end_block is not None and end_block > device.logical_range[1]
+    ):
+        raise ValueError(
+            f"Block range {start_block}-{end_block} must be contained"
+            f" within the logical range {device.logical_range[0]}-{device.logical_range[1]}."
+        )
+
+    if start_block is None:
+        start_block = device.logical_range[0]
+    if end_block is None:
+        end_block = device.logical_range[1]
+
+    return start_block, end_block
+
+
+def get_n_spare_sectors_from_args(args: UpdateArgs, device: DeviceConfig) -> int:
+    sector_size = device.sector_size
+
+    if args["spare_space"] is None:
+        spare_blocks = None
+    else:
+        spare_blocks = parse_bytes_to_sectors(
+            parse_memory_number_to_bytes(args["spare_space"], sector_size=sector_size),
+            sector_size,
+        )
+
+    if spare_blocks is not None:
+        return spare_blocks
+    else:
+        if device.spare_sectors is None:
+            raise ValueError(
+                "Should specify spare_sectors the first time you map badblocks"
+            )
+        return device.spare_sectors
+
+
+def get_starting_badblocks(args: UpdateArgs, device: DeviceConfig) -> set[int]:
+    known_badblocks: set[int] = set(device.badblocks)
+
+    # If known badblocks file is passed, merge it into our known_badblocks
+    if args["external_known_badblocks_file"]:
+        known_badblocks.update(
+            read_known_badblocks(
+                known_badblocks_file=args["external_known_badblocks_file"]
+            )
+        )
+
+    return known_badblocks
+
+
+def get_current_badblocks(
+    args: UpdateArgs,
+    known_badblocks: set[int],
+    device: DeviceConfig,
+    block_range: tuple[int, int],
+) -> set[int]:
+    if args["mode"] == "skip":
+        return known_badblocks
+
+    with NamedTemporaryFile() as known_badblocks_file:
+        known_badblocks_file.write("\n".join(map(str, known_badblocks)).encode("utf-8"))
+        known_badblocks_file.flush()
+
+        badblocks: set[int] = set(
+            get_all_badblocks(
+                device.path,
+                device.sector_size,
+                mode=args["mode"],
+                known_badblocks_file=Path(known_badblocks_file.name),
+                block_range=block_range,
+            )
+        )
+
+        if args["output"]:
+            pipe_lines_to_file(badblocks, args["output"])
+
+        return badblocks
+
+
+def update_mapping(
+    args: UpdateArgs,
+    device: DeviceConfig,
+    badblocks: Iterable[int],
+    n_spare_sectors: int,
+) -> Mapping:
+    new_mapping: Iterable[tuple[int, int, int]]
+
+    if (device.mapping is None) or (args["reset_mapping"]):
+        good_ranges = compute_good_ranges(
+            badblocks, available_range=device.logical_range
+        )
+        good_ranges = reserve_space_from_good_ranges(good_ranges, n_spare_sectors)
+        good_ranges = reserve_space_from_good_ranges(
+            good_ranges,
+            int(100 * 1024 / device.sector_size),
+        )
+        new_mapping = generate_mapping(
+            good_ranges,
+        )
+    else:
+        spare_sectors = iter_all_spare_sectors(n_spare_sectors, device.logical_range[0])
+        n_used_spare_sectors = count_used_spare_sectors(device.mapping, spare_sectors)
+        current_mapping = map(lambda x: x.to_tuple(), device.mapping)
+        new_mapping = remap_badblocks(
+            current_mapping,
+            badblocks,
+            spare_sectors=iter_free_spare_sectors(spare_sectors, n_used_spare_sectors),
+        )
+
+    new_mapping = simplify_mapping(list(new_mapping))
+    return Mapping(set(map(MappingElement.from_tuple, new_mapping)))
+
+
+def update(dc: DevicesConfig, _args: Namespace) -> None:
+    """
+    Update the mapping of a device in the database.
+    """
+    args = parse_args(_args)
+
+    device: DeviceConfig = get_device_from_args(args, dc)
+    block_range = get_update_block_range_from_args(args, device)
+    n_spare_sectors = get_n_spare_sectors_from_args(args, device)
+
+    known_badblocks = get_starting_badblocks(args, device)
+    badblocks = get_current_badblocks(args, known_badblocks, device, block_range)
+
+    print(
+        f"Badblocks computed. Badblocks went from {len(known_badblocks)} to {len(badblocks)}."
+    )
+
+    dc.update_device(
+        id=device.id,
+        badblocks=badblocks,
+    )
+
+    print("Badblocks list updated.")
+
+    new_mapping = update_mapping(args, device, badblocks, n_spare_sectors)
+    dc.update_device(id=device.id, mapping=new_mapping, spare_sectors=n_spare_sectors)
+
+    print("Mapping updated.")

cli/commands/version.py

@@ -0,0 +1,15 @@
+import os
+from tomllib import load
+
+
+def get_version() -> str:
+
+    pyproject = os.path.join(
+        os.path.dirname(os.path.abspath(__name__)), "pyproject.toml"
+    )
+    with open(pyproject, "rb") as f:
+        return load(f)["project"]["version"]
+
+
+def version():
+    print(f"remap-badblocks version v{get_version()}")

remap_badblocks/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()

remap_badblocks/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"]

remap_badblocks/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
+        ),
+    )

remap_badblocks/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)

remap_badblocks/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))