#!/opt/alt/python35/bin/python3.5
#
# imunify360-pam        Python script to manage imunify360 pam module
#                       enabled/diabled state.
#

from __future__ import print_function
from configparser import ConfigParser
from collections import OrderedDict
import os
import platform
import re
import shutil
import sys
import subprocess

import yaml

CONFIG = '/etc/pam_imunify/i360.ini'
CONFIG_DOVECOT = '/etc/dovecot/dovecot.conf'
CONFIG_DOVECOT_TMPL = '/var/cpanel/templates/dovecot2.3/main.default'
CONFIG_DOVECOT_LOCAL = '/var/cpanel/templates/dovecot2.3/main.local'
CONFIG_PAM_DOVECOT = '/etc/pam.d/dovecot_imunify'
CONFIG_PAM_DOVECOT_DOMAINOWNER = '/etc/pam.d/dovecot_imunify_domainowner'

PAM_UNIX_REGEX = re.compile(r'auth\s+.+?\s+pam_unix\.so')


class Output:
    def status_changed(self, services):
        enabled = []
        already_enabled = []
        disabled = []
        already_disabled = []

        services = OrderedDict(sorted(services.items(), key=lambda x: x[0]))
        for key, value in services.items():
            enabled_prev, enabled_now = value
            if enabled_now:
                if enabled_prev:
                    already_enabled.append(key)
                else:
                    enabled.append(key)
            else:
                if not enabled_prev:
                    already_disabled.append(key)
                else:
                    disabled.append(key)

        message = None
        if len(enabled) > 0:
            message = 'imunify360-pam (%s) is now enabled.' \
                % ', '.join(enabled)

        if len(already_enabled) > 0:
            message = 'imunify360-pam (%s) is already enabled.' \
                % ', '.join(already_enabled)

        if len(disabled) > 0:
            message = 'imunify360-pam (%s) is now disabled.' \
                % ', '.join(disabled)

        if len(already_disabled) > 0:
            message = 'imunify360-pam (%s) is already disabled.' \
                % ', '.join(already_disabled)

        if message:
            print(message)

    def status(self, services):
        services = OrderedDict(sorted(services.items(), key=lambda x: x[1]))
        enabled = [key for key, value in services.items() if value]
        if len(enabled) > 0:
            print('status: enabled (%s)' % ', '.join(enabled))
        else:
            print('status: disabled')

    def warning(self, *args, **kwargs):
        print('[WARNING]', *args, **kwargs)

    def error(self, *args, **kwargs):
        print('[ERROR]', *args, **kwargs)


class YamlOutput(Output):
    def __init__(self):
        self.buffer = {}

    def __del__(self):
        print(yaml.safe_dump(self.buffer, default_flow_style=False))

    def status_changed(self, services):
        for service, value in services.items():
            enabled_prev, enabled_now = value
            if 'status_changed' not in self.buffer:
                self.buffer['status_changed'] = {}
            self.buffer['status_changed'][service] = {
                'from': 'enabled' if enabled_prev else 'disabled',
                'to': 'enabled' if enabled_now else 'disabled'
            }

    def status(self, services):
        for service, enabled in services.items():
            if 'status' not in self.buffer:
                self.buffer['status'] = {}
            self.buffer['status'][service] =  \
                'enabled' if enabled else 'disabled'

    def warning(self, *args, **kwargs):
        if 'warnings' not in self.buffer:
            self.buffer['warnings'] = []
        self.buffer['warnings'].append(' '.join(args))

    def error(self, *args, **kwargs):
        if 'errors' not in self.buffer:
            self.buffer['errors'] = []
        self.buffer['errors'].append(' '.join(args))


# This function get CP name only
def get_cp_name():
    panel = None

    # cPanel check
    if os.path.isfile('/usr/local/cpanel/cpanel'):
        panel = 'cpanel'

    # Plesk check
    elif os.path.isfile('/usr/local/psa/version'):
        panel = 'plesk'

    # DirectAdmin check
    elif os.path.isfile('/usr/local/directadmin/directadmin'):
        panel = 'directadmin'

    return panel


