#!/usr/bin/python3

#
# Copyright (c) 2020-2021 Virtuozzo International GmbH. All rights reserved.
#
# Our contact details: Virtuozzo International GmbH, Vordergasse 59, 8200
# Schaffhausen, Switzerland.

import argparse
import subprocess
import datetime
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool
import threading
import sys
import resource

# It is ok for these packages to be added due to differences between vzlinux8 and centos8/almalinux8 templates
# and default vzlinux requirements
ADDED_PKGS_IGNORE = ['annobin.x86_64', 'cpp.x86_64', 'dnf-plugins-core.noarch', 'fstrm.x86_64', 'gcc.x86_64',
                 'glibc-devel.x86_64', 'glibc-headers.x86_64', 'hwdata.noarch', 'isl.x86_64', 'kernel-headers.x86_64',
                 'libgomp.x86_64', 'libibverbs.x86_64', 'libmpc.x86_64', 'libnl3.x86_64', 'libpkgconf.x86_64',
                 'libxcrypt-devel.x86_64', 'lmdb-libs.x86_64', 'pciutils-libs.x86_64', 'pciutils.x86_64', 'pkgconf-m4.noarch',
                 'pkgconf-pkg-config.x86_64', 'pkgconf.x86_64', 'protobuf-c.x86_64', 'python3-dateutil.noarch',
                 'python3-dnf-plugins-core.noarch', 'rdma-core.x86_64', 'vzlinux-release.x86_64', 'zstd.x86_64',
                 'vzlinux-logos.noarch', 'vzlinux-logos-httpd.noarch',
                 ]

# The following packages are ok to drop - they are replaced by vzlinux ones
DROPPED_PKGS_IGNORE = ['centos-linux-release.noarch', 'centos-linux-repos.noarch', 'centos-logos-httpd.noarch',
                        'centos-logos.noarch', 'centos-gpg-keys.noarch', 'almalinux-release.noarch', 'almalinux-repos.noarch',
                        'almalinux-logos-httpd.noarch', 'almalinux-logos.noarch', 'almalinux-gpg-keys.noarch',
                        'almalinux-release.x86_64', 'centos-indexhtml.noarch', 'centos-stream-repos.noarch', 'centos-stream-release.noarch'
                        ]

# Packages to be added after upgrade from CentOS 7
POST_UPGRADE_ADD_PKGS = ['binutils', 'perl', 'python2', 'python2-chardet']

# Some C7 templates are not available in C8 (though some of their packages might still exist)
REMOVED_TMPL_IGNORE = ['tomcat', 'cloud-init']

# There is no harm if some processes became active/inacitve after upgrade.
# mandb & logrotate can be launched by cron;
# 'sh' also comes from rpm_clean cron task.
# 'systemct' (there is no typo here) looks like a ghost in vzps output
CHANGED_PS_IGNORE = ['mandb', 'sh', 'systemct', 'disp_helper']

# Forbid update if found installed package containing one of the following strings in name
BLOCKER_PKGS = {'plesk': 'Plesk', 'cpanel': 'cPanel'}

# Minimum free space we want to see inside the container
SPACE_LIMIT = 1000000

# We also supports all remi* repos, but we just check for 'remi' prefix in the code
SUPPORTED_REPOS = ['ct-preset', 'appstream', 'baseos', 'extras', 'powertools', 'epel', 'epel-modular',
                    'base/7/x86_64', 'extras/7/x86_64', 'updates/7/x86_64']
'''
Simple log wrapper to print messages to eithe STDOUT or to logfile
'''
def log_info(msg, ct_log=None):
    global logfile
    global lock

    if not ct_log:
        lock.acquire()
    for l in str(msg).split('\n'):
        if 'warning! rpmdb:' in l or 'ERROR: ld.so:' in l:
            continue
        if ct_log:
            ct_log.write(l + "\n")
        elif logfile:
            logfile.write(l + "\n")
        print(l)
        sys.stdout.flush()
    if not ct_log:
        lock.release()


