#!/usr/bin/python3
#
# 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 re
from scylla_util import *

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

    parser = argparse.ArgumentParser(description='Optimize NTP client setting for Scylla.')
    if is_redhat_variant():
        parser.add_argument('--subdomain',
                            help='specify subdomain of pool.ntp.org (ex: centos, fedora or amazon)')
    args = parser.parse_args()

    if is_debian_variant():
        run('apt-get install -y ntp ntpdate')
        with open('/etc/ntp.conf') as f:
            conf = f.read()
        if not re.match(r'^server ', conf, flags=re.MULTILINE):
            conf2 = re.findall(r'^(?!pool).*\n', conf, flags=re.MULTILINE)
            with open('/etc/ntp.conf', 'w') as f:
                f.writelines(conf2)
                f.write('server ntp.ubuntu.com iburst')
        if is_systemd():
            ntp = systemd_unit('ntp.service')
            ntp.stop()
            # ignore error, ntpd may able to adjust clock later
            run('ntpdate ntp.ubuntu.com', exception=False)
            ntp.start()
        else:
            run('service ntp stop')
            run('ntpdate ntp.ubuntu.com')
            run('service ntp start')

    if is_gentoo_variant():
        run('emerge -uq net-misc/ntp')
        with open('/etc/ntp.conf') as f:
            conf = f.read()
        match = re.search(r'^server\s+(\S*)(\s+\S+)?', conf, flags=re.MULTILINE)
        server = match.group(1)
        if is_systemd():
            sntpd = systemd_unit('sntpd.service')
            sntpd.stop()
            run('ntpdate {}'.format(server))
            sntpd.enable()
            sntpd.start()
        else:
            run('rc-service ntpd stop', exception=False)
            # ignore error, ntpd may able to adjust clock later
            run('ntpdate {}'.format(server), exception=False)
            run('rc-update add ntpd default')
            run('rc-service ntpd start')

    if is_redhat_variant():
        run('yum install -y ntp ntpdate')
        with open('/etc/ntp.conf') as f:
            conf = f.read()
        if args.subdomain:
            conf2 = re.sub(r'server\s+([0-9]+)\.(\S+)\.pool\.ntp\.org', 'server \\1.{}.pool.ntp.org'.format(args.subdomain), conf, flags=re.MULTILINE)
            with open('/etc/ntp.conf', 'w') as f:
                f.write(conf2)
            conf = conf2
        match = re.search(r'^server\s+(\S*)(\s+\S+)?', conf, flags=re.MULTILINE)
        server = match.group(1)
        ntpd = systemd_unit('ntpd.service')
        ntpd.stop()
        # ignore error, ntpd may able to adjust clock later
        run('ntpdate {}'.format(server), exception=False)
        ntpd.enable()
        ntpd.start()
