#!/usr/bin/env python3
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You may
# not use this file except in compliance with the License. A copy of the
# License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

"""Enable SELinux enforcing mode"""

import io
import re
import subprocess
from textwrap import dedent

from logging import Logger

from cloudinit.cloud import Cloud
from cloudinit.config.schema import MetaSchema, get_meta_doc
from cloudinit.distros import ALL_DISTROS
from cloudinit.settings import PER_INSTANCE
from cloudinit import util

MODULE_DESCRIPTION = """\
Toggle SELinux enforcing mode.
"""

CONFIG_FILE = '/etc/selinux/config'
REBOOT_FLAG = '/run/cloud-init-selinux-reboot'

meta: MetaSchema = {
    "id": "cc_selinux",
    "name": "SELinux",
    "title": "Controls a system's SELinux configuration",
    "description": MODULE_DESCRIPTION,
    "distros": [ALL_DISTROS],
    "frequency": PER_INSTANCE,
    "examples": [
        dedent(
            """
            selinux:
              mode: enforcing
              selinux_no_reboot: true"""
        )
    ],
}

__doc__ = get_meta_doc(meta)


def get_selinux_state():
    state = {}
    match_spec = [
        ('status', r'SELinux status'),
        ('root_directory', r'SELinux root directory'),
        ('mount_point', r'SELinuxfs mount'),
        ('loaded_policy_name', r'Loaded policy name'),
        ('current_mode', r'Current mode'),
        ('config_mode', r'Mode from config file'),
        ('policy_mls_status', r'Policy MLS status'),
        ('policy_deny_unknown_status', r'Policy deny_unknown status'),
        ('memory_protection_checking', r'Memory protection checking'),
        ('max_kernel_policy_version', r'Max kernel policy version'),
    ]
    with subprocess.Popen("sestatus", stdout=subprocess.PIPE, text=True) as p:
        for line in p.stdout.readlines():
            for r in match_spec:
                m = re.match(r"^(?P<%s>%s):\s+(.*)" % r, line)
                if m:
                    state[r[0]] = m.group(2)
    return state


def _update_config_mode(mode: str):
    if mode not in ("enforcing", "permissive"):
        raise ValueError("Invalid mode")
    content = io.StringIO()
    with open(CONFIG_FILE) as infd:
        for line in infd.readlines():
            if re.match(r'^SELINUX=', line):
                content.write("# Added by cloud-init\n")
                content.write("SELINUX=%s\n" % mode)
            elif line.startswith("# Added by cloud-init"):
                # skip our own comments
                True
            else:
                content.write(line)
    util.write_file(CONFIG_FILE, content.getvalue())


def secure_mode_policyload(setmode=None):
    if setmode:
        if setmode not in ("on", "off"):
            raise ValueError("Invalid value for setmode")
        rv = subprocess.run(
            ["setsebool", "-P", "secure_mode_policyload", setmode],
            check=True)
    rv = subprocess.run(["getsebool", "secure_mode_policyload"],
                        stdout=subprocess.PIPE, text=True, check=True)
    out = re.match(r'secure_mode_policyload --> (\w+)', str(rv.stdout))
    if out.group(1) == "on":
        return True
    return False


def _update_kernel_param(mode: str, log):
    if mode == "enforcing":
        param = "enforcing=1"
    elif mode == "permissive":
        param = "enforcing=0"
    else:
        raise ValueError("Invalid value")

    rv = subprocess.run(
        ["grubby", "--update-kernel", "ALL", "--args", param],
        check=True)
    log.debug("grubby exited with status %d", rv.returncode)


def lock_config(state: dict, log):
    rv = subprocess.run(["chattr", "+i", CONFIG_FILE], check=True)
    log.debug("chattr +i %s exited with status %d", CONFIG_FILE, rv.returncode)
    secure_mode_policyload('on')


def unlock_config(state: dict, log):
    rv = subprocess.run(["chattr", "-i", CONFIG_FILE], check=True)
    log.debug("chattr -i %s exited with status %d", CONFIG_FILE, rv.returncode)


def set_selinux_mode(mode: str, state: dict, log: Logger, lock=True):
    if mode == "enforcing":
        param = 1
        if state['current_mode'] == 'enforcing':
            # we're already in enforcing mode, so we're unlikely to be
            # able to do anything even if we try
            log.debug("SELinux enforcing mode is already active, skipping.")
            return
    elif mode == "permissive":
        param = 0
    else:
        raise ValueError("Invalid SELinux mode provided")

    if mode == "enforcing":
        _update_config_mode(mode)
        _update_kernel_param(mode, log)
        if lock:
            lock_config(state, log)
        _ = subprocess.run(["setenforce", str(param)], check=True)
    elif mode == "permissive":
        # We may need to reboot to get back to permissive mode if
        # secure_mode_policyload is active
        rv = subprocess.run(["setenforce", str(param)])
        if rv.returncode != 0:
            log.info("Unable to switch to permissive mode. Reboot required.")
        _update_kernel_param(mode, log)
        unlock_config(state, log)
        _update_config_mode(mode)

    # return the updated state
    return get_selinux_state()


def handle(name: str, cfg: dict, cloud: Cloud, log: Logger, args: list):
    state = get_selinux_state()
    if 'status' in state:
        log.debug("Current SELinux state is %s" % state['status'])
        log.debug("Current SELinux mode is %s" % state['current_mode'])
        log.debug("Configured SELinux mode is %s" % state['config_mode'])
        log.debug("Current SELinux policy is %s" % state['loaded_policy_name'])
    else:
        log.warn("Unable to determine current SELinux status")
        return
    selinux_cfg = util.get_cfg_by_path(cfg, ("selinux"))
    if selinux_cfg is None:
        log.debug("No SELinux configuration, skipping")
        return
    mode = util.get_cfg_option_str(selinux_cfg, "mode", None)
    if mode is None:
        log.warn("No mode setting in selinux configuration")
        return
    elif mode == state['current_mode']:
        log.debug("SELinux is already in %s mode, no changes needed" % mode)
        return
    lock = util.get_cfg_option_bool(selinux_cfg, "mode_lock", True)
    state = set_selinux_mode(mode, state, log, lock)
    skip_reboot = util.get_cfg_option_bool(selinux_cfg,
                                           "selinux_no_reboot", False)
    if not skip_reboot:
        log.info("Creating %s" % REBOOT_FLAG)
        open(REBOOT_FLAG, "w")
    else:
        log.info("Not creating %s per configuration" % REBOOT_FLAG)
    log.info("Setting SELinux to %s mode" % state['current_mode'])

# Local variables:
# mode: python
# indent-tabs-mode: nil
# tab-width: 4
# end: