#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2020 ScyllaDB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import glob
import re
import shutil
import shlex
import tarfile
import argparse
import subprocess
import platform
try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO

my_env = os.environ.copy()
my_env['DEBIAN_FRONTEND']='noninteractive'
def run(cmd, shell=False):
    if not shell:
        cmd = shlex.split(cmd)
    return subprocess.check_call(cmd, shell=shell, env=my_env)

def out(cmd, shell=False):
    if not shell:
        cmd = shlex.split(cmd)
    return subprocess.check_output(cmd, shell=shell, env=my_env).decode('utf-8')

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

def is_centos():
    return os.path.exists('/etc/redhat-release')

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    distro = 'centos' if is_centos() else 'ubuntu'
    homedir = os.path.abspath(os.path.join(__file__, os.pardir))
    parser = argparse.ArgumentParser(description='Construct AMI')
    parser.add_argument('--localrpm', action='store_true', default=False,
                        help='deploy locally built rpms')
    parser.add_argument('--product',
                        help='name of the product', default='scylla')
    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')
    parser.add_argument('--target-cloud', choices=['aws', 'gce', 'azure'], help='specify target cloud')
    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)

    if distro == 'centos':
        run('yum update -y')

        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('{}/scylla*.*.rpm'.format(homedir))
            run('yum install -y {}'.format(' '.join(rpms)))
        else:
            run('yum install -y {0}-machine-image {0}-debuginfo'.format(args.product))

        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))
    else:
        run('apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 5e08fbd8b5d6ec9c')
        if args.repo_for_install:
            run('curl -L -o /etc/apt/sources.list.d/scylla_install.list {REPO_FOR_INSTALL}'.format(REPO_FOR_INSTALL=args.repo_for_install))
        elif args.localrpm:
            with open('/etc/apt/sources.list.d/scylla_install.list', 'w') as f:
                f.write('deb file:{} ./'.format(homedir))
        else:
            print('no scylla package found.')
            sys.exit(1)

        run('apt-get update --allow-insecure-repositories -y')
        run('apt-get upgrade -y')
        run('apt-get install -y --auto-remove --allow-unauthenticated {0}-machine-image {0}-server-dbg systemd-coredump'.format(args.product))

        os.remove('/etc/apt/sources.list.d/scylla_install.list')
        if args.repo_for_update:
            run('curl -L -o /etc/apt/sources.list.d/scylla.list {REPO_FOR_UPDATE}'.format(REPO_FOR_UPDATE=args.repo_for_update))

        # disable unattended-upgrades
        run('apt-get purge -y unattended-upgrades update-notifier-common')

    if args.target_cloud == 'aws':
        # use easy_install instead of pip, because easy_install requires
        # fewer dependencies, no need to install gcc, it makes our AMI smaller.
        if distro == 'centos':
            run('yum install -y epel-release')
            run('yum install -y pystache python-daemon python-setuptools')
            # Default install prefix on CentOS7 is /usr not /usr/local, but we need
            # to align path for executable, so specify script-dir to /usr/local/bin
            run('python -m easy_install --script-dir /usr/local/bin https://s3.amazonaws.com/cloudformation-examples/aws-cfn-bootstrap-latest.tar.gz')
            pyver = ''
        else:
            # XXX: we want install dependencies from distro repo, but
            # python-daemon-2.2 on Ubuntu 20.04 is too new for aws-cfn-bootstrap,
            # so install python-daemon-2.1 from pypi
            run('apt-get install -y python3-pystache python3-lockfile python3-setuptools')
            run('python3 -m easy_install https://s3.amazonaws.com/cloudformation-examples/aws-cfn-bootstrap-py3-latest.tar.gz')
            pyver = '3'
        cfnpath = out('python{} -c "import cfnbootstrap as _; print(_.__path__[0])"'.format(pyver))
        initpath = os.path.abspath('{}/../init/{}/cfn-hup'.format(cfnpath, 'redhat' if distro == 'centos' else 'ubuntu'))
        shutil.copyfile(initpath, '/etc/init.d/cfn-hup')
        os.chmod('/etc/init.d/cfn-hup', 0o755)

        setup_opt = '--ntp-domain amazon'
        sysconfig_opt = ''
    elif args.target_cloud == 'gce':
        setup_opt = ''
        sysconfig_opt = '--disable-writeback-cache'
    elif args.target_cloud == 'azure':
        setup_opt = ''
        sysconfig_opt = '--disable-writeback-cache'

    run('systemctl daemon-reload')
    run('systemctl enable scylla-image-setup.service')
    run('/opt/scylladb/scripts/scylla_setup --no-coredump-setup --no-sysconfig-setup --no-raid-setup --no-io-setup --no-bootparam-setup --no-ec2-check --no-swap-setup --no-cpuscaling-setup')

    # The difference comes from two different services to configure CPU scaling.
    # On CentOS 'cpupower' fails when CPU scaling is not supported, so we can't
    # enable it here.
    # On Ubuntu 'cpufrequtils' never fails even CPU scaling is not supported,
    # so we want to enable it here.
    if distro == 'centos':
        run('yum install -y cpupowerutils')
    else:
        run('/opt/scylladb/scripts/scylla_cpuscaling_setup --force')

    run('/opt/scylladb/scripts/scylla_sysconfig_setup --ami --set-clocksource {}'.format(sysconfig_opt))
    run('/opt/scylladb/scripts/scylla_bootparam_setup --ami')
    os.remove('{}/.ssh/authorized_keys'.format(homedir))
    os.remove('/var/lib/scylla-housekeeping/housekeeping.uuid')

    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)
    if distro == 'centos':
        run('grub2-mkconfig -o /boot/grub2/grub.cfg')
    else:
        run('grub-mkconfig -o /boot/grub/grub.cfg')

    if distro == 'centos':
        profile = '/etc/skel/.bash_profile'
    else:
        profile = '/etc/skel/.profile'
    with open(profile, 'a') as f:
        f.write('\n\n/opt/scylladb/scylla-machine-image/scylla_login\n')

    # On AWS, ssh user is statically created at AMI building time, so we need to
    # change it to 'scyllaadm`.
    # However, on GCE and Azure ssh user is dynamically created at instance startup
    # time, and username is specified while launching the instance, we have nothing
    # to do.
    if args.target_cloud == 'aws':
        with open('/etc/cloud/cloud.cfg') as f:
            lines = f.readlines()
        s = StringIO()
        for l in lines:
            match_groups = re.match(r'^\s+groups: \[(.+)\]\n$', l)
            if match_groups:
                groups = match_groups.group(1).replace(' ', '')
            if not re.match(r'^ - mounts\n$', l):
                s.write(l)
        cfg = s.getvalue()
        s.close()
        if distro == 'centos':
            cfg = re.sub('^ssh_deletekeys:   0', 'ssh_deletekeys:   1', cfg, flags=re.MULTILINE)
            cfg = re.sub('^    name: centos', '    name: scyllaadm', cfg, flags=re.MULTILINE)
        else:
            cfg = re.sub('^preserve_hostname: false', 'preserve_hostname: false\n\nssh_deletekeys: true', cfg, flags=re.MULTILINE)
            cfg = re.sub('^     name: ubuntu', '     name: scyllaadm', cfg, flags=re.MULTILINE)
        with open('/etc/cloud/cloud.cfg', 'w') as f:
            f.write(cfg)
        # before deleting home directory, need to change current directory
        os.chdir('/tmp')
        run('userdel -r -f {}'.format(distro))
        run('cloud-init clean')
        run('cloud-init init')
        for skel in glob.glob('/etc/skel/.*'):
            shutil.copy(skel, '/home/scyllaadm')
            os.chown(skel, 1000, 1000)
        run('useradd -o -u 1000 -g scyllaadm -G {} -s /bin/bash -d /home/scyllaadm centos'.format(groups))
        run('groupadd -o -g 1000 centos')
        os.symlink('/home/scyllaadm', '/home/centos')
        if distro == 'ubuntu':
            run('useradd -o -u 1000 -g scyllaadm -G {} -s /bin/bash -d /home/scyllaadm ubuntu'.format(groups))
            run('groupadd -o -g 1000 ubuntu')
            os.symlink('/home/scyllaadm', '/home/ubuntu')

    if args.target_cloud == 'azure':
        with open('/etc/hosts', 'a') as f:
            f.write('\n\n169.254.169.254    metadata.azure.internal\n')

    # GCP CentOS is 8, don't need to install LTS kernel
    if args.target_cloud == 'aws' and distro == 'centos':
        # We need to install LTS kernel for the AMI.
        # However, elrepo-kernel only provides x86_64 package,
        # and centos-kernel (altarch/7/kernel) only provides aarch64 package.
        # So we have to use different kernel here.
        if platform.machine() == 'x86_64':
            run('yum -y install grubby')
            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')
            mlkver = get_kver('/boot/vmlinuz-*el7.elrepo.x86_64')
            run('dracut --add-drivers "xen-blkfront xen-netfront ixgbevf ena nvme" --force /boot/initramfs-{mlkver}.img {mlkver}'.format(mlkver=mlkver))
            run('grubby --grub2 --set-default /boot/vmlinuz-{mlkver}'.format(mlkver=mlkver))
            with open('/etc/yum.conf', 'a') as f:
                f.write(u'exclude=kernel kernel-devel')
        elif platform.machine() == 'aarch64':
            run('yum -y --enablerepo=centos-kernel update kernel')