def parse_command_line():
    global args

    parser = argparse.ArgumentParser(description='VzLinux Converter')
    subs = parser.add_subparsers(dest='list or convert')
    subs.required = True

    list_parser = subs.add_parser('list', help='List convertable CentOS 7 / CentOS 8 / AlmaLinux 8 containers')
    list_parser.set_defaults(func=get_upgradable)

    upgrade_parser = subs.add_parser('convert', help='Convert the specified containers to VzLinux 8')
    upgrade_parser.add_argument('CT', metavar='CT', nargs='+', help='UUID, CTID, or name of the container to convert')
    upgrade_parser.add_argument('--dry-run', action='store_true', help='Check that conversion is possible, do not actually perform it')
    upgrade_parser.add_argument('-q', '--quiet', action='store_true', help='Be quiet')
    upgrade_parser.add_argument('-v', '--verbose', action='store_true', help='Be verbose')
    upgrade_parser.add_argument('--log', help='Dump all messages to the specified log file. Detailed messages for every container will be dumped to separate files with the same prefix')
    upgrade_parser.add_argument('--parallel', metavar='parallel', type=int, choices=range(1, 101), nargs='?', help='The number of concurrent conversions to perform')
    upgrade_parser.add_argument('--strict', action='store_true', help='Treat some of the precheck warnings as errors that block conversion')
    upgrade_parser.set_defaults(func=process_cts)

    args = parser.parse_args()


'''
Check CT config - if OSTEMPLATE is set to supported distro
Return value:
    0 - distro is not supported
    1 - distro is supported, belongs to CentOS 8 family and can be converted directly to vzlinux 8
    2 - distro is supported, belongs to CentOS 7 family and requires preliminary upgrade to CentOS 8
'''
def check_config(ctid):
    vz_private = None
    try:
        with open("/etc/vz/vz.conf") as f:
            for l in f.readlines():
                if l.startswith("VE_PRIVATE"):
                    vz_private = l.strip().split("=")[1].replace('/$VEID', '').replace('"', '').replace("'", '')
                    break
    except Exception as e:
        log_info(ctid + ": Unable to check container config: " + str(e))
        return 0

    try:
        f = open(vz_private + "/" + ctid + "/ve.conf", "r")
    except Exception as e:
        log_info(ctid + ": Unable to check container config: " + str(e))
        return 0

    c7_available = check_vzl8_tmp_ver()
    if not c7_available:
        print("WARNING: Your version of vzlinux-8-x86_64-ez template is too old, upgrade of CentOS 7 CTs is not available.\n")

    for l in f.readlines():
        if 'OSTEMPLATE=".centos-8.stream-x86_64"' in l or 'OSTEMPLATE="centos-8.stream-x86_64"' in l \
                or 'OSTEMPLATE=".centos-8.stream"' in l or 'OSTEMPLATE="centos-8.stream"' in l \
                or 'OSTEMPLATE=".centos-8-x86_64"' in l or 'OSTEMPLATE="centos-8-x86_64"' in l \
                or 'OSTEMPLATE=".centos-8"' in l or 'OSTEMPLATE="centos-8"' in l \
                or 'OSTEMPLATE=".rockylinux-8-x86_64"' in l or 'OSTEMPLATE="rockylinux-8-x86_64"' in l \
                or 'OSTEMPLATE=".almalinux-8-x86_64"' in l or 'OSTEMPLATE="almalinux-8-x86_64"' in l \
                or 'OSTEMPLATE=".almalinux-8"' in l or 'OSTEMPLATE="almalinux-8"' in l:
            f.close()
            return 1

        if c7_available:
            if 'OSTEMPLATE=".centos-7-x86_64"' in l or 'OSTEMPLATE="centos-7-x86_64"' in l \
                    or 'OSTEMPLATE=".centos-7"' in l or 'OSTEMPLATE="centos-7"' in l :
                f.close()
                return 2

    f.close()
    return 0


