146 lines
No EOL
3.8 KiB
Python
146 lines
No EOL
3.8 KiB
Python
import psutil
|
|
import subprocess
|
|
|
|
def get_drive_health(device):
|
|
"""
|
|
Get the health of a storage device using smartctl.
|
|
Returns a score based on the drive's SMART health status.
|
|
"""
|
|
try:
|
|
result = subprocess.run(['smartctl', '-H', device], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
output = result.stdout.decode()
|
|
if "PASSED" in output:
|
|
return 10 # Healthy drive
|
|
elif "FAILED" in output:
|
|
return 0 # Failed drive
|
|
else:
|
|
return 5 # Unknown or problematic drive
|
|
except Exception as e:
|
|
print(f"Error checking SMART status for {device}: {e}")
|
|
return 5 # Default score for uncheckable devices
|
|
|
|
def get_drive_type(device):
|
|
"""
|
|
Determine if a device is an SSD or HDD based on its device type.
|
|
"""
|
|
if 'NVME' in device:
|
|
return 15 # SSDs are optimal
|
|
elif 'SSD' in device:
|
|
return 10 # SSDs are optimal
|
|
else:
|
|
return 5 # HDDs are less optimal for boot drives
|
|
|
|
def get_drive_size(device):
|
|
"""
|
|
Get size of a block device using lsblk (works for disks).
|
|
Always returns an integer score.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["lsblk", "-dn", "-o", "SIZE", device],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
check=True
|
|
)
|
|
|
|
size_str = result.stdout.strip().upper()
|
|
|
|
if not size_str:
|
|
return 5 # fallback
|
|
|
|
size_str = size_str.replace(",", ".")
|
|
|
|
# Convert to GB
|
|
if size_str.endswith("T"):
|
|
size_gb = float(size_str[:-1]) * 1024
|
|
elif size_str.endswith("G"):
|
|
size_gb = float(size_str[:-1])
|
|
elif size_str.endswith("M"):
|
|
size_gb = float(size_str[:-1]) / 1024
|
|
else:
|
|
return 5 # unknown format
|
|
|
|
if size_gb < 128:
|
|
return 5
|
|
elif size_gb < 512:
|
|
return 7
|
|
else:
|
|
return 10
|
|
|
|
except Exception as e:
|
|
print(f"Error getting size for {device}: {e}")
|
|
return 5 # ALWAYS return something
|
|
|
|
def get_device_score(device):
|
|
"""
|
|
Calculate a suitability score for each drive based on:
|
|
- Type (SSD/HDD)
|
|
- Health (SMART status)
|
|
- Size (GB)
|
|
"""
|
|
score = 0
|
|
|
|
# Type
|
|
score += get_drive_type(device)
|
|
|
|
# Health
|
|
score += get_drive_health(device)
|
|
|
|
# Size
|
|
score += get_drive_size(device)
|
|
|
|
return score
|
|
|
|
|
|
def list_storage_devices():
|
|
"""
|
|
List all physical storage devices (e.g. /dev/sda, /dev/nvme0n1)
|
|
and return them with their computed scores.
|
|
|
|
Compatible with the existing scoring pipeline.
|
|
"""
|
|
devices = []
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["lsblk", "-dn", "-o", "NAME"], # -d = disks only, no partitions
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
check=True
|
|
)
|
|
|
|
for line in result.stdout.strip().split("\n"):
|
|
if not line:
|
|
continue
|
|
|
|
device = f"/dev/{line.strip()}"
|
|
score = get_device_score(device) # <-- reuses your existing logic
|
|
devices.append((device, score))
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error listing devices: {e}")
|
|
|
|
return devices
|
|
|
|
def main():
|
|
devices = list_storage_devices()
|
|
print(devices)
|
|
|
|
if not devices:
|
|
print("No storage devices found.")
|
|
return
|
|
|
|
print(f"\n{'Device':<20} {'Score'}")
|
|
print("-" * 30)
|
|
|
|
for device, score in devices:
|
|
print(f"{device:<20} {score}")
|
|
|
|
# Find the highest scoring drive
|
|
best_device = max(devices, key=lambda x: x[1])
|
|
print(f"\nBest drive for boot: {best_device[0]} with score: {best_device[1]}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |