327 lines
10 KiB
Python
327 lines
10 KiB
Python
import subprocess
|
|
import logging
|
|
import os
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Import network logging for critical disk operations
|
|
from .network_logging import log_to_discord
|
|
|
|
|
|
def log_disk_operation(operation, status="start", details=""):
|
|
log_to_discord("INFO", f"DISK_{operation}_{status}: {details}", module="disk")
|
|
|
|
|
|
class CommandResult:
|
|
def __init__(self, stdout, stderr, returncode):
|
|
self.stdout = stdout
|
|
self.stderr = stderr
|
|
self.returncode = returncode
|
|
|
|
|
|
def run_command(cmd, check=True):
|
|
logger.info(f"Running command: {' '.join(cmd)}")
|
|
|
|
process = subprocess.Popen(
|
|
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1
|
|
)
|
|
|
|
stdout_lines = []
|
|
stderr_lines = []
|
|
|
|
def read_stream(stream, line_list, log_level):
|
|
for line in stream:
|
|
line_clean = line.strip()
|
|
if line_clean:
|
|
log_level(line_clean)
|
|
line_list.append(line)
|
|
|
|
import threading
|
|
|
|
t1 = threading.Thread(
|
|
target=read_stream, args=(process.stdout, stdout_lines, logger.info)
|
|
)
|
|
t2 = threading.Thread(
|
|
target=read_stream, args=(process.stderr, stderr_lines, logger.error)
|
|
)
|
|
|
|
t1.start()
|
|
t2.start()
|
|
|
|
t1.join()
|
|
t2.join()
|
|
|
|
returncode = process.wait()
|
|
|
|
stdout_str = "".join(stdout_lines)
|
|
stderr_str = "".join(stderr_lines)
|
|
|
|
if check and returncode != 0:
|
|
raise subprocess.CalledProcessError(
|
|
returncode, cmd, output=stdout_str, stderr=stderr_str
|
|
)
|
|
|
|
return CommandResult(stdout_str, stderr_str, returncode)
|
|
|
|
|
|
def get_partition_device(disk_device, partition_number):
|
|
"""
|
|
Returns the partition device path.
|
|
Handles NVMe style (p1) vs sd style (1).
|
|
"""
|
|
if disk_device[-1].isdigit():
|
|
return f"{disk_device}p{partition_number}"
|
|
return f"{disk_device}{partition_number}"
|
|
|
|
|
|
def get_total_memory():
|
|
"""Returns total system memory in bytes."""
|
|
try:
|
|
return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
|
|
except (ValueError, OSError):
|
|
with open("/proc/meminfo", "r") as f:
|
|
for line in f:
|
|
if line.startswith("MemTotal:"):
|
|
parts = line.split()
|
|
return int(parts[1]) * 1024
|
|
return 0
|
|
|
|
|
|
def get_disk_size(disk_device):
|
|
"""Returns disk size in bytes using blockdev."""
|
|
try:
|
|
res = run_command(["blockdev", "--getsize64", disk_device])
|
|
return int(res.stdout.strip())
|
|
except Exception as e:
|
|
logger.error(f"Failed to get disk size: {e}")
|
|
return 0
|
|
|
|
|
|
def ensure_unmounted(disk_device):
|
|
"""
|
|
Finds all mounted partitions belonging to disk_device and unmounts them.
|
|
"""
|
|
logger.info(f"Ensuring all partitions on {disk_device} are unmounted...")
|
|
try:
|
|
# Get all partitions for this disk from lsblk
|
|
res = run_command(["lsblk", "-nlo", "NAME,MOUNTPOINT", disk_device])
|
|
for line in res.stdout.splitlines():
|
|
parts = line.split()
|
|
if len(parts) >= 2:
|
|
# NAME is usually the leaf name (e.g. vda1), we need full path
|
|
# But lsblk -lp gives full paths
|
|
pass
|
|
|
|
# Simpler approach: findmnt or just look at /proc/mounts
|
|
# Let's use lsblk -lp to get full paths
|
|
res = run_command(["lsblk", "-lpno", "NAME,MOUNTPOINT", disk_device])
|
|
for line in res.stdout.splitlines():
|
|
# line looks like "/dev/vda1 /mnt/boot" or "/dev/vda1"
|
|
parts = line.split()
|
|
if len(parts) >= 2:
|
|
dev_path = parts[0]
|
|
mount_point = parts[1]
|
|
if mount_point and mount_point != "None":
|
|
logger.info(f"Unmounting {dev_path} from {mount_point}...")
|
|
run_command(["umount", "-l", dev_path])
|
|
except Exception as e:
|
|
logger.warning(f"Error during pre-partition unmount: {e}")
|
|
|
|
|
|
def is_uefi():
|
|
"""Checks if the system is booted in UEFI mode."""
|
|
return os.path.exists("/sys/firmware/efi")
|
|
|
|
|
|
def auto_partition_disk(disk_device):
|
|
"""
|
|
Automatically partitions the disk based on boot mode:
|
|
UEFI (GPT): ESP (1GB), Root (Remaining), Swap (RAM+2GB)
|
|
BIOS (MBR): Root (Remaining, Bootable), Swap (RAM+2GB)
|
|
"""
|
|
uefi = is_uefi()
|
|
logger.info(f"Starting auto-partitioning on {disk_device} (UEFI: {uefi})")
|
|
log_disk_operation("PARTITION", "start", f"Disk: {disk_device}, UEFI: {uefi}")
|
|
|
|
ensure_unmounted(disk_device)
|
|
|
|
disk_size = get_disk_size(disk_device)
|
|
ram_size = get_total_memory()
|
|
|
|
# Calculate sizes
|
|
swap_mb = int((ram_size / (1024 * 1024)) + 2048)
|
|
# Cap swap at 16GB for sanity on large RAM systems unless disk is huge
|
|
if disk_size < 100 * 1024**3: # < 100GB
|
|
swap_mb = min(swap_mb, 4096)
|
|
|
|
disk_mb = disk_size / (1024 * 1024)
|
|
use_swap = disk_mb > (10240 + swap_mb) # Only use swap if we have at least 10GB for root
|
|
|
|
if uefi:
|
|
# GPT Layout
|
|
logger.info("Using GPT layout for UEFI")
|
|
run_command(["sgdisk", "-Z", disk_device])
|
|
run_command(["sgdisk", "-o", disk_device])
|
|
|
|
# 1. ESP (1GB)
|
|
run_command(["sgdisk", "-n", "1:0:+1024M", "-t", "1:ef00", "-c", "1:EFI System", disk_device])
|
|
|
|
# 2. Swap (if enabled)
|
|
if use_swap:
|
|
run_command(["sgdisk", "-n", f"3:-{swap_mb}M:0", "-t", "3:8200", "-c", "3:Swap", disk_device])
|
|
|
|
# 3. Root
|
|
run_command(["sgdisk", "-n", "2:0:0", "-t", "2:8300", "-c", "2:Root", disk_device])
|
|
else:
|
|
# MBR Layout for BIOS
|
|
logger.info("Using MBR layout for BIOS")
|
|
# Wipe disk using dd to be sure MBR is cleared
|
|
run_command(["dd", "if=/dev/zero", f"of={disk_device}", "bs=512", "count=100"])
|
|
|
|
# Use parted for MBR
|
|
run_command(["parted", "-s", disk_device, "mklabel", "msdos"])
|
|
|
|
if use_swap:
|
|
# Root first, then swap at end
|
|
root_end = int(disk_mb - swap_mb)
|
|
run_command(["parted", "-s", disk_device, "mkpart", "primary", "ext4", "1MiB", f"{root_end}MiB"])
|
|
run_command(["parted", "-s", disk_device, "mkpart", "primary", "linux-swap", f"{root_end}MiB", "100%"])
|
|
run_command(["parted", "-s", disk_device, "set", "1", "boot", "on"])
|
|
else:
|
|
run_command(["parted", "-s", disk_device, "mkpart", "primary", "ext4", "1MiB", "100%"])
|
|
run_command(["parted", "-s", disk_device, "set", "1", "boot", "on"])
|
|
|
|
run_command(["partprobe", disk_device])
|
|
import time
|
|
time.sleep(2)
|
|
|
|
# Identify partitions
|
|
if uefi:
|
|
efi_part = get_partition_device(disk_device, 1)
|
|
root_part = get_partition_device(disk_device, 2)
|
|
swap_part = get_partition_device(disk_device, 3) if use_swap else None
|
|
else:
|
|
efi_part = None # BIOS doesn't use ESP
|
|
root_part = get_partition_device(disk_device, 1)
|
|
swap_part = get_partition_device(disk_device, 2) if use_swap else None
|
|
|
|
# Format
|
|
if uefi:
|
|
logger.info(f"Formatting EFI partition {efi_part}...")
|
|
run_command(["mkfs.vfat", "-F32", efi_part])
|
|
|
|
if use_swap:
|
|
logger.info(f"Formatting Swap partition {swap_part}...")
|
|
run_command(["mkswap", swap_part])
|
|
|
|
logger.info(f"Formatting Root partition {root_part}...")
|
|
run_command(["mkfs.ext4", "-F", root_part])
|
|
|
|
result = {"efi": efi_part, "root": root_part, "swap": swap_part}
|
|
log_disk_operation("PARTITION", "complete", str(result))
|
|
return result
|
|
|
|
|
|
def mount_partitions(partition_info, mount_root="/mnt"):
|
|
"""
|
|
Mounts the partitions into mount_root.
|
|
"""
|
|
import os
|
|
log_disk_operation("MOUNT", "start", f"Mounting to {mount_root}")
|
|
|
|
# 1. Mount Root
|
|
if not os.path.exists(mount_root):
|
|
os.makedirs(mount_root)
|
|
run_command(["mount", partition_info["root"], mount_root])
|
|
|
|
# 2. Mount EFI (if exists)
|
|
if partition_info.get("efi"):
|
|
efi_mount = os.path.join(mount_root, "boot/efi")
|
|
os.makedirs(efi_mount, exist_ok=True)
|
|
run_command(["mount", partition_info["efi"], efi_mount])
|
|
else:
|
|
# For BIOS, /boot is just a directory on root or we could make a separate /boot
|
|
# Current logic keeps it on root.
|
|
pass
|
|
|
|
# 3. Enable Swap
|
|
if partition_info.get("swap"):
|
|
run_command(["swapon", partition_info["swap"]])
|
|
|
|
logger.info(f"Partitions mounted at {mount_root}")
|
|
log_disk_operation("MOUNT", "complete", f"Root: {partition_info['root']}")
|
|
|
|
|
|
def create_partition(
|
|
disk_device, size_mb, type_code="8300", name="Linux filesystem", fstype=None
|
|
):
|
|
"""
|
|
Creates a new partition on the disk and formats it.
|
|
type_code: ef00 (EFI), 8200 (Swap), 8300 (Linux)
|
|
fstype: ext4, fat32, swap
|
|
"""
|
|
# Find next available partition number
|
|
res = run_command(["sgdisk", "-p", disk_device])
|
|
# Very basic parsing to find the next number
|
|
existing_nums = []
|
|
for line in res.stdout.splitlines():
|
|
parts = line.split()
|
|
if parts and parts[0].isdigit():
|
|
existing_nums.append(int(parts[0]))
|
|
|
|
next_num = 1
|
|
while next_num in existing_nums:
|
|
next_num += 1
|
|
|
|
# Use -g (mbrtogpt) to ensure we can work on the disk if it was MBR
|
|
# size_mb=0 means use all available space in the gap
|
|
size_spec = f"+{size_mb}M" if size_mb > 0 else "0"
|
|
|
|
run_command(
|
|
[
|
|
"sgdisk",
|
|
"-g",
|
|
"-n",
|
|
f"{next_num}:0:{size_spec}",
|
|
"-t",
|
|
f"{next_num}:{type_code}",
|
|
"-c",
|
|
f"{next_num}:{name}",
|
|
disk_device,
|
|
]
|
|
)
|
|
|
|
run_command(["partprobe", disk_device])
|
|
|
|
# Wait for partition node to appear
|
|
try:
|
|
run_command(["udevadm", "settle", "--timeout=5"])
|
|
except Exception:
|
|
import time
|
|
|
|
time.sleep(2)
|
|
|
|
part_dev = get_partition_device(disk_device, next_num)
|
|
|
|
if fstype == "fat32":
|
|
run_command(["mkfs.vfat", "-F32", part_dev])
|
|
elif fstype == "ext4":
|
|
run_command(["mkfs.ext4", "-F", part_dev])
|
|
elif fstype == "swap":
|
|
run_command(["mkswap", part_dev])
|
|
|
|
return next_num
|
|
|
|
|
|
def delete_partition(disk_device, part_num):
|
|
"""Deletes a partition by number."""
|
|
run_command(["sgdisk", "-d", str(part_num), disk_device])
|
|
run_command(["partprobe", disk_device])
|
|
|
|
|
|
def wipe_disk(disk_device):
|
|
"""Zaps the disk and creates a new GPT table."""
|
|
run_command(["sgdisk", "-Z", disk_device])
|
|
run_command(["sgdisk", "-o", disk_device])
|
|
run_command(["partprobe", disk_device])
|