#!/usr/bin/env python2
# -*- mode: python; coding: utf-8 -*-

"""
grub-mkconfig helper script.  Intended as a replacement for /etc/grub.d/10_linux .

Replacement version by Jim Rees, March 2010.  The standard version doesn't
give you any way to have different sets of options for different kernels,
and doesn't let you customize the names and orders of the kernels.

-c <filename>: config file name, default /etc/grub.conf
-d <bootdir>: boot dir, default /boot
-n: list kernels but don't generate grub entries
-u: use file system uuid instead of device name
-U: generate Ubuntu grub.cfg, else Arch

Config file format is (for example):
grubdev (hd0,1)
root /dev/sda1
kernel 2.6.31-??-generic apparmor=0 nomodeset
kernel 2.6.31.6 nomodeset
kernel 2.6.31.6 nomodeset single

The second field is a glob pattern, everything after that is a set of kernel
options to apply to kernels matching that pattern.  It's ok for the file to
be missing, which is equivalent to "kernel *" (list all kernels, with no
options).

    Copyright 2010 Jim Rees

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""

import getopt, os, sys, fnmatch, re

confname = "/etc/grub.conf"
bootdirname = "/boot"

nflag = False
uflag = False
Uflag = False

arch_extra = ('load_video', 'set gfxpayload=keep',
              'insmod gzio', 'insmod part_msdos', 'insmod ext2')
ubuntu_extra = ('recordfail', 'insmod gzio', 'insmod ext2')
arch_initrd = 'initramfs-%s.img'
ubuntu_initrd = 'initrd.img-%s'

def kkey(name):
    """ convert kernel version string to sort key, "3.2.0-24-generic"
    Each part is alternatively a string of digits or non-digits
    """
    parts = [kcconv(c) for c in re.split('(\d+)', name) if len(c) > 0]
    return parts

def kcconv(c):
    """ numbers sort in reverse
    """
    return str(999999 - int(c)) if c.isdigit() else c

def print_kernel(k, grub_dev, grub_dev_uuid, linux_dev, kopts):
    global nflag, uflag, Uflag

    if nflag:
        print k
    else:
        print >> sys.stderr, k
        print 'menuentry "%s %s" {' % (k, kopts)

        extras = ubuntu_extra if Uflag else arch_extra
        for s in extras:
            print '  %s' % (s)

        print '  set root=%s' % (grub_dev)

        if uflag:
            print '  search --no-floppy --fs-uuid --set %s' % (grub_dev_uuid)

        print '  linux /boot/vmlinuz-%s root=%s %s' % (k, linux_dev, kopts)

        initrd_fmt = ubuntu_initrd if Uflag else arch_initrd
        print '  initrd /boot/' + initrd_fmt % (k)

        print '}'

def print_img(k, grub_dev):
    global nflag, uflag, Uflag

    if nflag:
        print k
    else:
        print >> sys.stderr, k
        print 'menuentry "%s" {' % (k)
        print '  insmod ext2'
        print '  set root=%s' % (grub_dev)
        print '  linux16 /boot/memdisk'
        print '  initrd16 /boot/%s' % (k)
        print '}'

def main():
    global confname, bootdirname, nflag, uflag, Uflag
    grub_dev = None
    linux_dev = None

    # Parse command line options

    opts, args = getopt.getopt(sys.argv[1:], "c:d:nuU", [])
    for option, value in opts:
        if option == '-c':
            confname = value
        elif option == '-d':
            bootdirname = value
        elif option == '-n':
            nflag = True
        elif option == '-u':
            uflag = True
        elif option == '-U':
            Uflag = True

    # Read the config file

    try:
        f = open(confname, 'r')
        conflines = f.readlines()
        f.close()
    except:
        conflines = ['kernel *']

    conf = []
    for l in conflines:
        if l[0] == '#':
            continue
        ls = l.split()
        if len(ls) == 0:
            continue
        if ls[0] == 'kernel':
            conf.append((ls[1], ' '.join(ls[2:])))
        elif ls[0] == 'grubdev':
            grub_dev = ls[1]
        elif ls[0] == 'root':
            linux_dev = ls[1]

    # Get device names and other grub info

    if grub_dev == None:
        grub_dev = '(hd0,1)'

    grub_dev_uuid = os.getenv('GRUB_DEVICE_UUID')

    if uflag:
        linux_dev = 'UUID=' + grub_dev_uuid
    elif linux_dev == None:
        f = os.popen("mount | grep ' on / '", 'r')
        linux_dev = f.readline().split()[0]
        f.close()

    # Get the list of available kernels in /boot

    klist = os.listdir(bootdirname)
    knlist = []
    for k in klist:
        if k.startswith('vmlinuz-') and not k.endswith('-') and not k.endswith('.old'):
            knlist.append(k.replace('vmlinuz-', ''))
    knlist.sort(key=kkey)

    # For each line in the config file, add all matching kernels

    for c in conf:
        for k in knlist:
            if fnmatch.fnmatch(k, c[0]):
                print_kernel(k, grub_dev, grub_dev_uuid, linux_dev, c[1])

    # Also put in Thinkpad bios updates (this part not configurable)

    for k in klist:
        if k.endswith('.IMG'):
            print_img(k, grub_dev)

    return 0

if __name__ == "__main__":
    sys.exit(main())