path: root/tests/wipedisk/test_wipedisk.py

"""Test scripts/wipedisk, the manual disk-erase helper.

Three defects, all of which make the script's final word untrue:

1. `sgdisk --zap-all` had its result discarded and the script printed
   "Disk erased." and exited 0 unconditionally. sgdisk fails on a busy device
   -- a disk with a mounted filesystem or an active md/LVM/ZFS holder, which is
   exactly the case a user hits when they pick the wrong disk -- so the tool
   claimed an erase it had not performed.

2. The prompt says "Select the disk id to use", and then listed every entry in
   /dev/disk/by-id, two thirds of which are partitions (18 entries on this
   machine, 12 of them -partN). The menu promised disks and offered partitions.

3. "Disk erased." overstates what the tool does. `sgdisk --zap-all` destroys
   partition tables, not data, and `blkdiscard -f || true` deliberately
   tolerates a device that cannot discard. On a disk without discard support
   the script erased the partition table and every byte of data remained
   readable -- while telling the user the disk was erased.

Method: run the REAL script with WIPEDISK_BY_ID pointed at a fixture directory
and fake blkdiscard/sgdisk on PATH, feeding the select menu and the confirm
prompt on stdin. The env-overridable path follows nvidia_preflight_report's
NVIDIA_DRM_GLOB.

Run from repo root:
    python3 -m unittest tests.wipedisk.test_wipedisk
"""

import os
import shutil
import stat
import subprocess
import tempfile
import unittest

REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
SCRIPT = os.path.join(REPO_ROOT, "scripts", "wipedisk")

# A realistic by-id listing: two whole disks, each with two partitions, plus
# the eui alias the kernel also publishes.
# Names that cannot collide with a real device. The first cut of this suite
# copied this machine's actual by-id entries, so a test that lost the
# WIPEDISK_BY_ID seam would have driven a menu of real disks and still passed.
DISKS = [
    "wipedisk-fixture-disk-A",
    "wipedisk-fixture-disk-B",
]
PARTS = [d + p for d in DISKS for p in ("-part1", "-part2")]


class Wipedisk(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="wipedisk-test-")
        self.bin = os.path.join(self.tmp, "bin")
        self.byid = os.path.join(self.tmp, "by-id")
        os.makedirs(self.bin)
        os.makedirs(self.byid)
        for name in DISKS + PARTS:
            open(os.path.join(self.byid, name), "w").close()

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def fake(self, name, body):
        path = os.path.join(self.bin, name)
        with open(path, "w") as f:
            f.write("#!/bin/bash\n" + body)
        os.chmod(path, os.stat(path).st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)

    def install_fakes(self, discard_ok=True, sgdisk_ok=True,
                      discard_err="blkdiscard: BLKDISCARD ioctl failed: "
                                  "Operation not supported"):
        log = os.path.join(self.tmp, "calls.log")
        self.fake("blkdiscard",
                  'echo "blkdiscard $*" >> "%s"\n'
                  '[ %d -eq 0 ] || echo "%s" >&2\nexit %d\n'
                  % (log, 0 if discard_ok else 1, discard_err,
                     0 if discard_ok else 1))
        self.fake("sgdisk", 'echo "sgdisk $*" >> "%s"\nexit %d\n'
                  % (log, 0 if sgdisk_ok else 1))
        self.calls_log = log

    def calls(self):
        if not os.path.exists(self.calls_log):
            return ""
        with open(self.calls_log) as f:
            return f.read()

    def run_script(self, stdin):
        env = dict(os.environ)
        env["PATH"] = self.bin + os.pathsep + env["PATH"]
        env["WIPEDISK_BY_ID"] = self.byid
        return subprocess.run(
            ["bash", SCRIPT], input=stdin, capture_output=True, text=True,
            env=env, timeout=20,
        )

    # ---------------------------------------------------------- normal ----
    def test_successful_wipe_exits_zero_and_calls_both_tools(self):
        self.install_fakes()
        r = self.run_script("1\ny\n")
        self.assertEqual(r.returncode, 0, r.stdout + r.stderr)
        # Assert the fixture PATH, not just the tool name: this is what
        # catches a lost WIPEDISK_BY_ID seam driving real devices.
        self.assertIn("blkdiscard %s/%s" % (self.byid, DISKS[0]), self.calls())
        self.assertIn("sgdisk --zap-all %s/%s" % (self.byid, DISKS[0]),
                      self.calls())

    def test_declining_the_confirmation_touches_nothing(self):
        self.install_fakes()
        r = self.run_script("1\nn\n")
        self.assertNotEqual(r.returncode, 0)
        self.assertEqual(self.calls(), "",
                         "a declined confirmation must not run either tool")

    # -------------------------------------------------------- boundary ----
    def test_menu_offers_whole_disks_only(self):
        # The prompt says "Select the disk id"; partitions are not disks.
        self.install_fakes()
        r = self.run_script("1\nn\n")
        menu = r.stdout + r.stderr
        for part in PARTS:
            self.assertNotIn(part, menu,
                             "the disk menu must not offer partition %s" % part)
        for disk in DISKS:
            self.assertIn(disk, menu)

    def test_discard_unsupported_does_not_claim_the_data_is_gone(self):
        # blkdiscard failing is tolerated by design (not every device can
        # discard) -- but then only the partition table went, and the message
        # must not tell the user the disk was erased.
        self.install_fakes(discard_ok=False)
        r = self.run_script("1\ny\n")
        out = r.stdout + r.stderr
        self.assertIn("sgdisk", self.calls(),
                      "a failed discard must not stop the partition-table zap")
        self.assertNotIn("Disk erased.", out,
                         "data was not erased -- only the partition table was")

    # ----------------------------------------------------------- error ----
    def test_failed_zap_exits_non_zero(self):
        # sgdisk fails on a busy device. Claiming success here is the defect.
        self.install_fakes(sgdisk_ok=False)
        r = self.run_script("1\ny\n")
        self.assertNotEqual(r.returncode, 0,
                            "a failed zap must not report success")
        self.assertNotIn("Disk erased.", r.stdout + r.stderr)

    def test_empty_by_id_directory_says_so(self):
        # Without the guard this still exits non-zero and runs nothing -- the
        # empty select menu falls through to the confirm prompt, which reads
        # EOF and declines. So asserting only on the exit code and the call log
        # is not a gate; it passes with the guard deleted. What the guard
        # actually buys is saying why, instead of an empty menu the user has to
        # break out of.
        self.install_fakes()
        for name in os.listdir(self.byid):
            os.unlink(os.path.join(self.byid, name))
        r = self.run_script("1\ny\n")
        self.assertNotEqual(r.returncode, 0)
        self.assertEqual(self.calls(), "",
                         "no disks means nothing to erase, not an unguarded run")
        self.assertIn("No disks found", r.stdout + r.stderr,
                      "an empty device list must be reported, not shown as an empty menu")


