Add module for switching kernel to HWE

Related-Prod: PROD-30103

Change-Id: Ice332409f7876683109e5134f6ce3496b93e1152
diff --git a/README.rst b/README.rst
index 611a1ee..ea0b419 100644
--- a/README.rst
+++ b/README.rst
@@ -2811,6 +2811,72 @@
 
    ip link set dev eth0 mtu 1400
 
+Switch Kernel from non-HWE to HWE
+==================================
+
+It is possible to switch Kernel from non-HWE to HWE by using module
+linux_kernel_switch. It has few methods:
+
+* check_hwe_kernel
+* switch_kernel
+* rollback_switch_kernel
+
+Method ``check_hwe_kernel`` allows to check whether HWE kernel installed
+or not:
+
+.. code-block:: bash
+
+   salt <target> linux_kernel_switch.check_hwe_kernel
+
+Output for case HWE is installed:
+
+.. code-bloc:: bash
+
+  kvm02.cluster-env.local:
+      ----------
+      linux-image-extra-virtual-hwe-16.04:
+          ----------
+          linux-image-extra-virtual-hwe-16.04:
+              ----------
+              architecture:
+                  amd64
+              description:
+                  Extra drivers for Virtual Linux kernel image
+                   This package will always depend on linux-image-generic.
+              group:
+                  kernel
+              install_date:
+                  2019-10-01T11:50:15Z
+              name:
+                  linux-image-extra-virtual-hwe-16.04
+              packager:
+                  Ubuntu Kernel Team <kernel-team@lists.ubuntu.com>
+              source:
+                  linux-meta-hwe
+              version:
+                  4.15.0.54.75
+          ...
+
+Output for case HWE is not installed:
+
+.. code-bloc:: bash
+
+  kvm02.cluster-env.local:
+      ----------
+      linux-image-extra-virtual-hwe-16.04:
+          Not installed!
+      linux-image-generic-hwe-16.04:
+          Not installed!
+
+Method ``switch_kernel`` allows you to switch from non-HWE to HWE. It has
+two options: ``dry_run`` - to check what packages are going to be installed or
+removed and ``only_kernel`` - install only Kernel image packages without other
+HWE packages.
+
+Method ``rollback_switch_kernel`` allows you to rollback method
+``switch_kernel`` which was executed successfully previously. Option
+``dry_run`` - to check what packages are going to be installed/removed.
+
 Read more
 =========
 
diff --git a/_modules/linux_kernel_switch.py b/_modules/linux_kernel_switch.py
new file mode 100644
index 0000000..3e634cb
--- /dev/null
+++ b/_modules/linux_kernel_switch.py
@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+'''
+Manage Kernel switch from generic to hwe
+
+'''
+import logging
+import json
+import os
+import re
+import salt.utils
+
+logger = logging.getLogger(__name__)
+stream = logging.StreamHandler()
+logger.addHandler(stream)
+kernel_state_backup = '/etc/salt/.kernel_state_backup'
+
+
+def _get_hwe_packages(only_kernel=True):
+    distribRelease = __salt__['grains.get']('lsb_distrib_release')
+    hwe_pkgs = [ 'linux-image-generic-hwe-{0}'.format(distribRelease),
+                 'linux-image-extra-virtual-hwe-{0}'.format(distribRelease) ]
+    if only_kernel:
+        return hwe_pkgs
+
+    return hwe_pkgs + [ 'linux-generic-hwe-{0}'.format(distribRelease),
+                    'linux-headers-virtual-hwe-{0}'.format(distribRelease) ]
+
+
+def check_hwe_kernel():
+    pgks_res = {}
+    for pkg in _get_hwe_packages():
+        try:
+            pgks_res[pkg] = __salt__['pkg.info_installed'](pkg)
+        except:
+            pgks_res[pkg] = 'Not installed'
+            continue
+    return pgks_res
+
+
+def switch_kernel(dry_run=False, only_kernel=True):
+    kernel_state = {}
+    kernel = __salt__['cmd.shell']('uname -r | cut -f 1 -d "-"')
+    hwe_pkgs_array = _get_hwe_packages(only_kernel)
+    hwe_pkgs = ','.join(hwe_pkgs_array)
+    if dry_run:
+        kernel_state['to_install'] = hwe_pkgs_array
+    else:
+        kernel_state['installed'] = __salt__['pkg.install'](hwe_pkgs)
+        with open(kernel_state_backup, 'w') as f:
+            f.write(json.dumps(kernel_state))
+
+    gen_pkgs_to_remove = [ 'linux-image-generic', 'linux-headers-generic',
+                        'linux-image-extra-virtual', 'linux-image-virtual',
+                        'linux-signed-generic', 'linux-signed-image-generic' ]
+    installed_gen_pkg = __salt__['pkg.list_pkgs']()
+    for pkg in installed_gen_pkg:
+        if (re.match("^linux-headers.*-{0}-.*".format(kernel), pkg) or
+            re.match("^linux-image.*-{0}-.*".format(kernel), pkg) or
+            re.match("^linux-modules.*-{0}-.*".format(kernel), pkg) or
+            re.match("^linux-signed-image.*-{0}-.*".format(kernel), pkg)):
+            gen_pkgs_to_remove.append(pkg)
+
+    if dry_run:
+        kernel_state['to_remove'] = gen_pkgs_to_remove
+    else:
+        pkgs = __salt__['pkg.purge'](','.join(gen_pkgs_to_remove))
+        kernel_state['removed'] = pkgs['removed'].copy()
+        kernel_state['removed'].update(pkgs['installed'])
+        with open(kernel_state_backup, 'w') as f:
+            f.write(json.dumps(kernel_state))
+    return kernel_state
+
+
+def rollback_switch_kernel(dry_run=False):
+    kernel_info = {}
+    kernel_state = {}
+    if not os.path.isfile(kernel_state_backup):
+        return 'Nothing to rollback.'
+
+    with open(kernel_state_backup, 'r') as f:
+        kernel_info = json.loads(f.read())
+
+    gen_pkgs = []
+    for pkg,ver in kernel_info['removed'].items():
+        gen_pkgs.append('{0}={1}'.format(pkg, ver['old']))
+    hwe_pkgs = []
+    for pkg,ver in kernel_info['installed'].items():
+        hwe_pkgs.append(pkg)
+
+    if dry_run:
+        kernel_state['to_install'] = gen_pkgs
+        kernel_state['to_remove'] = hwe_pkgs
+    else:
+        kernel_state['installed'] = __salt__['pkg.install'](','.join(gen_pkgs))
+        pkgs = __salt__['pkg.purge'](','.join(hwe_pkgs))
+        kernel_state['removed'] = pkgs['removed'].copy()
+        kernel_state['removed'].update(pkgs['installed'])
+        os.remove(kernel_state_backup)
+
+    return kernel_state