'''
Check if we have enough free space inside container.
Return False if not (or when can't check), True if yes
'''
def check_space(ctid):
    df_out = subprocess.check_output(['/sbin/vzctl', 'exec', ctid, 'df', '--output=avail', '/'])
    for l in df_out.decode('utf-8').split("\n"):
        if "vail" in l:
            continue
        free_space = int(l)
        if free_space < SPACE_LIMIT:
            log_info(ctid + ": Not enough free space in the container, at least 1 GB is required")
            return False
        else:
            return True
    log_info(ctid + ": Unable to check free space in the container!")
    return False

'''
Check if we have enabled repos not supported by upgrade.
Return False if yes (or when can't check), True if no
'''
def check_repos(ctid, distro_kind=1):
    try:
        if distro_kind == 1:
            dnf_out = subprocess.check_output(['/sbin/vzctl', 'exec', ctid, 'dnf', 'repolist', 'enabled'])
        else:
            dnf_out = subprocess.check_output(['/sbin/vzctl', 'exec', ctid, 'yum', 'repolist', 'enabled'])
    except Exception as e:
        log_info(ctid + ": Unable to check the repositories in the container! " + str(e))
        return False

    for l in dnf_out.decode('utf-8').split("\n"):
        if "repo name" in l:
            continue
        repo_id = l.split(" ")[0]
        if repo_id and repo_id not in SUPPORTED_REPOS:
            if repo_id.startswith('remi'):
                continue
            log_info(ctid + ": The following unsupported repository is enabled in the container: " + str(repo_id))
            return False

    return True

'''
Get list of ports open inside CT.
We get list of ports using lsof (since netstat is not available by default)
For every port we save command, protocol, node and port number
'''
def get_open_ports(ctid, distro_kind=1):
    # lsof has different locations in CentOS 7 and 8
    if distro_kind == 1:
        lsof_path = '/bin/lsof'
    else:
        lsof_path = '/sbin/lsof'

    l = subprocess.check_output(['/sbin/vzctl', 'exec', ctid, lsof_path, '-Pi'])
    all_ports = []
    for p in l.decode('utf-8').split('\n'):
        data = p.split()
        if len(data) > 8:
            all_ports.append((data[0], data[4], data[7], data[8]))
    return all_ports

def try_stop_ct(ctid):
    try:
        res = subprocess.check_output(['/sbin/vzctl', 'stop', ctid]).decode('utf-8')
        log_info(res)
    except Exception as e:
        log_info(ctid + ": Failed to stop the container: " + str(e))

'''
Perform conversion of containers specified in cmdline
'''
def process_cts():
    global args
    global logfile
    global lock

    # Increase 'ulimit -n' for ancestors
    resource.setrlimit(resource.RLIMIT_NOFILE, (131072, 131072))
    lock = threading.Lock()
    logfile = None
    if args.log:
        logfile = open(args.log, "w")

    if args.parallel:
        pool = ThreadPool(args.parallel)
    else:
        pool = ThreadPool(1)

    results = pool.map(process_single_ct, args.CT)

    pool.close()
    pool.join()

    if args.log:
        logfile.close()


'''
Upgrade CentOS 7 CT to CentOS 8 one
'''
def upgrade_c7_ct(ct):
    global args

    vzpkg_args = ['upgrade']
    if args.dry_run:
        vzpkg_args.append('-n')

    vzpkg_args.append(ct)

    pr = subprocess.Popen(['/sbin/vzpkg'] + vzpkg_args, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1)
    if args.log:
        ct_log = open(args.log + "." + ct, "w")
    else:
        ct_log = None
    while True:
        l = pr.stdout.readline()
        if l == '' and pr.poll() != None:
            break
        if l:
            log_info(l.strip(), ct_log)

    if pr.returncode and pr.returncode != 0:
        log_info("Failed to upgrade container from 7 to 8 family", ct_log)
        if args.log:
            ct_log.close()
        sys.exit(1)

    # In case of dry run, we only check upgrade itself
    if args.dry_run:
        if args.log:
            ct_log.close()
        return

    # Without dry run, we need to perform further update
    pr = subprocess.Popen(['/sbin/vzpkg', 'update', ct], universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1)
    while True:
        l = pr.stdout.readline()
        if l == '' and pr.poll() != None:
            break
        if l:
            log_info(l.strip(), ct_log)

    if pr.returncode and pr.returncode != 0:
        log_info("Failed to update container after conversion from 7 to 8 family", ct_log)
        if args.log:
            ct_log.close()
        sys.exit(1)

    for p in POST_UPGRADE_ADD_PKGS:
        pr = subprocess.check_output(['/sbin/vzpkg', 'install', ct, '-p', p])
        log_info(pr, ct_log)

    if args.log:
        ct_log.close()


