#!/usr/bin/python3 -su

# Copyright (C) 2025 - 2025 ENCRYPTED SUPPORT LLC <adrelanos@whonix.org>
# See the file COPYING for copying conditions.

# TODO: This code can probably be expanded significantly to take more edge
# cases into account. In particular:
# - It can't reliably handle the case where two different filesystems have
#   been mounted to the same mountpoint, with one hiding the other.
# - It can't handle software RAID yet.

from pathlib import Path
import sys
import subprocess
import os

def find_backing_items(path):
    # Finds the device or devices that back a particular mountpoint or file.
    # This uses a recursive algorithm to find the ultimate backing device. The
    # algorithm is, roughly:
    #
    # - If the path is one of /dev/mmc*, /dev/nvme*, /dev/sd*, /dev/vd*,
    #   /dev/xvd*, /dev/hd*, or /dev/sr*, we've found the final device, return
    #   it
    # - If the path starts with /dev/dm, we have a device-mapper device, find
    #   the devices that compose it and resolve their backing items
    # - If the path starts with /dev/loop*, we have a loop device, use
    #   losetup -l and find the underlying file, then resolve its backing
    #   item.
    # - If the path starts with /dev/mapper, we have a device-mapper
    #   symlink, resolve it to find the real device-mapper device and then
    #   resolve its backing item
    # - If the path is just to a normal file or directory on the file system,
    #   dig for the nearest mountpoint and resolve its backing item

    path_obj = Path(path)

    ## TODO: Good idea?
    #path_obj = path_obj.resolve()
    #path = str(path_obj)

    print(f"path: {path}", file=sys.stderr)

    output_objs = []

    ## Special-case looking for the backing device for the root partition when
    ## running in ISO live mode. In this instance, we want to look for the
    ## backing device of /run/initramfs/live, because both / and
    ## /run/rootfsbase go back to /dev/loop0, which goes to
    ## /live/filesystem.squashfs, which no longer exists at this point because
    ## the root filesystem has pivoted and /live is now /run/initramfs/live.
    ##
    ## NOTE: The missing single quote in the below string is intentional, to
    ## match iso-live, iso-live-semi-persistent, etc.
    if (
        "live_status_detected_live_mode_environment_machine='iso-live"
        in subprocess.run(
            "/usr/libexec/helper-scripts/live-mode.sh",
            capture_output=True,
            encoding="utf-8",
        ).stdout
        and path == "/"
    ):
        output_objs.extend(find_backing_items("/run/initramfs/live"))
        return output_objs

    if path_obj.is_block_device():
        print("is block device: yes", file=sys.stderr)
        if (
            path.startswith("/dev/mmc")
            or path.startswith("/dev/nvme")
            or path.startswith("/dev/sd")
            or path.startswith("/dev/vd")
            or path.startswith("/dev/xvd")
            or path.startswith("/dev/hd")
            or path.startswith("/dev/sr")
        ):
            if path.count("/") == 2:
                output_objs.append(path_obj)
                return output_objs
            else:
                raise ValueError("Failed to look up backing item! (code 1)")
        elif path.startswith("/dev/loop"):
            if path.count("/") != 2:
                raise ValueError("Failed to look up backing item! (code 2)")
            if " " in path:
                raise ValueError("Failed to look up backing item! (code 3)")
            loop_name = Path(path).name ## e.g. loop0
            backing_file = Path(f"/sys/class/block/{loop_name}/loop/backing_file")
            try:
                backing = backing_file.read_text().strip()
            except FileNotFoundError:
                raise ValueError("Failed to look up backing item! (code 4.1)")
            if not backing:
                raise ValueError("Failed to look up backing item! (code 4.2)")
            output_objs.extend(find_backing_items(backing))
            return output_objs
        elif path.startswith("/dev/dm"):
            if path.count("/") != 2:
                raise ValueError("Failed to look up backing item! (code 5)")
            if " " in path:
                raise ValueError("Failed to look up backing item! (code 6)")
            slaves_dir = Path(f"/sys/class/block/{path_obj.name}/slaves")
            print(f"dm slaves_dir: {slaves_dir}", file=sys.stderr)
            if not slaves_dir.exists():
                print(f"dm slaves_dir does not exist, ok.", file=sys.stderr)
                return output_objs
            for dev_name in slaves_dir.iterdir():
                dev_path = "/dev/" + dev_name.name
                output_objs.extend(find_backing_items(dev_path))
            return output_objs
        elif path.startswith("/dev/mapper"):
            # make sure we aren't dealing with a device named something
            # like /dev/mapperabc, we're looking for a device named
            # something like /dev/mapper/luks-UUID
            if path.count("/") != 3:
                raise ValueError("Failed to look up backing item! (code 7)")
            if " " in path:
                raise ValueError("Failed to look up backing item! (code 8)")
            mapper_name = path.split("/")[3]
            if mapper_name == "":
                raise ValueError("Failed to look up backing item! (code 9)")
            # need to be in /dev/mapper to properly resolve symlinks
            os.chdir("/dev/mapper")
            try:
                real_dev_path = str(path_obj.resolve(strict=True))
            except FileNotFoundError:
                raise ValueError("Failed to look up backing item! (code 10.1)")
            if real_dev_path == "":
                raise ValueError("Failed to look up backing item! (code 10.2)")
            print(f"real_dev_path: {real_dev_path}", file=sys.stderr)
            output_objs.extend(find_backing_items(real_dev_path))
            return output_objs
        else:
            raise ValueError("Failed to look up backing item! (code 11)")
    else:
        print("is block device: no", file=sys.stderr)
        if not path_obj.exists():
            raise ValueError("Failed to look up backing item! (code 12)")
        # Walk up to the nearest mountpoint (which may be path_obj itself)
        while not path_obj.is_mount():
            temp_path_obj = path_obj.parent
            if str(temp_path_obj) == "/":
                # don't allow looping back around to the root dir
                raise ValueError("Failed to look up backing item! (code 13)")
            path_obj = temp_path_obj
        path = str(path_obj)
        mount_data = Path("/proc/self/mounts").read_text()
        if mount_data == "":
            raise ValueError("Failed to look up backing item! (code 14)")
        for mount_line in mount_data.splitlines():
            mount_parts = mount_line.split(" ")
            if mount_parts[1] != path:
                continue
            if mount_parts[0].startswith("/dev"):
                output_objs.extend(find_backing_items(mount_parts[0]))
                return output_objs
            elif mount_parts[2] == "overlay":
                overlay_data_parts = mount_parts[3].split(",")
                for overlay_data_part in overlay_data_parts:
                    if overlay_data_part.startswith("lowerdir="):
                        output_objs.extend(
                            find_backing_items(
                                overlay_data_part.split("=", maxsplit=1)[1]
                            )
                        )
                        return output_objs
            else:
                ## Out-commented. This may happen if mutliple filesystems are
                ## mounted to the same location and the mount we're looking
                ## for is shadowing an earlier mount. This scenario occurs
                ## when booting in live mode under at least Trixie.
                #raise ValueError("Failed to look up backing item! (code 15)")
                continue

    # If we get here, we weren't able to find the backing device.
    raise ValueError("Failed to look up backing item! (code 16)")

try:
    path_list = find_backing_items(sys.argv[1])
except Exception as e:
    print(f"Could not find backing device(s)! reason: {e}", file=sys.stderr)
    exit(1)

for path in path_list:
    print(str(path))
