diff --git a/PythonScripts/test_disk.py b/PythonScripts/test_disk.py new file mode 100644 index 0000000..11c3190 --- /dev/null +++ b/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()