'''
A thread function processing single CT
'''
def process_single_ct(ct):
    global args

    try:
        ctid = subprocess.check_output(['/sbin/vzlist', '-H', '-o', 'ctid', ct]).decode('utf-8')
    except Exception as e:
        log_info("Failed to get info for container %s" % ct)
        return

    ct = str(ctid).strip()

    distro_kind = check_config(ct)
    if not distro_kind:
        log_info(ct + ": Conversion aborted: Container's OS template is not supported")
        return

    need_stop = False
    ct_state = subprocess.check_output(['/sbin/vzctl', 'status', ct]).decode('utf-8')
    if not "running" in ct_state:
        log_info(ct + ": Container is stopped. Starting container...")
        try:
            res = subprocess.check_output(['/sbin/vzctl', 'start', ct, '--wait']).decode('utf-8')
            log_info(res)
        except Exception as e:
            log_info(ct + ": Failed to start: " + str(e))
            return
        need_stop = True

    old_pkg_list_raw = subprocess.check_output(['/sbin/vzpkg', 'list', '-p', ct]).decode('utf-8')
    proceed = True
    # TODO: check whether this works 
    for b in BLOCKER_PKGS:
        if b in old_pkg_list_raw:
            log_info(ct + ": " + "Conversion aborted: Software unsupported by VzLinux 8 detected: " + BLOCKER_PKGS[b])
            log_info("Please contact the software vendor and request VzLinux 8 support")
            proceed = False

    # DirectAdmin can't be detected by packages
    try:
        check_da = subprocess.check_output(['/sbin/vzctl', 'exec', ct, 'ls', '/usr/local/directadmin/directadmin'], stderr=subprocess.DEVNULL).decode('utf-8')
        if check_da:
            log_info(ct + ": " + "Conversion aborted: We have detected DirectAdmin Software inside the CT. It will not function after the upgrade.")
            proceed = False
    except:
        pass

    # Need to enable ssh root login explicitely
    try:
        check_ssh = subprocess.check_output(['/sbin/vzctl', 'exec', ct, 'grep', '^PermitRootLogin', '/etc/ssh/sshd_config']).decode('utf-8')
    except:
        try:
            subprocess.call(['/sbin/vzctl', 'exec', ct, 'sed', '"/#PermitRoot/aPermitRootLogin yes"', '-i', '/etc/ssh/sshd_config']).decode('utf-8')
        except:
            pass

    if not proceed:
        if need_stop:
            try_stop_ct(ct)
        return

    if not check_space(ct):
        if not args.strict:
            log_info(ct + ": Warning! May not be enough free space in the container")
        else:
            log_info(ct + ": Conversion aborted: May not be enough free space in the container")
            if need_stop:
                try_stop_ct(ct)
            return

    if not check_repos(ct, distro_kind):
        if not args.strict:
            log_info(ct + ": Warning! Unsupported repositories detected")
        else:
            log_info(ct + ": Conversion aborted: Unsupported repositories detected")
            if need_stop:
                try_stop_ct(ct)
            return

    if not args.dry_run:
        try:
            snaphost_out = subprocess.check_output(['/bin/prlctl', 'snapshot', ct, '-n', 'pre-vzlinux8'])
        except:
            snaphost_out = subprocess.check_output(['/sbin/vzctl', 'snapshot', ct, '--name', 'pre-vzlinux8'])

        log_info(ct + ": " + snaphost_out.decode('utf-8'))
        old_proc_list = sorted(set(subprocess.check_output(['/bin/vzps', '-Ao', 'fname', '-E', ct]).decode('utf-8').split("\n")))
        old_pkg_list = old_pkg_list_raw.split("\n")
        old_pkg_list = [p.split(" ")[0] for p in old_pkg_list]
        old_ports_list = sorted(set(get_open_ports(ct, distro_kind)))
        # App template list
        if distro_kind == 2:
            old_tmpl_list_raw = sorted(set(subprocess.check_output(['/sbin/vzpkg', 'list', ct]).decode('utf-8').split("\n")))
            old_tmpl_list = [p.split()[1] for p in old_tmpl_list_raw if len(p.split()) > 1]


    vzdeploy_args = ['-n']
    if args.quiet:
        vzdeploy_args.append('-q')
    elif args.verbose:
        vzdeploy_args.append('-v')
    if args.dry_run:
        vzdeploy_args.append('-d')

    vzdeploy_args.append('-c')
    vzdeploy_args.append(ct)

    log_info("Starting conversion: " + ct + " at " + str(datetime.datetime.now().time()))

    # For CentOS 7, we first convert to VzLinux7 and then use 'vzpkg upgrade'
    if distro_kind == 2:
        pr = subprocess.Popen(['/bin/vzdeploy7_ct'] + vzdeploy_args, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1)
        if args.log:
            ct_log = open(args.log + "." + ct, "w")
        else:
            ct_log = None
        while True:
            l = pr.stdout.readline()
            if l == '' and pr.poll() != None:
                break
            if 'Warning! rpmdb:' in l or 'ERROR: ld.so:' in l:
                continue
            if l:
                log_info(l.strip(), ct_log)
        if args.log:
            ct_log.close()

        # In case of dry-run, do not try to launch vzpkg upgrade
        if args.dry_run:
            if need_stop:
                try_stop_ct(ct)

            log_info(ct + ": Conversion successful at " + str(datetime.datetime.now().time()))
            return
        upgrade_c7_ct(ct)
    else:
        pr = subprocess.Popen(['/bin/vzdeploy8_ct'] + vzdeploy_args, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1)
        if args.log:
            ct_log = open(args.log + "." + ct, "w")
        else:
            ct_log = None
        while True:
            l = pr.stdout.readline()
            if l == '' and pr.poll() != None:
                break
            if 'Warning! rpmdb:' in l or 'ERROR: ld.so:' in l:
                continue
            if l:
                log_info(l.strip(), ct_log)
        if args.log:
            ct_log.close()

    if not args.dry_run:
        save_timestamp(ct, distro_kind)

        if distro_kind == 2:
            new_tmpl_list_raw = sorted(set(subprocess.check_output(['/sbin/vzpkg', 'list', ct]).decode('utf-8').split("\n")))
            new_tmpl_list = [p.split()[1] for p in new_tmpl_list_raw if len(p.split()) > 1]
            readd_tmpl = [p for p in old_tmpl_list if p not in new_tmpl_list and p not in REMOVED_TMPL_IGNORE]
            for t in readd_tmpl:
                try:
                    # There still can be templates missing in the new os and not covered in REMOVED_TMPL_IGNORE
                    pr = subprocess.check_output(['/sbin/vzpkg', 'install', ct, t])
                    log_info(pr)
                except:
                    pass

        new_proc_list = sorted(set(subprocess.check_output(['/bin/vzps', '-Ao', 'fname', '-E', ct]).decode('utf-8').split("\n")))
        new_pkg_list = subprocess.check_output(['/sbin/vzpkg', 'list', '-p', ct]).decode('utf-8').split("\n")
        new_pkg_list = [p.split(" ")[0] for p in new_pkg_list]
        new_ports_list = sorted(set(get_open_ports(ct)))
        added_pkgs = [p for p in new_pkg_list if p.replace("vl7", "el7") not in old_pkg_list and p not in ADDED_PKGS_IGNORE]
        removed_pkgs = [p for p in old_pkg_list if p.replace("el7", "vl7") not in new_pkg_list and p not in DROPPED_PKGS_IGNORE]
        added_ps = [p for p in new_proc_list if p not in old_proc_list and p not in CHANGED_PS_IGNORE]
        removed_ps = [p for p in old_proc_list if p not in new_proc_list and p not in CHANGED_PS_IGNORE]
        added_ports = [p for p in new_ports_list if p not in old_ports_list]
        removed_ports = [p for p in old_ports_list if p not in new_ports_list]
        if added_pkgs and distro_kind != 2:
            # After upgrade from c7, we have a ton of added/removed packages, no sense to report them
            msg = ct + ": Warning!\nThe following packages were added compared to the old container state:" + str(added_pkgs)
            log_info(msg)
        if removed_pkgs and distro_kind != 2:
            msg = ct + ": Warning!\nThe following packages were removed compared to the old container state:" + str(removed_pkgs)
            log_info(msg)
        if added_ps:
            msg = ct + ": Warning!\nThe following processes became active compared to the old container state:" + str(added_ps)
            log_info(msg)
        if removed_ps:
            msg = ct + ": Warning!\nThe following processes became inactive compared to the old container state:" + str(removed_ps)
            log_info(msg)
        if added_ports:
            msg = ct + ": Warning!\nThe following ports were open compared to the old container state:" + str(added_ports)
            log_info(msg)
        if removed_ports:
            msg = ct + ": Warning!\nThe following ports were closed compared to the old container state:" + str(removed_ports)
            log_info(msg)

    if need_stop:
        try_stop_ct(ct)

    log_info(ct + ": Conversion successful at " + str(datetime.datetime.now().time()))


