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

#
# This file is part of Scylla.
#
# Scylla is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Scylla is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Scylla.  If not, see <http://www.gnu.org/licenses/>.

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

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    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')
    parser.add_argument('--compress', action='store_true', default=False,
                        help='enable compress on systemd-coredump')
    args = parser.parse_args()

# Gentoo may uses OpenRC
    if is_gentoo_variant():
        run('sysctl -p /etc/sysctl.d/99-scylla-coredump.conf')
# Other distributions can use systemd-coredump, so setup it
    else:
        if is_debian_variant():
            run('apt-get install -y systemd-coredump')
        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=scylla-server.service
After=local-fs.target

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

[Install]
WantedBy=multi-user.target
'''[1:-1]
            with open('/etc/systemd/system/var-lib-systemd-coredump.mount', 'w') as f:
                f.write(dot_mount)
            makedirs('/var/lib/scylla/coredump')
            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')
        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')

        fp = tempfile.NamedTemporaryFile()
        fp.write(b'kill -SEGV $$')
        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:
            run('coredumpctl --no-pager --no-legend info {}'.format(pid))
            print('\nsystemd-coredump is working finely.')

            # get last coredump generated by bash and remove it, ignore inaccessaible ones
            corefile = out(cmd=r'coredumpctl -1 --no-legend dump 2>&1 | grep "bash" | '
                               r'grep -v "inaccessible" | grep "Storage:\|Coredump:"',
                           shell=True, exception=False)
            if corefile:
                corefile = corefile.split()[-1]
                run('rm -f {}'.format(corefile))
        except subprocess.CalledProcessError as e:
            print('Does not able to detect coredump, failed to configure systemd-coredump.')
            sys.exit(1)