def readlink_f(filename):
    """
    Pythonic way of doing /bin/readlink --canonicalize filename
    and is needed for cPanel /etc/pam.d symlinks.
    """
    try:
        result = os.readlink(filename)
    except OSError:
        # not a symlink
        return filename

    if os.path.isabs(result):
        return result
    else:
        return os.path.join(os.path.dirname(filename), result)


def detect_conffiles(output=None):
    if not output:
        output = Output()

    distname = platform.linux_distribution()[0].split()[0].upper()
    if distname in ('DEBIAN', 'UBUNTU'):
        conffiles = '/etc/pam.d/common-auth',
    else:
        conffiles = '/etc/pam.d/password-auth', '/etc/pam.d/system-auth'

    if not all(os.path.exists(conf) for conf in conffiles):
        output.error("PAM configuration file(s) not found: %s" %
                     ' '.join(conffiles))
        sys.exit(1)

    return [readlink_f(fn) for fn in conffiles]


def atomic_rewrite(filename, content):
    """
    Atomically rewrites filename with given content to
    avoid possible "No space left on device"
    and unintenrional PAM module break.
    """
    shutil.copy(filename, filename + '.i360bak')
    tmp = filename + '.i360edit'
    with open(tmp, 'w') as tf:
        tf.write(content)
        st = os.stat(filename)
        os.fchmod(tf.fileno(), st.st_mode)
        os.fchown(tf.fileno(), st.st_uid, st.st_gid)
    os.rename(tmp, filename)


class i360RPatch:
    def __init__(self, conf_filename, output=None):
        self._conf_filename = conf_filename
        self.output = Output() if not output else output

    def filename(self):
        return os.path.join(
            os.path.dirname(self._conf_filename),
            '.%s.i360patch' % os.path.basename(self._conf_filename))

    def create_upon(self, content):
        cmd = ['/usr/bin/diff', '--unified=1', self._conf_filename, '-']
        proc = subprocess.Popen(cmd,
                                stdin=subprocess.PIPE,
                                stdout=open(self.filename(), 'w'))
        proc.communicate(content.encode())

        if proc.returncode != 1:
            # not a big deal: will use .i360bak as the last resort
            self.output.warning("'diff -u' error", file=sys.stderr)
            os.unlink(self.filename())

    def apply(self):
        """
        :raise CalledProcessError:
        """
        cmd = ['/usr/bin/patch', '--reverse', self._conf_filename]
        subprocess.check_call(cmd,
                              stdin=open(self.filename()),
                              stdout=open('/dev/null', 'w'))
        os.unlink(self.filename())


def read_config():
    ini_pairs = open(CONFIG).read()
    config = ConfigParser()
    config.read_string("[DEFAULT]\n" + ini_pairs)
    return config['DEFAULT']


def pam_unix_patch_around(pamconfig_lines, pam_unix_ln):
    match_offset = re.search(r'success=(\d)\s+default=ignore',
                             pamconfig_lines[pam_unix_ln])

    patch_simple(pamconfig_lines, pam_unix_ln)
    pam_unix_ln += 1

    if match_offset:
        fix_offset(pamconfig_lines,
                   pam_unix_ln,
                   int(match_offset.group(1)))


