Throughly test a USB flash drive for capacity, reliability and authenticity

usb_drive_health.py

#!/usr/bin/env python3
"""
usb_drive_health.py

Runs the full verification plan for a used USB flash drive:
  1. lsblk device identification
  2. dmesg check (identity / USB class / errors)
  3. wipefs (clean signature baseline)
  4. f3probe (true vs reported capacity - catches capacity fraud)
  5. badblocks (raw read/write bad-block scan)
  6. mkfs.exfat (only if steps 4 and 5 both pass)
  7. f3write / f3read (full surface write/read verify)
  8. speed benchmark (dd write + read)

Must be run as root (sudo), since most of these operations require it.

Usage:
    sudo python3 test_usb_drive.py /dev/sdX
    sudo python3 test_usb_drive.py /dev/sdX --verbose

Device path can be obtained using lsblk (before and after inserting the USB
flash drive).

WARNING: This script is DESTRUCTIVE. It will wipe and reformat the
given device. Double- and triple-check the device path before
confirming. Never point this at your system disk.

TODO:
- Add timestamps to outputs
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime

VERBOSE = False
_APT_UPDATED = False


class Result:
    def __init__(self, name):
        self.name = name
        self.status = "SKIPPED"
        self.detail = ""


def announce(name):
    print(f"\n=== Running: {name} ===")


def report_status(r):
    print(f"[STATUS] {r.name}: {r.status}")
    if VERBOSE and r.detail:
        print(f"----- RAW OUTPUT: {r.name} -----")
        print(r.detail.strip())
        print(f"----- END: {r.name} -----")


def stream_command(cmd):
    """Run a command, streaming its output live, and return (returncode, full_output).
    Returns (127, detail) if the binary itself is not found, rather than raising."""
    is_shell = isinstance(cmd, str)
    try:
        proc = subprocess.Popen(
            cmd,
            shell=is_shell,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
    except FileNotFoundError:
        missing = cmd if is_shell else cmd[0]
        msg = f"Command not found: {missing}\n"
        print(f"    {msg}", end="")
        return 127, msg
    lines = []
    for line in proc.stdout:
        lines.append(line)
        print(f"    {line}", end="")
    proc.wait()
    return proc.returncode, "".join(lines)


def run_quiet(cmd):
    """Run a short command without streaming, return (returncode, stdout, stderr)."""
    proc = subprocess.run(cmd, shell=isinstance(cmd, str), capture_output=True, text=True)
    return proc.returncode, proc.stdout or "", proc.stderr or ""


def require_root():
    if os.geteuid() != 0:
        sys.exit("This script must be run as root, e.g.: sudo python3 test_usb_drive.py /dev/sdX")


def tool_available(name):
    return shutil.which(name) is not None


def apt_update():
    """Refresh apt's package lists once per run, before the first install attempt."""
    global _APT_UPDATED
    if not _APT_UPDATED:
        announce("apt update")
        stream_command(["apt", "update"])
        _APT_UPDATED = True


def ensure_installed(candidates, binary_name):
    """Try each apt package name in `candidates` in turn until `binary_name` is
    available. Returns True/False rather than raising, so callers can record a
    clean FAIL result and continue instead of crashing the whole run."""
    if tool_available(binary_name):
        return True
    apt_update()
    for pkg in candidates:
        announce(f"Installing {pkg}")
        stream_command(["apt", "install", "-y", pkg])
        if tool_available(binary_name):
            return True
    return False


def get_device_info(device):
    code, out, err = run_quiet(["lsblk", "-J", "-o", "NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT,TYPE,MODEL", device])
    if code != 0:
        sys.exit(f"Could not read device info for {device}: {err.strip()}")
    return json.loads(out)


