#!/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 glob
import re
import shutil
import shlex
import tarfile
import argparse
import subprocess

def run(cmd, shell=False):
    if not shell:
        cmd = shlex.split(cmd)
    return subprocess.check_call(cmd, shell=shell)

def get_kver(pattern):
    for k in glob.glob(pattern):
        return re.sub(r'^/boot/vmlinuz-(.+)$', r'\1', k)

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    parser = argparse.ArgumentParser(description='Construct AMI')
    parser.add_argument('--localrpm', action='store_true', default=False,
                        help='deploy locally built rpms')
    parser.add_argument('--repo',
                        help='repository for both install and update, specify .repo/.list file URL')
    parser.add_argument('--repo-for-install',
                        help='repository for install, specify .repo/.list file URL')
    parser.add_argument('--repo-for-update',
                        help='repository for update, specify .repo/.list file URL')
    args = parser.parse_args()

    if args.repo:
        args.repo_for_install = args.repo_for_update = args.repo

    if not args.localrpm and not args.repo_for_install:
        print('Error: need to specify --localrpm or --repo/--repo-for-install')
        sys.exit(1)

    run('yum update -y')

    run('curl -L -o /etc/yum.repos.d/scylla-ami-drivers.repo https://copr.fedorainfracloud.org/coprs/scylladb/scylla-ami-drivers/repo/epel-7/scylladb-scylla-ami-drivers-epel-7.repo')
    if args.repo_for_install:
        run('curl -L -o /etc/yum.repos.d/scylla_install.repo {REPO_FOR_INSTALL}'.format(REPO_FOR_INSTALL=args.repo_for_install))

    if args.localrpm:
        rpms = glob.glob('/home/centos/scylla*.*.rpm')
        run('yum install -y {}'.format(' '.join(rpms)))
    else:
        run('yum install -y scylla-ami')

    if args.repo_for_install:
        os.remove('/etc/yum.repos.d/scylla_install.repo')
    if args.repo_for_update:
        run('curl -L -o /etc/yum.repos.d/scylla.repo {REPO_FOR_UPDATE}'.format(REPO_FOR_UPDATE=args.repo_for_update))

    run('systemctl daemon-reload')
    run('systemctl enable scylla-ami-setup.service')
    with open('/etc/cloud/cloud.cfg') as f:
        lines = f.readlines()
    with open('/etc/cloud/cloud.cfg', 'w') as f:
        for l in lines:
            if not re.match(r'^ - mounts\n$', l):
                f.write(l)
    run('/usr/lib/scylla/scylla_setup --ntp-domain amazon --no-coredump-setup --no-sysconfig-setup --no-raid-setup --no-io-setup --no-bootparam-setup --no-ec2-check')
    run('/usr/lib/scylla/scylla_sysconfig_setup --ami')
    run('/usr/lib/scylla/scylla_bootparam_setup --ami')
    os.remove('/home/centos/.ssh/authorized_keys')
    os.remove('/var/lib/scylla-housekeeping/housekeeping.uuid')

    run('yum -y update kernel')
    run('yum -y install dkms git grubby kernel-devel')
    run('rpm -e kernel-$(uname -r)', shell=True)
    run('rpm --import https://www.elrepo.org/RPM-GPG-KEY-elrepo.org')
    run('rpm -Uvh http://www.elrepo.org/elrepo-release-7.0-3.el7.elrepo.noarch.rpm')
    run('yum -y --enablerepo=elrepo-kernel install kernel-ml kernel-ml-devel')
    os.remove('/etc/udev/rules.d/80-net-name-slot.rules')
    with open('/etc/default/grub') as f:
        grub = f.read()
    grub = re.sub(r'^GRUB_CMDLINE_LINUX="(.+)"$', r'GRUB_CMDLINE_LINUX="\1 net.ifnames=0"', grub, flags=re.MULTILINE)
    with open('/etc/default/grub', 'w') as f:
        f.write(grub)
    run('grub2-mkconfig -o /boot/grub2/grub.cfg')
    c7kver = get_kver('/boot/vmlinuz-*el7.x86_64')
    mlkver = get_kver('/boot/vmlinuz-*el7.elrepo.x86_64')
    run('grubby --grub2 --set-default /boot/vmlinuz-{mlkver}'.format(mlkver=mlkver))
    with open('/etc/yum.conf', 'a') as f:
        f.write('exclude=kernel kernel-devel')

    run('dracut --verbose --add-drivers "ixgbevf nvme ena" --force --kver {c7kver}'.format(c7kver=c7kver))
    run('dracut --verbose --add-drivers "ixgbevf nvme ena" --force --kver {mlkver}'.format(mlkver=mlkver))

    with open('/etc/cloud/cloud.cfg') as f:
        cfg = f.read()
    cfg2 = re.sub('^ssh_deletekeys:   0', 'ssh_deletekeys:   1', cfg, flags=re.MULTILINE)
    with open('/etc/cloud/cloud.cfg', 'w') as f:
        f.write(cfg2)
