grub

#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2014, William Leemans <willie@elaba.net>

import datetime
import os
import re
import subprocess
import shutil
import datetime

DOCUMENTATION = """
---
module: grub
author: William Leemans <willie@elaba.net>
short_description: modify GRUB configuration
description:
 - This module allows the manipulation of the grub configuration
options:
  state:
    required: false
    choices: [ present, absent ]
    default: "present"
    aliases: []
    description:
      - Whether the setting should be added/modified or removed

  default:
    required: false
    aliases: []
    description:
      - Which grub stanza is the default

  timeout:
    required: false
    aliases: []
    description:
      - Grub timeout to wait for user input

  koption:
    required: false
    aliases: []
    description:
      - Which kernel option to modify

  kvalue:
    required: false
    aliases: []
    description:
      - Add a value to the specified kernel option. If there's no value associated to 
        the option, then ommit this argument.
  
  grub:
    required: false
    description:
      - set the path to the grub config
      - /boot/grub/grub.conf in case of grub
      - /etc/default/grub in case of grub2

  allow_multi:
    required: false
    choices: [ yes, no ]
    description:
      - Whether an option can be defined multiple times
    
"""

EXAMPLES = r"""
# Add/Change root=/dev/sda1 to the kernel line(s)
- grub: koption=root kvalue=/dev/sda1

# Remove the quiet kernel option
- grub: state=absent koption=quiet

# Set the default boot stanza to 0
- grub: default=0

# Change the grub wait timeout to 15
- grub: timeout=15

"""

