#!/usr/bin/python3
import os
import re
import shutil

#
# подаем каталог на вход, 
# на выходе получаем уид раздела которому он принадлежит 
# и порядковый номер раздела (начинаем с нуля, как того требует ПНС) 
# и имя раздела стандартного вида (sda1,sdb1 etc)
#
# ожидается разумная подача аргумента, проверок нет)
#
def getPartitionInfo(directory = '/'):    
# 1. Определяем устройство из /proc/mounts
    device = False
    with open('/proc/mounts', 'r') as f:
        for line in f:            
            parts = line.split()
            if len(parts) >= 2 and parts[1] == directory:
                device = parts[0]
                break
    if not device:
        with open('/proc/mounts', 'r') as f:    
            if 'boot' in directory:
                directory = directory[:-4]                
                for line in f:                    
                    parts = line.split()                    
                    if len(parts) >= 2 and parts[1] == directory:
                        device = parts[0]                        
                        break
    if not device:
        raise RuntimeError("In /proc/mounts not found device for these directory")
    
# 2. Нормализуем путь устройства
# Если устройство начинается с /dev, убираем префикс
    if device.startswith('/dev/'):
        device_name = os.path.basename(device)
    else:
        device_name = device
    
# 3. Ищем UUID в /dev/disk/by-uuid
    uuid_dir = '/dev/disk/by-uuid'
    if os.path.exists(uuid_dir):
        for uuid in os.listdir(uuid_dir):
            uuid_path = os.path.join(uuid_dir, uuid)
            try:
# Получаем реальное устройство, на которое указывает симлинк
                real_device = os.path.realpath(uuid_path)
                real_device_name = os.path.basename(real_device)
                match = re.search(r'(\d+)$', real_device_name)
                partNum = -1
                if match:
                    partNum = int(match.group(1)) - 1
                else:
                    partNum = -1
# Проверяем, совпадает ли с нашим корневым устройством
                if real_device_name == device_name:
                    return [uuid,partNum,real_device_name]
            except OSError:
                continue
    return None

def getOsInfo(osReleaseFullName):
    
    osInfo = {}    
    with open(osReleaseFullName, 'r') as f:
        for line in f:
            if '=' in line:
                key, value = line.strip().split('=', 1)
                if key in ['NAME', 'VERSION']:
# Убираем кавычки из значения
                    osInfo[key] = value.strip('"')
    
    return f"{osInfo.get('NAME', '').split()[0]}_{osInfo.get('NAME', '').split()[1]}-{osInfo.get('VERSION', '')}"

def weInInstallationProcess():
# Ощущаем себя в процессе установки при наличии следующих каталогов
    testDir1 ='/mnt/sysimage'
    testDir2 ='/mnt/sysroot'    
    if os.path.exists(testDir1) and os.path.exists(testDir2):
        return True
    else:
        return False
    
def getKernels(bootPath = '/boot'):    
    files = os.listdir(bootPath)
    kernels = []
    for kernel in files:
        if kernel.startswith('vmlinuz-') and ('rescue' not in kernel):
            kernels.append(kernel)
    initramfses = []
    for initramfs in files:
        if initramfs.startswith('initramfs-') and initramfs.endswith('.img'):
            initramfses.append(initramfs)
    pairs = []
    def version_key(f):
 # Извлекаем числа из vmlinuz-A.B.C-D.red80.x86_64
            nums = re.findall(r'\d+', f)
            return tuple(map(int, nums)) if nums else (0,)
        
    kernels.sort(key=version_key, reverse=True)
    for kernel in kernels:
        for initramfs in initramfses:
            if kernel[8:] == initramfs[10:-4]:
                pairs.append([kernel,initramfs])
    return pairs

