""" 磁盘管理模块 提供分区检测、挂载、卸载等功能 """ import json import os import subprocess import tempfile from typing import List, Dict, Optional, Tuple from dataclasses import dataclass from utils import run_command, format_size @dataclass class Partition: """分区信息""" device: str # 设备名,如 /dev/sda1 mount_point: str # 挂载点,如 / fstype: str # 文件系统类型,如 ext4 size: int # 总大小(字节) used: int # 已用大小(字节) available: int # 可用大小(字节) use_percent: float # 使用百分比 label: str = "" # 卷标 is_system: bool = False # 是否为系统盘 @property def size_formatted(self) -> str: return format_size(self.size) @property def used_formatted(self) -> str: return format_size(self.used) @property def available_formatted(self) -> str: return format_size(self.available) @property def use_percent_formatted(self) -> str: return f"{self.use_percent:.1f}%" class DiskManager: """磁盘管理器""" # 排除的文件系统类型 EXCLUDED_FS_TYPES = { 'swap', 'vfat', # 通常是 EFI 分区 'iso9660', # 光盘镜像 'squashfs', # LiveCD 文件系统 'tmpfs', # 临时文件系统 'devtmpfs', # 设备文件系统 'sysfs', # 内核虚拟文件系统 'proc', # 进程虚拟文件系统 'devpts', # 设备终端 'securityfs', 'cgroup', 'pstore', 'debugfs', 'hugetlbfs', 'mqueue', 'fusectl', 'configfs', 'overlay', # Overlay 文件系统 } # 排除的挂载点前缀 EXCLUDED_MOUNT_PREFIXES = [ '/boot/efi', '/efi', '/boot/EFI', '/snap', '/cdrom', '/media/cdrom', ] def __init__(self, sudo_password: Optional[str] = None): self.sudo_password = sudo_password self.mounted_partitions: Dict[str, str] = {} # device -> mount_point def get_partitions(self) -> List[Partition]: """ 获取所有可用分区(排除 EFI、swap、LiveCD USB 等) """ partitions = [] # 获取 LiveCD 设备 livecd_device = self._detect_livecd_device() # 使用 lsblk 获取分区信息 stdout, stderr, returncode = run_command( ['lsblk', '-J', '-o', 'NAME,SIZE,TYPE,MOUNTPOINT,FSTYPE,LABEL,PKNAME'], use_sudo=True, sudo_password=self.sudo_password ) if returncode != 0: # 备用方案:解析 /proc/partitions return self._get_partitions_from_proc() try: data = json.loads(stdout) devices = data.get('blockdevices', []) for device in devices: self._parse_device(device, partitions, livecd_device) except (json.JSONDecodeError, KeyError) as e: print(f"解析 lsblk 输出失败: {e}") return self._get_partitions_from_proc() # 获取每个分区的使用情况 self._fill_usage_info(partitions) # 按使用率排序(高的在前) partitions.sort(key=lambda p: p.use_percent, reverse=True) return partitions def _parse_device( self, device: dict, partitions: List[Partition], livecd_device: Optional[str] ): """递归解析设备信息""" device_name = device.get('name', '') device_type = device.get('type', '') fstype = device.get('fstype', '') mount_point = device.get('mountpoint', '') label = device.get('label', '') size_str = device.get('size', '0') # 只处理分区类型 if device_type == 'part': device_path = f"/dev/{device_name}" # 排除条件 if self._should_exclude(device_path, fstype, mount_point, livecd_device): return # 解析大小 size = self._parse_size(size_str) partition = Partition( device=device_path, mount_point=mount_point or "", fstype=fstype or "unknown", size=size, used=0, available=0, use_percent=0.0, label=label or "", is_system=(mount_point == '/') ) partitions.append(partition) # 递归处理子设备 children = device.get('children', []) for child in children: self._parse_device(child, partitions, livecd_device) def _should_exclude( self, device: str, fstype: str, mount_point: str, livecd_device: Optional[str] ) -> bool: """判断是否应该排除此分区""" # 排除特定文件系统类型 if fstype in self.EXCLUDED_FS_TYPES: return True # 排除 EFI 分区 if mount_point and any(mount_point.startswith(prefix) for prefix in self.EXCLUDED_MOUNT_PREFIXES): return True # 排除 LiveCD USB 设备 if livecd_device and device.startswith(livecd_device): return True # 排除没有文件系统的设备 if not fstype: return True return False def _parse_size(self, size_str: str) -> int: """解析 lsblk 输出的大小字符串""" if not size_str: return 0 size_str = size_str.strip() # 提取数字和单位 try: if size_str.endswith('T'): return int(float(size_str[:-1]) * 1024 * 1024 * 1024 * 1024) elif size_str.endswith('G'): return int(float(size_str[:-1]) * 1024 * 1024 * 1024) elif size_str.endswith('M'): return int(float(size_str[:-1]) * 1024 * 1024) elif size_str.endswith('K'): return int(float(size_str[:-1]) * 1024) else: return int(size_str) except ValueError: return 0 def _fill_usage_info(self, partitions: List[Partition]): """填充分区的使用情况信息""" # 使用 df 命令获取使用情况 stdout, stderr, returncode = run_command( ['df', '-B1'], use_sudo=True, sudo_password=self.sudo_password ) if returncode != 0: return # 解析 df 输出 df_info = {} for line in stdout.strip().split('\n')[1:]: # 跳过标题行 parts = line.split() if len(parts) >= 6: device = parts[0] try: size = int(parts[1]) used = int(parts[2]) available = int(parts[3]) use_percent = float(parts[4].rstrip('%')) df_info[device] = { 'size': size, 'used': used, 'available': available, 'use_percent': use_percent } except (ValueError, IndexError): pass # 填充分区信息 for partition in partitions: if partition.device in df_info: info = df_info[partition.device] partition.size = info['size'] partition.used = info['used'] partition.available = info['available'] partition.use_percent = info['use_percent'] def _detect_livecd_device(self) -> Optional[str]: """ 检测 LiveCD USB 设备 返回设备名前缀,如 /dev/sdb """ # 方法1: 检查 /proc/cmdline try: with open('/proc/cmdline', 'r') as f: cmdline = f.read() # 查找 boot= 参数 for param in cmdline.split(): if param.startswith('boot='): boot_device = param.split('=', 1)[1] # 提取设备名(去掉分区号) if boot_device.startswith('/dev/'): return boot_device.rstrip('0123456789') except Exception: pass # 方法2: 检查根文件系统类型 try: stdout, _, _ = run_command(['findmnt', '-n', '-o', 'SOURCE', '/']) root_device = stdout.strip() if root_device: # 如果根文件系统是 overlay 或 tmpfs,很可能是 LiveCD stdout2, _, _ = run_command(['findmnt', '-n', '-o', 'FSTYPE', '/']) fstype = stdout2.strip() if fstype in ('overlay', 'squashfs', 'tmpfs'): # 查找 USB 设备 return self._find_usb_device() except Exception: pass # 方法3: 查找标记为 LiveCD 的设备 return self._find_usb_device() def _find_usb_device(self) -> Optional[str]: """查找 USB 设备""" try: stdout, _, _ = run_command( ['lsblk', '-J', '-o', 'NAME,TRAN,TYPE'], use_sudo=True, sudo_password=self.sudo_password ) data = json.loads(stdout) devices = data.get('blockdevices', []) for device in devices: if device.get('tran') == 'usb' and device.get('type') == 'disk': return f"/dev/{device['name']}" except Exception: pass return None def _get_partitions_from_proc(self) -> List[Partition]: """从 /proc/partitions 获取分区信息(备用方案)""" partitions = [] try: stdout, _, returncode = run_command( ['cat', '/proc/partitions'], use_sudo=True, sudo_password=self.sudo_password ) if returncode != 0: return [] livecd_device = self._detect_livecd_device() for line in stdout.strip().split('\n')[2:]: # 跳过标题行 parts = line.split() if len(parts) >= 4: name = parts[3] # 只处理分区(名称以数字结尾) if name[-1].isdigit(): device = f"/dev/{name}" # 排除 LiveCD 设备 if livecd_device and device.startswith(livecd_device): continue # 获取文件系统类型 fstype = self._get_fstype(device) # 排除不需要的文件系统类型 if fstype in self.EXCLUDED_FS_TYPES or not fstype: continue # 获取挂载点 mount_point = self._get_mount_point(device) # 排除 EFI 分区 if mount_point and any( mount_point.startswith(prefix) for prefix in self.EXCLUDED_MOUNT_PREFIXES ): continue size = int(parts[2]) * 1024 # 转换为字节 partition = Partition( device=device, mount_point=mount_point or "", fstype=fstype, size=size, used=0, available=0, use_percent=0.0, is_system=(mount_point == '/') ) partitions.append(partition) except Exception as e: print(f"从 /proc/partitions 获取信息失败: {e}") # 填充使用情况 self._fill_usage_info(partitions) return partitions def _get_fstype(self, device: str) -> str: """获取设备的文件系统类型""" stdout, _, _ = run_command( ['blkid', '-o', 'value', '-s', 'TYPE', device], use_sudo=True, sudo_password=self.sudo_password ) return stdout.strip() def _get_mount_point(self, device: str) -> str: """获取设备的挂载点""" stdout, _, _ = run_command( ['findmnt', '-n', '-o', 'TARGET', device], use_sudo=True, sudo_password=self.sudo_password ) return stdout.strip() def mount_partition(self, partition: Partition, mount_point: Optional[str] = None) -> Tuple[bool, str]: """ 挂载分区到临时目录 参数: partition: 分区信息 mount_point: 指定挂载点(可选) 返回: (success, mount_point_or_error_message) """ # 如果已经挂载,返回现有挂载点 if partition.mount_point: return True, partition.mount_point # 如果已挂载到临时目录 if partition.device in self.mounted_partitions: return True, self.mounted_partitions[partition.device] # 创建临时挂载点 if mount_point is None: mount_point = tempfile.mkdtemp(prefix='log_cleaner_') # 挂载分区 stdout, stderr, returncode = run_command( ['mount', partition.device, mount_point], use_sudo=True, sudo_password=self.sudo_password ) if returncode == 0: partition.mount_point = mount_point self.mounted_partitions[partition.device] = mount_point return True, mount_point else: return False, f"挂载失败: {stderr}" def unmount_partition(self, partition: Partition) -> Tuple[bool, str]: """ 卸载分区 参数: partition: 分区信息 返回: (success, message) """ # 如果是系统挂载点,不卸载 if partition.is_system: return False, "不能卸载系统分区" mount_point = partition.mount_point if not mount_point: return False, "分区未挂载" # 如果是临时挂载点,卸载 if partition.device in self.mounted_partitions: stdout, stderr, returncode = run_command( ['umount', mount_point], use_sudo=True, sudo_password=self.sudo_password ) if returncode == 0: del self.mounted_partitions[partition.device] partition.mount_point = "" # 删除临时目录 try: os.rmdir(mount_point) except OSError: pass return True, "卸载成功" else: return False, f"卸载失败: {stderr}" return False, "分区不是由本工具挂载的" def unmount_all(self): """卸载所有由本工具挂载的分区""" for device in list(self.mounted_partitions.keys()): mount_point = self.mounted_partitions[device] run_command( ['umount', mount_point], use_sudo=True, sudo_password=self.sudo_password ) try: os.rmdir(mount_point) except OSError: pass self.mounted_partitions.clear() def get_partition_display_name(self, partition: Partition) -> str: """获取分区的显示名称""" parts = [partition.device] if partition.mount_point: parts.append(f"({partition.mount_point})") elif partition.label: parts.append(f"({partition.label})") if partition.is_system: parts.append("[系统盘]") return " ".join(parts)