if __name__ == "__main__":
    unittest.main()


class WipediskBusyDisk(unittest.TestCase):
    """The wrong-disk case, which is the whole reason this tool checks anything.

    blkdiscard opens the device O_EXCL by default (util-linux >= 2.36) and
    refuses a disk something is holding -- a mounted filesystem, an md member,
    an LVM PV. That refusal is the safety gate, so it has to come first. Forcing
    past it with -f destroyed a live filesystem and only then let sgdisk fail,
    leaving the user told that nothing had happened and to try again.

    A discard that fails because the hardware cannot discard is a different
    answer entirely, and must not stop the run.
    """

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="wipedisk-busy-test-")
        self.bin = os.path.join(self.tmp, "bin")
        self.byid = os.path.join(self.tmp, "by-id")
        os.makedirs(self.bin)
        os.makedirs(self.byid)
        for name in DISKS + PARTS:
            open(os.path.join(self.byid, name), "w").close()
        self.log = os.path.join(self.tmp, "calls.log")

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def fake(self, name, body):
        path = os.path.join(self.bin, name)
        with open(path, "w") as f:
            f.write("#!/bin/bash\n" + body)
        os.chmod(path, 0o755)

    def install(self, discard_rc, discard_err, sgdisk_rc=0):
        self.fake("blkdiscard",
                  'echo "blkdiscard $*" >> "%s"\n'
                  '[ %d -eq 0 ] || echo "%s" >&2\nexit %d\n'
                  % (self.log, discard_rc, discard_err, discard_rc))
        self.fake("sgdisk", 'echo "sgdisk $*" >> "%s"\nexit %d\n'
                  % (self.log, sgdisk_rc))

    def calls(self):
        if not os.path.exists(self.log):
            return ""
        with open(self.log) as f:
            return f.read()

    def run_script(self, stdin="1\ny\n"):
        env = dict(os.environ)
        env["PATH"] = self.bin + os.pathsep + env["PATH"]
        env["WIPEDISK_BY_ID"] = self.byid
        return subprocess.run(["bash", SCRIPT], input=stdin,
                              capture_output=True, text=True, env=env,
                              timeout=20)

    # ---------------------------------------------------------- normal ----
    def test_blkdiscard_is_not_forced_past_the_exclusive_open(self):
        # -f is what turned "refuse a busy disk" into "destroy it anyway".
        self.install(0, "")
        self.run_script()
        lines = [l for l in self.calls().splitlines()
                 if l.startswith("blkdiscard ")]
        self.assertEqual(len(lines), 1, self.calls())
        # Exact argv tokens: "-f" is a substring of the fixture disk name, so
        # assertNotIn would pass on the fixture path and prove nothing.
        args = lines[0].split()[1:]
        self.assertNotIn("-f", args)
        self.assertNotIn("--force", args)

    # -------------------------------------------------------- boundary ----
    def test_a_device_that_cannot_discard_still_gets_its_table_cleared(self):
        self.install(1, "blkdiscard: BLKDISCARD ioctl failed: "
                        "Operation not supported")
        r = self.run_script()
        self.assertEqual(r.returncode, 0, r.stdout + r.stderr)
        self.assertIn("sgdisk", self.calls())
        self.assertIn("still", r.stdout)   # the honest data-still-present note

    # ----------------------------------------------------------- error ----
    def test_a_busy_disk_stops_before_sgdisk_is_ever_called(self):
        self.install(1, "blkdiscard: /dev/sda: BLKDISCARD ioctl failed: "
                        "Device or resource busy")
        r = self.run_script()
        self.assertNotEqual(r.returncode, 0)
        self.assertNotIn("sgdisk", self.calls(),
                         "a refused discard must not fall through to sgdisk")

    def test_a_busy_disk_says_nothing_was_erased(self):
        self.install(1, "blkdiscard: /dev/sda: BLKDISCARD ioctl failed: "
                        "Device or resource busy")
        r = self.run_script()
        out = r.stdout + r.stderr
        self.assertIn("Nothing was erased", out)
        self.assertNotIn("Disk erased", out)

    def test_a_busy_disk_never_reports_a_cleared_partition_table(self):
        # The defect this class exists for: the old order discarded first and
        # then said "could not clear the partition table ... run this again",
        # which reads as "nothing happened" after data was already gone.
        self.install(1, "blkdiscard: /dev/sda: BLKDISCARD ioctl failed: "
                        "Device or resource busy")
        r = self.run_script()
        self.assertNotIn("Partition table cleared", r.stdout + r.stderr)