#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2018-present ScyllaDB
#

#
# SPDX-License-Identifier: AGPL-3.0-or-later

import os
import sys
import argparse
import subprocess
import time
import tempfile
import shutil
from scylla_util import *
from subprocess import run


def has_zstd():
    coredump_version = out('coredumpctl --version').split('\n')
    features = coredump_version[1].split(' ')
    return '+ZSTD' in features

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)

    if is_debian_variant() or is_suse_variant():
        if not shutil.which('coredumpctl'):
            pkg_install('systemd-coredump')

    parser = argparse.ArgumentParser(description='Optimize coredump settings for Scylla.')
    parser.add_argument('--dump-to-raiddir', action='store_true', default=False,
                        help='store coredump to /var/lib/scylla')
    # Enable compress by default when zstd support available
    parser.add_argument('--compress', action='store_true', default=has_zstd(),
                        help='enable compress on systemd-coredump')
    args = parser.parse_args()

    # Seems like a specific version of systemd package on RHEL9 has a bug on
    # SELinux configuration, it introduced "systemd-container-coredump" module
    # to provide rule for systemd-coredump but not enabled by default.
    # We have to manually load it, otherwise it causes permission error.
    # (#19325)
    if is_redhat_variant() and distro.major_version() == '9':
        if not shutil.which('getenforce'):
            pkg_install('libselinux-utils')
        if not shutil.which('semodule'):
            pkg_install('policycoreutils')
        enforce = out('getenforce')
        if enforce != "Disabled":
            if os.path.exists('/usr/share/selinux/packages/targeted/systemd-container-coredump.pp.bz2'):
                modules = out('semodule -l')
                match = re.match(r'^systemd-container-coredump$', modules, re.MULTILINE)
                if not match:
                    run('semodule -v -i /usr/share/selinux/packages/targeted/systemd-container-coredump.pp.bz2', shell=True, check=True)
                    run('semodule -v -e systemd-container-coredump', shell=True, check=True)

    # abrt-ccpp.service needs to stop before enabling systemd-coredump,
    # since both will try to install kernel coredump handler
    # (This will only requires for abrt < 2.14)
    if systemd_unit.available('abrt-ccpp.service'):
        abrt_ccpp = systemd_unit('abrt-ccpp.service')
        abrt_ccpp.disable()
        abrt_ccpp.stop()

# Gentoo may uses OpenRC
    if is_gentoo():
        run('sysctl -p /etc/sysctl.d/99-scylla-coredump.conf', shell=True, check=True)
# Other distributions can use systemd-coredump, so setup it
    else:
        if is_suse_variant():
            systemd_unit('systemd-coredump.socket').restart()
        # Some older distribution does not have this unit
        if systemd_unit.available('systemd-coredump@.service'):
            dropin = '''
[Service]
RuntimeMaxSec=infinity
TimeoutSec=infinity
'''[1:-1]
            os.makedirs('/etc/systemd/system/systemd-coredump@.service.d', exist_ok=True)
            with open('/etc/systemd/system/systemd-coredump@.service.d/timeout.conf', 'w') as f:
                f.write(dropin)
        conf_data = '''
[Coredump]
Storage=external
Compress={compress}
ProcessSizeMax=1024G
ExternalSizeMax=1024G
'''[1:-1].format(compress = 'yes' if args.compress else 'no')
        with open('/etc/systemd/coredump.conf', 'w') as f:
            conf = f.write(conf_data)
        if args.dump_to_raiddir:
            dot_mount = '''
[Unit]
Description=Save coredump to scylla data directory
Conflicts=umount.target
Before=local-fs.target scylla-server.service
DefaultDependencies=no

[Mount]
What=/var/lib/scylla/coredump
Where=/var/lib/systemd/coredump
Type=none
Options=bind

[Install]
WantedBy=local-fs.target scylla-server.service
'''[1:-1]
            with open('/etc/systemd/system/var-lib-systemd-coredump.mount', 'w') as f:
                f.write(dot_mount)
            os.makedirs('/var/lib/scylla/coredump', exist_ok=True)
            systemd_unit.reload()
            systemd_unit('var-lib-systemd-coredump.mount').enable()
            systemd_unit('var-lib-systemd-coredump.mount').start()
        if os.path.exists('/usr/lib/sysctl.d/50-coredump.conf'):
            run('sysctl -p /usr/lib/sysctl.d/50-coredump.conf', shell=True, check=True)
        else:
            with open('/etc/sysctl.d/99-scylla-coredump.conf', 'w') as f:
                f.write('kernel.core_pattern=|/usr/lib/systemd/systemd-coredump %p %u %g %s %t %e')
            run('sysctl -p /etc/sysctl.d/99-scylla-coredump.conf', shell=True, check=True)

        fp = tempfile.NamedTemporaryFile()
        fp.write(b'ulimit -c unlimited\n')
        fp.write(b'kill -SEGV $$\n')
        fp.flush()
        p = subprocess.Popen(['/bin/bash', fp.name], stdout=subprocess.PIPE)
        pid = p.pid
        p.wait()
        fp.close()

        print('Generating coredump to test systemd-coredump...\n')
        # need to wait for systemd-coredump to complete collecting coredump
        time.sleep(3)
        try:
            coreinfo = run('coredumpctl --no-pager --no-legend info {}'.format(pid), shell=True, check=True, capture_output=True, encoding='utf-8').stdout.strip()
        except subprocess.CalledProcessError:
            print('Unable to detect coredump, failed to configure systemd-coredump.')
            sys.exit(1)

        print(coreinfo)
        print()

        # "coredumpctl info" behavior had been changed since systemd-v232,
        # we need to support both version.
        #
        # Before systemd-v232, it was simple.
        # It print 'Coredump' field only when the coredump exists on filesystem.
        # Otherwise print nothing.
        #
        # After the change made on systemd-v232, it become more complex.
        # It always print 'Storage' field even the coredump does not exists.
        # Not just available/unavailable, it describe more:
        #  - Storage: none
        #  - Storage: journal
        #  - Storage: /path/to/file (inaccessible)
        #  - Storage: /path/to/file
        #
        # After systemd-v248, available coredump file output changed like this:
        #  - Storage: /path/to/file (present)
        # We need to support both versions.
        #
        # reference: https://github.com/systemd/systemd/commit/47f50642075a7a215c9f7b600599cbfee81a2913

        corefail = False
        res = re.findall(r'Storage: (\S+)(?: \(.+\))?$', coreinfo, flags=re.MULTILINE)
        # v232 or later
        if res:
            corepath = res[0]
            if corepath == 'none' or corepath == 'journal' or corepath.endswith('(inaccessible)'):
                corefail = True
        # before v232
        else:
            res = re.findall(r'Coredump: (.*)$', coreinfo, flags=re.MULTILINE)
            if res:
                corepath = res[0]
            else:
                corefail = True

        if not corefail:
            try:
                os.remove(corepath)
            except FileNotFoundError:
                corefail = True

        if corefail:
            print('Does not able to detect coredump file, failed to configure systemd-coredump.')
            sys.exit(1)

        print('\nsystemd-coredump is working finely.')