def toggle_dovecot_support(enable=True, output=None):
    """
    Enable or disable pam_imunify support for Dovecot
    """
    conf = CONFIG_DOVECOT_TMPL
    if not os.path.isfile(conf):
        return

    if not output:
        output = Output()

    passdb_regex = re.compile(r'^\s*passdb\s*\{.*?\}\s*$',
                              re.DOTALL | re.MULTILINE)

    data = open(conf).read()

    if enable:
        def passdb_replace(match):
            repl = re.sub(r'driver\s*=.*', 'driver = pam', match.group(0))
            repl = re.sub(r'args\s*=.*',
                          r'args = '
                          r'[% IF allow_domainowner_mail_pass %]'
                          r'dovecot_imunify_domainowner'
                          r'[% ELSE %]dovecot_imunify[% END %]',
                          repl)
            return repl

        data = re.sub(passdb_regex, passdb_replace, data)

        with open(CONFIG_DOVECOT_LOCAL, 'w') as f:
            f.write(data)
    else:
        os.unlink(CONFIG_DOVECOT_LOCAL)

    subprocess.call(["/scripts/builddovecotconf"])


def toggle_sshd_support(conffiles, enable=True, output=None):
    """
    Enable or disable pam_imunify module for sshd authentication
    """
    if not output:
        output = Output()

    if enable:
        for conf in conffiles:
            lines = open(conf).readlines()

            try:
                pam_unix_ln = next(
                    ln for ln, line in enumerate(lines)
                    if PAM_UNIX_REGEX.search(line)
                    )
            except StopIteration:
                output.error("PAM configuration file %s parse error" % conf)
                sys.exit(1)

            pam_unix_patch_around(lines, pam_unix_ln)

            content = ''.join(lines)
            i360RPatch(conf, output).create_upon(content)
            atomic_rewrite(conf, content)
    else:
        for conf in conffiles:
            rpatch = i360RPatch(conf, output)
            if os.path.exists(rpatch.filename()):
                try:
                    rpatch.apply()
                    continue
                except subprocess.CalledProcessError as e:
                    output.warning("'patch -R' was not successful: %s" % e,
                                   file=sys.stderr)
            else:
                output.warning("File not found: %s" % rpatch.filename(),
                               file=sys.stderr)

            atomic_rewrite(conf, open(conf + '.i360bak').read())


def set_panel_integration(panel, output=None):
    confs = [CONFIG_PAM_DOVECOT, CONFIG_PAM_DOVECOT_DOMAINOWNER]
    imunify_regex = re.compile(r'auth\s+sufficient\s+pam_imunify\.so')

    if not panel:
        return

    if not output:
        output = Output()

    for conf in confs:
        lines = open(conf).readlines()

        try:
            imunify_ln = next(
                ln for ln, line in enumerate(lines)
                if imunify_regex.search(line)
                )
        except StopIteration:
            output.error("PAM configuration file %s parse error" % conf)
            sys.exit(1)

        if 'domain_owner_mail_pass' in lines[imunify_ln]:
            lines[imunify_ln] = 'auth\tsufficient\tpam_imunify.so ' \
                                'domain_owner_mail_pass %s\n' % panel
        else:
            lines[imunify_ln] = 'auth\tsufficient\tpam_imunify.so %s\n' % panel

        content = ''.join(lines)
        atomic_rewrite(conf, content)
        os.unlink(conf + '.i360bak')


def patch_simple(pamconfig_lines, pam_unix_ln):
    pamconfig_lines.insert(pam_unix_ln + 1,
                           'auth\trequired\tpam_imunify.so\n')
    pamconfig_lines.insert(pam_unix_ln,
                           'auth\trequired\tpam_imunify.so\tcheck_only\n')


def fix_offset(pamconfig_lines, pam_unix_ln, pam_unix_success_offset):
    bump_to = pam_unix_success_offset + 1
    pamconfig_lines[pam_unix_ln] = re.sub(r'success=\d',
                                          'success=%d' % bump_to,
                                          pamconfig_lines[pam_unix_ln])