def confirm_device(device, info):
    dev = info["blockdevices"][0]
    print("\nAbout to run DESTRUCTIVE tests against the following device:")
    print(f"  Path:  {device}")
    print(f"  Size:  {dev.get('size')}")
    print(f"  Model: {dev.get('model', 'unknown')}")
    print(f"  FS:    {dev.get('fstype', 'none')}")
    print("\nThis will erase all data on the device.")
    typed = input(f"Type the device path again to confirm ({device}): ").strip()
    if typed != device:
        sys.exit("Confirmation did not match. Aborting for safety.")


def check_dmesg(results):
    r = Result("dmesg check")
    announce(r.name)
    code, out, err = run_quiet("dmesg | tail -30")
    r.detail = out.strip()
    lowered = out.lower()
    if "reset" in lowered and lowered.count("reset") > 3:
        r.status = "WARN - multiple reset events (possible flaky connector/port)"
    else:
        r.status = "INFO"
    results.append(r)
    report_status(r)


def wipe_signatures(device, results):
    r = Result("wipefs")
    announce(r.name)
    code, out = stream_command(["wipefs", "-a", device])
    r.status = "OK" if code == 0 else "FAIL"
    r.detail = out
    results.append(r)
    report_status(r)


def run_f3probe(device, results):
    r = Result("f3probe (capacity fraud check)")
    announce(r.name)
    if not ensure_installed(["f3"], "f3probe"):
        r.status = "FAIL - could not install f3 (package 'f3')"
        results.append(r)
        report_status(r)
        return False
    code, out = stream_command(["f3probe", "--destructive", "--time-ops", device])
    r.detail = out
    lowered = out.lower()
    if "good news" in lowered:
        r.status = "PASS"
    elif "fake" in lowered or "counterfeit" in lowered:
        r.status = "FAIL - POSSIBLE FAKE CAPACITY"
    elif code == 0:
        r.status = "PASS (review detail)"
    else:
        r.status = "FAIL"
    results.append(r)
    report_status(r)
    return r.status.startswith("PASS")


def run_badblocks(device, results):
    r = Result("badblocks (bad sector scan)")
    announce(r.name)
    if not ensure_installed(["e2fsprogs"], "badblocks"):
        r.status = "FAIL - could not install e2fsprogs (provides badblocks)"
        results.append(r)
        report_status(r)
        return False
    code, out = stream_command(["badblocks", "-wsv", device])
    r.detail = out
    match = re.search(r"Pass completed, (\d+) bad blocks found", out)
    bad_count = int(match.group(1)) if match else None
    if bad_count is None:
        r.status = "PASS (0 bad blocks assumed - check detail)" if code == 0 else "FAIL"
    elif bad_count == 0:
        r.status = "PASS"
    else:
        r.status = f"FAIL - {bad_count} bad blocks found"
    results.append(r)
    report_status(r)
    return r.status.startswith("PASS")


def format_exfat(device, results):
    r = Result("mkfs.exfat")
    announce(r.name)
    if not ensure_installed(["exfatprogs", "exfat-utils"], "mkfs.exfat"):
        r.status = "FAIL - could not install exfatprogs or exfat-utils"
        r.detail = ("Neither package provided mkfs.exfat on this system. Try "
                    "`sudo apt update` manually, check the package name for your "
                    "distro, or install exfat-utils/exfatprogs by hand and re-run.")
        results.append(r)
        report_status(r)
        return False
    code, out = stream_command(["mkfs.exfat", device])
    r.status = "OK" if code == 0 else "FAIL"
    r.detail = out
    results.append(r)
    report_status(r)
    return code == 0


def mount_device(device):
    mountpoint = tempfile.mkdtemp(prefix="usbtest_")
    run_quiet(["mount", device, mountpoint])
    return mountpoint


def unmount_device(mountpoint):
    run_quiet(["umount", mountpoint])
    try:
        os.rmdir(mountpoint)
    except OSError:
        pass


