261 lines
7.8 KiB
Python
261 lines
7.8 KiB
Python
import subprocess
|
|
import logging
|
|
import os
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
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 auto_partition_disk(disk_device):
|
|
"""
|
|
Automatically partitions the disk:
|
|
1. EFI System Partition (2GB standard, 1GB small)
|
|
2. Root (Remaining)
|
|
3. Swap (RAM + 2GB, or 0 if small)
|
|
|
|
Layout: P1=EFI, P3=Swap (End), P2=Root (Middle)
|
|
"""
|
|
logger.info(f"Starting auto-partitioning on {disk_device}")
|
|
|
|
# Calculate sizes
|
|
disk_size = get_disk_size(disk_device)
|
|
ram_size = get_total_memory()
|
|
|
|
# Defaults
|
|
efi_mb = 2048
|
|
swap_mb = int((ram_size / (1024*1024)) + 2048)
|
|
min_root_mb = 10240 # 10GB
|
|
|
|
total_required_mb = efi_mb + swap_mb + min_root_mb
|
|
disk_mb = disk_size / (1024*1024)
|
|
|
|
use_swap = True
|
|
|
|
if disk_mb < total_required_mb:
|
|
logger.warning("Disk too small for standard layout. Adjusting...")
|
|
efi_mb = 1024
|
|
use_swap = False
|
|
|
|
# Check minimal viability
|
|
if disk_mb < (efi_mb + min_root_mb):
|
|
raise Exception("Disk too small for installation (Need ~11GB)")
|
|
|
|
# 1. Zap the disk (destroy all data)
|
|
run_command(["sgdisk", "-Z", disk_device])
|
|
|
|
# 2. Create new GPT table
|
|
run_command(["sgdisk", "-o", disk_device])
|
|
|
|
# 3. Create EFI Partition (Part 1, Start)
|
|
run_command(["sgdisk", "-n", f"1:0:+{efi_mb}M", "-t", "1:ef00", "-c", "1:EFI System", disk_device])
|
|
|
|
# 4. Create Swap Partition (Part 3, End) - If enabled
|
|
if use_swap:
|
|
# sgdisk negative start is from end of disk
|
|
# We use partition 3 for Swap
|
|
run_command(["sgdisk", "-n", f"3:-{swap_mb}M:0", "-t", "3:8200", "-c", "3:Swap", disk_device])
|
|
|
|
# 5. Create Root Partition (Part 2, Fill Gap)
|
|
# This fills the space between P1 and P3 (or end if no swap)
|
|
run_command(["sgdisk", "-n", "2:0:0", "-t", "2:8300", "-c", "2:Root", disk_device])
|
|
|
|
# Inform kernel of changes
|
|
run_command(["partprobe", disk_device])
|
|
|
|
import time
|
|
time.sleep(1)
|
|
|
|
# 6. Format Partitions
|
|
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
|
|
|
|
logger.info("Formatting EFI partition...")
|
|
run_command(["mkfs.vfat", "-F32", efi_part])
|
|
|
|
if use_swap:
|
|
logger.info("Formatting Swap partition...")
|
|
run_command(["mkswap", swap_part])
|
|
|
|
logger.info("Formatting Root partition...")
|
|
run_command(["mkfs.ext4", "-F", root_part])
|
|
|
|
logger.info("Partitioning and formatting complete.")
|
|
|
|
result = {
|
|
"efi": efi_part,
|
|
"root": root_part
|
|
}
|
|
if use_swap:
|
|
result["swap"] = swap_part
|
|
else:
|
|
# If no swap, we should probably return None or handle it in fstab generation
|
|
# backend/os_install.py expects "swap" key for UUID.
|
|
# We should probably pass a dummy or None, and update os_install to handle it.
|
|
result["swap"] = None
|
|
|
|
return result
|
|
|
|
def mount_partitions(partition_info, mount_root="/mnt"):
|
|
"""
|
|
Mounts the partitions into mount_root.
|
|
"""
|
|
import os
|
|
|
|
# 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
|
|
efi_mount = os.path.join(mount_root, "boot/efi")
|
|
if not os.path.exists(efi_mount):
|
|
os.makedirs(efi_mount, exist_ok=True)
|
|
|
|
run_command(["mount", partition_info["efi"], efi_mount])
|
|
|
|
# 3. Enable Swap
|
|
if partition_info.get("swap"):
|
|
run_command(["swapon", partition_info["swap"]])
|
|
|
|
logger.info(f"Partitions mounted at {mount_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])
|