'''
Save conversion timestamp inside container.
Useful for CEP to know that the container was converted
'''
def save_timestamp(ctid, distro_kind=1):
    if distro_kind == 1:
        subprocess.call(['/sbin/vzctl', 'exec', ctid, '/bin/touch', '/var/log/vzconvert8.stamp'])
    else:
        subprocess.call(['/sbin/vzctl', 'exec', ctid, '/bin/touch', '/var/log/vzconvert7.stamp'])

'''
Check if version of centos-8.stream template package is suitable for c7 -> vzl8 upgrade
'''
def check_vzl8_tmp_ver():
    try:
        vzl8_tmpl_ver = subprocess.check_output(['/bin/rpm', '-q', '--qf', '%{RELEASE}', 'vzlinux-8-x86_64-ez']).decode('utf-8')
    except:
        return False

    if not vzl8_tmpl_ver:
        return False
    vzl8_tmpl_ver = vzl8_tmpl_ver.replace(".vz7", "").replace(".vz8", "").replace(".vz9", "")
    if int(vzl8_tmpl_ver) < 18:
        return False
    return True

'''
Get a list of containers that can be subjected for upgrade
'''
def get_upgradable():
    c7_available = check_vzl8_tmp_ver()
    if not c7_available:
        print("WARNING: Your version of vzlinux-8-x86_64-ez template is too old, upgrade of CentOS 7 CTs is not available.\n")
    all_ct = subprocess.check_output(['/sbin/vzlist', '-a', '-o', 'ostemplate,ctid,name'])
    for l in all_ct.decode('utf-8').split("\n"):
        try:
            parts = l.split()
            ostemplate = parts[0]
        except:
            continue
        if ostemplate.endswith('centos-8.stream-x86_64') or ostemplate.endswith('centos-8.stream') \
                or ostemplate.endswith('centos-8-x86_64') or ostemplate.endswith('centos-8') \
                or (c7_available and (ostemplate.endswith('centos-7-x86_64') or ostemplate.endswith('centos-7'))) \
                or ostemplate.endswith('almalinux-8-x86_64') or ostemplate.endswith('almalinux-8'):
            print("%s (%s)" % (parts[2], parts[1]))

if __name__ == '__main__':
    global args
    parse_command_line()
    args.func()