class Grub(object):

    def __init__(self, module):
        self.module       = module
        self.state        = module.params['state']
        self.default      = module.params['default']
        self.timeout      = module.params['timeout']
        self.koption      = module.params['koption']
        self.kvalue       = module.params['kvalue']
        self.config_file  = module.params['grub']
        self.allow_multi  = module.params['allow_multi']

        self.exists = True
        self.config_old = ""
        self.config_new = ""
        self.grub_bin_install  = None
        self.grub_bin_mkconfig = None
        self.grub_bin_version  = -1

        self._get_grub_version()

        self.grub_bin_version = 0.99

        if self.config_file is None and self.grub_bin_version >= 1:
            self.config_file = "/etc/default/grub"
        elif self.config_file is None and self.grub_bin_version > 0.9 and self.grub_bin_version < 1:
            self.config_file = "/boot/grub/grub.conf"

        if os.path.isdir(os.path.dirname(self.config_file)) and os.path.isfile(self.config_file):
            with open(self.config_file, "r") as f:
                self.config_old = f.read()
        else:
            self.exists = False

        self.config_new = self.config_old

    def _get_grub_version(self):
        for path in os.environ['PATH'].split(os.pathsep):
            path = path.strip('"')
            for binname in [ 'grub-install', 'grub2-install' ]:
                binpath = os.path.join(path, binname)
                if os.path.isfile(binpath) and os.access(binpath, os.X_OK):
                    self.grub_bin_install = binpath
            for binname in [ 'grub-mkconfig', 'grub2-mkconfig' ]:
                binpath = os.path.join(path, binname)
                if os.path.isfile(binpath) and os.access(binpath, os.X_OK):
                    self.grub_bin_mkconfig = binpath
        if self.grub_bin_install is not None:
            proc = subprocess.Popen([self.grub_bin_install,"--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out, err = proc.communicate()
            self.grub_bin_version = float(re.search('[0-9]+\.[0-9]+',out).group(0))

    def _kernel_options_tolist(self, options):
        t = re.findall(r'([^\s]*=\$\([^\(]*\)|\$\([^\(]*\)|[^\s]*)', options)
        r = []
        for v in t:
            if v != '':
                if v.find('$(') >= 0 and v.find('$(') < v.find("="):
                    akey = v
                    aval = ""
                else:
                    akey = v.partition("=")[0]
                    aval = v.partition("=")[2]
            
                replaced = False
                if akey == self.koption and self.allow_multi is not True and akey[:1] != '$':
                    for i, el in enumerate(r):
                        if el[0] == akey and el[0]:
                            r[i] = (akey, aval)
                            replaced = True
                if replaced is not True:
                    r.append((akey, aval))
        return r

    def _kernel_options_tostring(self, options):
        t = []
        for key, val in options:
            if val == '' or val is None:
                t.append(key)
            else:
                t.append('%s=%s' % (key, val))
        t = ' '.join(t)
        return t

    def _kernel_remove_option(self, options):
        t = self._kernel_options_tolist(options)
        t = [ (key, val) for key, val in t if key != self.koption or ( key == self.koption and val != self.kvalue ) ]
        return self._kernel_options_tostring(t)

    def _kernel_set_option(self, options):
        t = self._kernel_options_tolist(options)

        found = False
        for i, el in enumerate(t):
            if el[0] == self.koption and el[1] == self.kvalue:
                found = True
        if found is True:
            return self._kernel_options_tostring(t)

        replaced = False
        for i, el in enumerate(t):
            found = False
            if el[0] == self.koption and self.allow_multi is not True:
                t[i] = (el[0], self.kvalue)
                replaced = True
        if replaced is not True:
            t.append((self.koption, self.kvalue))
        return self._kernel_options_tostring(t)

    def _grub2_modify_kernel_options(self):
        newfile = ""
        for line in self.config_new.splitlines():
            if re.compile(r'^\s*GRUB_CMDLINE_LINUX=').search(line) is not None:
                t = line.partition("=")[2].strip()
                if t[:1] == '"' and t[-1:] == '"':
                    t = t[1:-1]
                if self.state == "present":
                    t = self._kernel_set_option(t)
                else:
                    t = self._kernel_remove_option(t)
                line = re.sub(r'(\s*GRUB_CMDLINE_LINUX=).*$', r'\1"%s"' % t, line)
            newfile = newfile + line + "\n"
        self.config_new = newfile
        
    def _grub_modify_kernel_options(self):
        newfile = ""
        
        for line in self.config_new:
            if re.compile(r'^\s*kernel ').search(line) is not None:
                t = line.split()
                k = t[1]
                if self.state == "present":
                    o = self._kernel_set_option(" ".join(t[2:]))
                else:
                    o = self._kernel_remove_option(" ".join(t[2:]))
                line = re.sub(r'(\s*kernel ).*$', r'\1 %s %s' % ( k, o ), line)
            newfile = newfile + line + "\n"
        self.config_new = newfile

    def _grub2_refresh_menu(self):
        p = subprocess.Popen([self.grub_bin_mkconfig, "-o", "/boot/grub2/grub.cfg"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = p.communicate()
        if err != 0:
            return ( True, "" )
        else:
            return ( False, "Couln't make new config" )

    def _grub2_set_default(self):
        newfile = ""
        for line in self.config_new.splitlines():
            if re.compile(r'^\s*GRUB_DEFAULT=').search(line) is not None:
                line = re.sub(r'(\s*GRUB_DEFAULT=).*$', r'\1"%s"' % self.default, line)
            newfile = newfile + line + "\n"
        self.config_new = newfile

    def _grub2_set_timeout(self):
        newfile = ""
        for line in self.config_new.splitlines():
            if re.compile(r'^\s*GRUB_TIMEOUT=').search(line) is not None:
                line = re.sub(r'(\s*GRUB_TIMEOUT)=.*$', r'\1="%s"' % self.timeout, line)
            newfile = newfile + line + "\n"
        self.config_new = newfile

    def _grub_set_default(self):
        newfile = ""
        for line in self.config_new.splitlines():
            if re.compile(r'^\s*default=').search(line) is not None:
                line = re.sub(r'(\s*default=).*$', r'\1%s' % self.default, line)
            newfile = newfile + line + "\n"
        self.config_new = newfile
        
    def _grub_set_timeout(self):
        newfile = ""
        for line in self.config_new.splitlines():
            if re.compile(r'^\s*timeout=').search(line) is not None:
                line = re.sub(r'(\s*timeout)=.*$', r'\1=%s' % self.timeout, line)
            newfile = newfile + line + "\n"
        self.config_new = newfile
        
    def remove_kernel_option(self):
        if self.grub_bin_version >= 1:
            self._grub2_modify_kernel_options()
        elif self.grub_bin_version > 0.9 and self.grub_bin_version < 1:
            self._grub_modify_kernel_options()

    def set_kernel_option(self):
        if self.grub_bin_version >= 1:
            self._grub2_modify_kernel_options()
        elif self.grub_bin_version > 0.9 and self.grub_bin_version < 1:
            self._grub_modify_kernel_options()

    def set_default(self):
        if self.grub_bin_version >= 1:
            self._grub2_set_default()
        elif self.grub_bin_version > 0.9 and self.grub_bin_version < 1:
            self._grub_set_default()
        
    def set_timeout(self):
        if self.grub_bin_version >= 1:
            self._grub2_set_timeout()
        elif self.grub_bin_version > 0.9 and self.grub_bin_version < 1:
            self._grub_set_timeout()
        
    def save_config(self, contents):
        try: 
            with open(self.config_file, "w") as f:
                f.write(contents)
        except:
            return ( False, "Couldn't save config" )

        if self.grub_bin_version >= 1:
            return self._grub2_refresh_menu()
    
    def backup_config(self):
        try:
            shutil.copy(self.config_file, "%s.%s" % ( self.config_file, datetime.datetime.now().strftime("%y%m%d%H%M")))
            return True
        except:
            return False

    def haschanged(self):
        return self.config_old != self.config_new
        

def main():
    module = AnsibleModule(
        argument_spec = dict(
            state        = dict(default = 'present', choices = ['present', 'absent'], type = 'str'),
            default      = dict(default = None, type = 'int'),
            timeout      = dict(default = None, type = 'int'),
            koption      = dict(default = None, type = 'str'),
            kvalue       = dict(default = None, type = 'str'),
            grub         = dict(default = None, type = 'str'),
            allow_multi  = dict(default = False, type = 'bool'),
        ),
        supports_check_mode = True
    )

    grub = Grub(module)
    if not grub.exists:
        module.fail_json(msg="%s couldn't be found" % grub.config_file )

    out = ''
    err  = ''
    result = {}
    result['state'] = grub.state
    msg = ""

    if grub.state == 'absent':
        if grub.koption is not None:
            grub.remove_kernel_option()
        else:
            module.fail_json(msg="No kernel option specified to modify")

    elif grub.state == 'present':
        if grub.default is not None:
            grub.set_default()
        if grub.timeout is not None:
            grub.set_timeout()
        if grub.koption is not None:
            grub.set_kernel_option()

    result['diff'] = { 'before': grub.config_old,
                       'after': grub.config_new }

    if not module.check_mode and grub.haschanged():
        if not grub.backup_config():
            module.fail_json(msg="Couldn't create backup of config")
        (result, msg) = grub.save_config()
        if not result:
            module.fail_json(msg=msg)

    module.exit_json(changed=grub.haschanged(), msg=msg, **result)

# import module snippets
from ansible.module_utils.basic import *
main()