Python Scripts: add first draft of disk tester

PythonScripts/test_disk.py

@@ -0,0 +1,140 @@
+import argparse
+import subprocess
+import time
+import sys
+import os
+from datetime import datetime
+
+
+def run_cmd(cmd, log_file, check=True, capture_output=True):
+    print(f"Running: {cmd}")
+    result = subprocess.run(cmd, shell=True, capture_output=capture_output, text=True)
+    log_file.write(f"\n$ {cmd}\n")
+    if result.stdout:
+        log_file.write(result.stdout)
+    if result.stderr:
+        log_file.write(result.stderr)
+    if check and result.returncode != 0:
+        print(f"Command failed: {cmd}\n{result.stderr}")
+        sys.exit(1)
+    return result.stdout.strip()
+
+
+def is_nvme(device):
+    return "nvme" in os.path.basename(device)
+
+
+def prompt_overwrite():
+    print("\nWARNING: Disk already has a partition table!")
+    confirm = input("Type 'yes' to overwrite the entire disk: ").strip().lower()
+    while confirm != "yes":
+        confirm = (
+            input("Incorrect input. Please type 'yes' to proceed: ").strip().lower()
+        )
+
+
+def get_device_info(device):
+    print("Fetching device model and serial...")
+    output = subprocess.run(
+        f"sudo smartctl -i {device}", shell=True, capture_output=True, text=True
+    ).stdout
+    model = serial = "unknown"
+    for line in output.splitlines():
+        if any(x in line for x in ["Device Model", "Model Family", "Model Number"]):
+            model = line.split(":", 1)[-1].strip().replace(" ", "_")
+        elif "Serial Number" in line:
+            serial = line.split(":", 1)[-1].strip()
+    if model == "unknown" or serial == "unknown":
+        print("Failed to extract model or serial. Using timestamp for filename.")
+        return f"drive_test_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
+    return f"{model}_{serial}_test.log"
+
+
+def show_smart_info(device, log_file):
+    print("Gathering extended SMART info...")
+    cmd = f"sudo smartctl -x {'-d nvme' if is_nvme(device) else ''} {device}"
+    run_cmd(cmd, log_file)
+
+
+def run_smart_test(device, log_file):
+    print("Running short SMART self-test (2 minutes)...")
+    cmd = f"sudo smartctl -t short {'-d nvme' if is_nvme(device) else ''} {device}"
+    run_cmd(cmd, log_file, check=False)
+    time.sleep(120)
+    print("SMART Self-Test Results:")
+    cmd = f"sudo smartctl -l selftest {'-d nvme' if is_nvme(device) else ''} {device}"
+    run_cmd(cmd, log_file)
+
+
+def check_partition_table(device, log_file):
+    print("Checking for existing partition table...")
+    output = run_cmd(f"sudo parted -s {device} print", log_file, check=False)
+    return not (
+        "Partition Table: unknown" in output or "unrecognised disk label" in output
+    )
+
+
+def run_badblocks(device, log_file):
+    print("Running non-destructive read-only badblocks check (this may take time)...")
+    run_cmd(f"sudo badblocks -sv {device}", log_file, check=False)
+    print("Badblocks check complete.")
+
+
+def fill_random(device, log_file):
+    print("Overwriting disk with random data...")
+    run_cmd(
+        f"sudo dd if=/dev/urandom of={device} bs=1M status=progress",
+        log_file,
+        check=False,
+    )
+    print("Random data written to disk.")
+
+
+def run_fio(device, rw_mode, log_file, bs="4k", runtime=30):
+    print(f"Running fio test: {rw_mode}")
+    cmd = (
+        f"sudo fio --name={rw_mode}_test "
+        f"--filename={device} --direct=1 --rw={rw_mode} "
+        f"--bs={bs} --iodepth=16 --numjobs=1 "
+        f"--time_based --runtime={runtime} --group_reporting"
+    )
+    run_cmd(cmd, log_file)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test used drive with SMART, badblocks, and fio."
+    )
+    parser.add_argument("device", help="Target device (e.g., /dev/sdX)")
+    args = parser.parse_args()
+    device = args.device
+
+    if not os.path.exists(device):
+        print(f"Device {device} not found.")
+        sys.exit(1)
+
+    log_filename = get_device_info(device)
+    with open(log_filename, "w") as log_file:
+        print(f"Logging to {log_filename}...")
+
+        show_smart_info(device, log_file)
+        run_smart_test(device, log_file)
+        run_badblocks(device, log_file)
+
+        if check_partition_table(device, log_file):
+            prompt_overwrite()
+            fill_random(device, log_file)
+        else:
+            fill_random(device, log_file)
+
+        run_fio(device, "write", log_file)
+        run_fio(device, "read", log_file)
+        run_fio(device, "randread", log_file)
+        run_fio(device, "randwrite", log_file)
+        run_fio(device, "randread", log_file, bs="512", runtime=10)  # IOPS test
+
+        print(f"\nTest completed. Results saved to {log_filename}")
+
+
+if __name__ == "__main__":
+    main()