def run_f3write_read(mountpoint, results):
    r = Result("f3write/f3read (surface verify)")
    announce(r.name)
    if not ensure_installed(["f3"], "f3write"):
        r.status = "FAIL - could not install f3 (package 'f3')"
        results.append(r)
        report_status(r)
        return
    stream_command(["f3write", mountpoint])
    code, out = stream_command(["f3read", mountpoint])
    r.detail = out
    lowered = out.lower()
    corrupt_match = re.search(r"(\d+)\s+corrupted", lowered)
    corrupted = int(corrupt_match.group(1)) if corrupt_match else 0
    if corrupted > 0:
        r.status = f"FAIL - {corrupted} corrupted files"
    else:
        r.status = "PASS"
    results.append(r)
    report_status(r)
    # Clean up test files so the drive is left usable
    for name in os.listdir(mountpoint):
        if name.endswith(".h2w") or name.endswith(".fff"):
            try:
                os.remove(os.path.join(mountpoint, name))
            except OSError:
                pass


def speed_test(mountpoint, results):
    r = Result("speed benchmark")
    announce(r.name)
    testfile = os.path.join(mountpoint, "speedtest.bin")
    run_quiet(["sync"])
    code, out = stream_command(
        f"dd if=/dev/zero of={testfile} bs=1M count=256 conv=fdatasync 2>&1"
    )
    m = re.search(r"([\d.]+)\s*MB/s", out)
    write_speed = m.group(1) if m else "unknown"

    run_quiet(["sync"])
    run_quiet("echo 3 > /proc/sys/vm/drop_caches")
    code2, out2 = stream_command(f"dd if={testfile} of=/dev/null bs=1M 2>&1")
    m2 = re.search(r"([\d.]+)\s*MB/s", out2)
    read_speed = m2.group(1) if m2 else "unknown"

    os.remove(testfile)
    r.detail = f"Write: {write_speed} MB/s, Read: {read_speed} MB/s\n{out}\n{out2}"
    r.status = f"INFO - write {write_speed} MB/s, read {read_speed} MB/s"
    results.append(r)
    report_status(r)


def print_report(device, results):
    print("\n" + "=" * 60)
    print(f"TEST REPORT - {device} - {datetime.now().isoformat(timespec='seconds')}")
    print("=" * 60)
    for r in results:
        print(f"- {r.name}: {r.status}")
    print("=" * 60)
    print("Reminder: physically inspect the connector and test in a")
    print("second USB port / machine - not automated by this script.")
    if not VERBOSE:
        print("(Re-run with --verbose to see full raw command output per step.)")


def main():
    global VERBOSE
    parser = argparse.ArgumentParser(description="Full verification test for a used USB flash drive.")
    parser.add_argument("device", help="Device path, e.g. /dev/sdb")
    parser.add_argument(
        "--verbose", action="store_true",
        help="Print full raw output for each step (in addition to live progress), "
             "useful for copy-pasting back for further analysis."
    )
    args = parser.parse_args()
    VERBOSE = args.verbose

    require_root()
    device = args.device

    info = get_device_info(device)
    confirm_device(device, info)

    results = []
    try:
        check_dmesg(results)
        wipe_signatures(device, results)

        capacity_ok = run_f3probe(device, results)
        blocks_ok = run_badblocks(device, results)

        if capacity_ok and blocks_ok:
            if format_exfat(device, results):
                mountpoint = mount_device(device)
                try:
                    run_f3write_read(mountpoint, results)
                    speed_test(mountpoint, results)
                finally:
                    unmount_device(mountpoint)
        else:
            print("\nSkipping format/surface tests - capacity or bad-block check failed.")
            print("Review the report below before deciding whether to raise an")
            print("eBay item-not-as-described case.")
    except Exception as exc:  # noqa: BLE001 - deliberately broad: we must still report
        print(f"\n[UNEXPECTED ERROR] {exc}")
        print("Continuing to report - results gathered before the error are still valid.")
    finally:
        print_report(device, results)


if __name__ == "__main__":
    main()