class Cmd:
    @classmethod
    def enable(cls, conffiles, output=None):
        if not output:
            output = Output()

        if any('pam_imunify.so' in open(conf).read() for conf in conffiles):
            cls._cphulk_check()
            output.status_changed({'sshd': (True, True)})
            return

        toggle_sshd_support(conffiles, True, output)

        cls._cphulk_check(output)
        output.status_changed({'sshd': (False, True)})

    @staticmethod
    def enable_dovecot(conffiles, output=None):
        if not output:
            output = Output()

        panel = get_cp_name()
        if panel:
            set_panel_integration(panel, output)
        else:
            output.error('No supported panel found.')
            return

        dovecot_enabled = os.path.isfile(CONFIG_DOVECOT) and \
            'dovecot_imunify' in open(CONFIG_DOVECOT).read()

        if not dovecot_enabled:
            toggle_dovecot_support(True, output)

        output.status_changed({'dovecot': (dovecot_enabled, True)})

    @staticmethod
    def disable_all(conffiles, output=None):
        if not output:
            output = Output()

        dovecot_enabled = os.path.isfile(CONFIG_DOVECOT) and \
            'dovecot_imunify' in open(CONFIG_DOVECOT).read()
        sshd_enabled = any('pam_imunify.so' in open(conf).read()
                           for conf in conffiles)

        if dovecot_enabled:
            toggle_dovecot_support(False, output)

        if sshd_enabled:
            toggle_sshd_support(conffiles, False, output)

        output.status_changed({
            'sshd': (sshd_enabled, False),
            'dovecot': (dovecot_enabled, False)
        })

    @staticmethod
    def disable(conffiles, output=None):
        if not output:
            output = Output()

        sshd_enabled = any('pam_imunify.so' in open(conf).read()
                           for conf in conffiles)

        if sshd_enabled:
            toggle_sshd_support(conffiles, False, output)

        output.status_changed({'sshd': (sshd_enabled, False)})

    @staticmethod
    def disable_dovecot(conffiles, output=None):
        if not output:
            output = Output()

        dovecot_enabled = os.path.isfile(CONFIG_DOVECOT) and \
            'dovecot_imunify' in open(CONFIG_DOVECOT).read()
        if dovecot_enabled:
            toggle_dovecot_support(False, output)

        output.status_changed({'dovecot': (dovecot_enabled, False)})

    @classmethod
    def status(cls, conffiles, output=None):
        if not output:
            output = Output()

        dovecot_enabled = os.path.isfile(CONFIG_DOVECOT) and \
            'dovecot_imunify' in open(CONFIG_DOVECOT).read()
        sshd_enabled = any('pam_imunify.so' in open(conf).read()
                           for conf in conffiles)

        if dovecot_enabled or sshd_enabled:
            cls._cphulk_check(output)

        output.status({'sshd': sshd_enabled, 'dovecot': dovecot_enabled})

    @staticmethod
    def _cphulk_check(output=None):
        if not os.path.isfile("/usr/sbin/whmapi1"):
            return
        if not read_config().getboolean('verbose'):
            return

        if not output:
            output = Output()

        proc = subprocess.run(['/usr/sbin/whmapi1',
                               'servicestatus',
                               'service=cphulkd'],
                              stdin=subprocess.DEVNULL,
                              stdout=subprocess.PIPE)
        if proc.returncode != 0:
            # we expect err dump is printed to stderr
            return

        try:
            status = yaml.load(proc.stdout)
            if status['data']['service'][0]['enabled']:
                output.warning("cPHulk is enabled", file=sys.stderr)
        except (yaml.YAMLError, IndexError, KeyError) as e:
            output.warning("whmapi error:", e, file=sys.stderr)


if __name__ == '__main__':
    output = Output()
    try:
        cmd = getattr(Cmd, sys.argv[1].replace('-', '_'))
        if len(sys.argv) > 2 and sys.argv[2] == '--yaml':
            output = YamlOutput()
    except (AttributeError, IndexError):
        sys.exit("Usage: %s {enable|enable-dovecot|"
                 "disable|disable-dovecot|disable-all|status}" % sys.argv[0])
    else:
        try:
            cmd(detect_conffiles(output), output)
        except SystemExit:
            del output
            raise
