120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
import os
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import psutil
|
|
from ._protocols import MountInfo, DiskInfo
|
|
|
|
class SystemAPI:
|
|
|
|
def query_mounts(self) -> list[MountInfo]:
|
|
mounts = []
|
|
for partition in psutil.disk_partitions(all=False):
|
|
mount_info = MountInfo(device=partition.device, mount_point=partition.mountpoint, fs_type=partition.fstype, options=partition.opts)
|
|
mounts.append(mount_info)
|
|
return mounts
|
|
|
|
def query_nvmes(self) -> list[DiskInfo]:
|
|
disks = []
|
|
try:
|
|
result = subprocess.run(['lsblk', '-ndo', 'NAME,MODEL,SIZE,SERIAL', '-b'], capture_output=True, text=True, check=False)
|
|
if result.returncode == 0:
|
|
for line in result.stdout.strip().split('\n'):
|
|
if not line.strip():
|
|
continue
|
|
parts = line.split(maxsplit=3)
|
|
if len(parts) >= 3:
|
|
device = f'/dev/{parts[0]}'
|
|
model = parts[1] if len(parts) > 1 else 'Unknown'
|
|
size_str = parts[2] if len(parts) > 2 else '0'
|
|
serial = parts[3] if len(parts) > 3 else 'Unknown'
|
|
try:
|
|
size = int(size_str)
|
|
except ValueError:
|
|
size = 0
|
|
disk_info = DiskInfo(device=device, model=model, size=size, serial=serial)
|
|
disks.append(disk_info)
|
|
except FileNotFoundError:
|
|
pass
|
|
if not disks:
|
|
disks = self._query_disks_fallback()
|
|
return disks
|
|
|
|
def _query_disks_fallback(self) -> list[DiskInfo]:
|
|
disks = []
|
|
seen_devices = set()
|
|
for partition in psutil.disk_partitions(all=True):
|
|
device = partition.device
|
|
if not device.startswith('/dev/'):
|
|
continue
|
|
base_device = self._get_base_device(device)
|
|
if base_device in seen_devices:
|
|
continue
|
|
seen_devices.add(base_device)
|
|
try:
|
|
usage = psutil.disk_usage(partition.mountpoint)
|
|
size = usage.total
|
|
except (PermissionError, OSError):
|
|
size = 0
|
|
disk_info = DiskInfo(device=base_device, model='Unknown', size=size, serial='Unknown')
|
|
disks.append(disk_info)
|
|
return disks
|
|
|
|
def _get_base_device(self, device: str) -> str:
|
|
if 'nvme' in device:
|
|
if 'p' in device:
|
|
return device.rsplit('p', 1)[0]
|
|
return device
|
|
import re
|
|
match = re.match('(/dev/[a-z]+)', device)
|
|
if match:
|
|
return match.group(1)
|
|
return device
|
|
|
|
def get_disk_for_path(self, path: Path) -> Optional[str]:
|
|
path = path.resolve()
|
|
best_match = None
|
|
best_match_len = 0
|
|
for partition in psutil.disk_partitions():
|
|
mount_point = Path(partition.mountpoint)
|
|
try:
|
|
if path == mount_point or mount_point in path.parents:
|
|
mount_len = len(str(mount_point))
|
|
if mount_len > best_match_len:
|
|
best_match = partition.device
|
|
best_match_len = mount_len
|
|
except (ValueError, OSError):
|
|
continue
|
|
return best_match
|
|
|
|
def get_disk_usage(self, path: Path) -> tuple[int, int, int]:
|
|
try:
|
|
usage = psutil.disk_usage(str(path))
|
|
return (usage.total, usage.used, usage.free)
|
|
except (PermissionError, OSError):
|
|
return (0, 0, 0)
|
|
|
|
def get_mount_point(self, path: Path) -> Optional[Path]:
|
|
path = path.resolve()
|
|
best_match = None
|
|
best_match_len = 0
|
|
for partition in psutil.disk_partitions():
|
|
mount_point = Path(partition.mountpoint)
|
|
try:
|
|
if path == mount_point or mount_point in path.parents:
|
|
mount_len = len(str(mount_point))
|
|
if mount_len > best_match_len:
|
|
best_match = mount_point
|
|
best_match_len = mount_len
|
|
except (ValueError, OSError):
|
|
continue
|
|
return best_match
|
|
|
|
def is_same_filesystem(self, path1: Path, path2: Path) -> bool:
|
|
try:
|
|
stat1 = path1.stat()
|
|
stat2 = path2.stat()
|
|
return stat1.st_dev == stat2.st_dev
|
|
except (OSError, PermissionError):
|
|
return False
|