def parseBootBonf(filePath):
    entries = []
    currentEntry = {}
    
    with open(filePath, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    for line in lines:
        line = line.strip()
        
        if not line:
            continue
        
        if line.startswith('label='):
 # Если у нас уже есть запись, сохраняем её
            if currentEntry:
                entries.append(currentEntry)
                currentEntry = {}
            
 # Убираем 'label=' и возможные кавычки
            labelValue = line[6:].strip()
            if (labelValue.startswith('"') and labelValue.endswith('"')) or \
               (labelValue.startswith("'") and labelValue.endswith("'")):
                labelValue = labelValue[1:-1]
            
            currentEntry['label'] = labelValue
        
        elif '=' in line:
 # Обработка остальных параметров
            key, value = line.split('=', 1)
            key = key.strip()
            value = value.strip()
            
 # Убираем кавычки если есть
            if (value.startswith('"') and value.endswith('"')) or \
               (value.startswith("'") and value.endswith("'")):
                value = value[1:-1]
            
            currentEntry[key] = value
    
 # Добавляем последнюю запись
    if currentEntry:
        entries.append(currentEntry)
    
    return entries


# Шапка
defaultOptions = 'audit=0 quiet rhgb ro'
title = '#\n# This file autogenerated by boot_conf_generator - tool for generating boot.conf for Elbrus platform\n'
title2 = '# Please do not modify it manually!\n#\n\n'

timeout = 'timeout=10\n\n'
# Проверяем не в процессе ли мы установки и соответственно выставляем точки отсчета
prefix = '/'
if weInInstallationProcess():
    prefix = '/mnt/sysroot/'
else:
    prefix = '/'
if prefix == '/mnt/sysroot/':
    rootDeviceInfo = getPartitionInfo(prefix[:-1])
else:
    rootDeviceInfo = getPartitionInfo(prefix)
bootDeviceInfo = getPartitionInfo(prefix + 'boot')
pairsKernelInitramfs = getKernels(prefix + 'boot')
osNameVersion = getOsInfo(prefix + 'etc/os-release')

# Делаем новый список меток загрузки
newEntries = []
for line in pairsKernelInitramfs:
    newEntries.append({'label':osNameVersion + '_(' + line[0][8:-10] + ')_('+ rootDeviceInfo[2] +')',
                       'partition':bootDeviceInfo[1],
                       'image':'/'+line[0],
                       'cmdline':'root=UUID='+rootDeviceInfo[0]+' '+defaultOptions,
                       'initrd':'/'+line[1]})
# Проверяем есть ли файл бут конф и при наличии загружаем оттуда старый список меток загрузки
# После чего сразу делаем бекап
oldEntries = []
if os.path.exists(prefix + 'boot/boot.conf'):
    oldEntries = parseBootBonf(prefix + 'boot/boot.conf')
    srcFile = prefix + 'boot/boot.conf'
    dstFile = prefix + 'boot/boot.conf.BACKUP'
    shutil.copy2(srcFile, dstFile)

# Сверяем списки, если находим в старом метку аналогичную новой, 
# но с измененными параметрами ядра и при условии что это ядро сейчас тоже существует, 
# оставляем старую метку, 
# это для того чтобы не испортить что-то что отредактировал пользователь
if len(oldEntries) > 0:
    for pair in pairsKernelInitramfs:
        kernel = '/' + pair[0]
        for oldEntry in oldEntries:
            for newEntry in newEntries:            
                if (oldEntry.get('image') == kernel) and (newEntry.get('image') == kernel):                
                    if (oldEntry.get('partition') != newEntry.get('partition')) or \
                       (oldEntry.get('cmdline') != newEntry.get('cmdline')) or \
                       (oldEntry.get('initrd') != newEntry.get('initrd')):
                        newEntries[newEntries.index(newEntry)]=oldEntry

# Пишем полученный список в файл
with open('boot.conf', 'w', encoding='utf-8') as f:
    f.write(title)
    f.write(title2)
    f.write('default='+newEntries[0].get('label')+'\n')
    f.write(timeout)
    for entry in newEntries:
        f.write(f"label={entry['label']}\n")
        f.write(f"    partition={entry.get('partition', '0')}\n")
        f.write(f"    image={entry.get('image', '')}\n")
        f.write(f"    cmdline={entry.get('cmdline', '')}\n")
        f.write(f"    initrd={entry.get('initrd', '')}\n")        
        f.write(f"\n") 

