#!/usr/bin/python
#
# A tool to manage Vz Auto Update service
#

import configparser
import subprocess
import logging
import datetime
import time
import sys
import os
import argparse
import requests
import xml.etree.ElementTree as ET
import psutil
from shutil import copyfile
from email.mime.text import MIMEText
import smtplib

CONFIG_NAME = "/etc/vz/vzautoupdate.conf"
LOG_FILE = "/var/log/vz_auto_update.log"
PATCHBASENAME = "readykernel-patch"
USER_VISIBLE_POLICIES = ['fast', 'stable', 'slow', 'auto']

# Thresholds for system resource usage - if we exceed them, then we won't run update
CPU_THRESHOLD = 95
RAM_THRESHOLD = 95

def parse_command_line():
    global command_line

    parser = argparse.ArgumentParser(description='Vz Auto Update management utility')
    subparsers = parser.add_subparsers(title='action')

    subparser = subparsers.add_parser('get-policy', help='Show the current Vz Auto Update policy')
    subparser.set_defaults(func=print_policy)

    subparser = subparsers.add_parser('list-policies', help='Show available Vz Auto Update policies')
    subparser.set_defaults(func=list_policies)

    subparser = subparsers.add_parser('set-policy',
                                       help='Set auto update policy for the node. Available values: fast, stable and auto.')
    subparser.add_argument('policy', action='store', choices=USER_VISIBLE_POLICIES, help='Policy to use')
    subparser.set_defaults(func=set_policy)

    subparser = subparsers.add_parser('get-schedule', help='Show the current Vz Auto Update schedule')
    subparser.set_defaults(func=get_schedule)

    subparser = subparsers.add_parser('set-schedule', help='Set Vz Auto Update schedule')
    subparser.add_argument('schedule', action='store', help='Schedule in a format accessible by systemd timer Calendar, see Calendar Events in man systemd.time')
    subparser.set_defaults(func=set_schedule)

    subparser = subparsers.add_parser('get-builds', help='Get build numbers in different repos')
    subparser.set_defaults(func=get_builds)

    subparser = subparsers.add_parser('available-updates',
                                       help='List updates available for the specified policy.')
    subparser.add_argument('policy', action='store', choices=['fast', 'slow', 'stable'], help='Policy to check')
    subparser.set_defaults(func=available_updates)

    subparser = subparsers.add_parser('update', help='Launch update')
    subparser.set_defaults(func=update)

    if len(sys.argv)==1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    # We are requested to accept commands starting with '--' as ordinary commands,
    # but '--help' is an exception
    if len(sys.argv)==2 and sys.argv[1] == '--help':
        parser.print_help(sys.stderr)
        sys.exit(0)

    command_line = parser.parse_args([a.replace('--','') for a in sys.argv[1:]])

# Read config from file
def read_config():
    config = configparser.ConfigParser()
    config.read(CONFIG_NAME)
    return config


# Get policy from config
def get_policy(config):
    if 'Default' in config and 'Policy' in config['Default']:
        return config['Default']['Policy']
    return 'auto'


# List available policies
def list_policies():
    print("Available update policies:")
    for p in USER_VISIBLE_POLICIES:
        print("  * " + p)


# List updates available for the ring specified
def available_updates():
    global command_line

    if not command_line.policy:
        policy = get_policy(read_config())
    else:
        policy = command_line.policy

    hwid = get_hwid()
    (REPO_OS, REPO_UP, REPO_RK) = get_repo_urls(policy, hwid)

    p = subprocess.Popen(["yum", "-q", "--disablerepo=*", \
                                         "--enablerepo=" + REPO_OS, \
                                         "--enablerepo=" + REPO_UP, \
                                         "--enablerepo=" + REPO_RK, \
                                         "check-update"], \
                             env=os.environ, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    (out, err) = p.communicate()
    if p.returncode == 100:
        print("Updates available for %s policy:" % policy)
        print(out)
    elif p.returncode == 0:
        print("No updates found for the %s policy" % policy)
    else:
        print("Failed to check updates for the %s policy" % policy)
        print(err)


# Get update launch schedule
def get_schedule():
    CONF_FILE = "/etc/systemd/system/vzautoupdate.timer.d/override.conf"
    if not os.path.isfile(CONF_FILE):
        print("WARNING: Override file /etc/systemd/system/vzautoupdate.timer.d/override.conf is not found, schedule from /lib/systemd/system/vzautoupdate.timer is used")
        CONF_FILE = "/lib/systemd/system/vzautoupdate.timer"
    if not os.path.isfile(CONF_FILE):
        print("/lib/systemd/system/vzautoupdate.timer is not found. No timer - no schedule. Please verify your installation")
        return
    with open(CONF_FILE) as f:
        for l in f.readlines():
            if l.startswith("OnCalendar=") and l.strip() != "OnCalendar=":
                print("vzautoupdate time schedule: " + l.replace("OnCalendar=", ""))
                break


# Set update launch schedule
# We will verify the schedule provided by user using systemd-analyze
# and revert the old value if systemd doesn't like the new one
def set_schedule():
    global command_line

    CONF_FILE = "/etc/systemd/system/vzautoupdate.timer.d/override.conf"
    if not os.path.isfile(CONF_FILE):
        print("WARNING: Override file /etc/systemd/system/vzautoupdate.timer.d/override.conf doen't exist, it will be created")
        try:
            os.makedirs(getSysroot() + "/etc/systemd/system/vzautoupdate.timer.d")
        except:
            pass

        with open(CONF_FILE, "w") as f:
           # Need to empty OnCalendar first to really override values from the unit
           f.write("[Timer]\nOnCalendar=\nOnCalendar=*-*-* 00:00:00")

    schedule = command_line.schedule
    copyfile(CONF_FILE, CONF_FILE + ".orig")
    new_conf = open(CONF_FILE + ".new", "w")
    with open(CONF_FILE) as f:
        for l in f.readlines():
            if l.startswith("OnCalendar=") and l.strip() != "OnCalendar=":
                new_conf.write("OnCalendar=" + schedule + "\n")
            else:
                new_conf.write(l)
    new_conf.close()
    copyfile(CONF_FILE + ".new", CONF_FILE)

    try:
        subprocess.check_call(["systemd-analyze", "verify", "vzautoupdate.timer"])
        subprocess.check_call(["systemctl", "daemon-reload"])
    except:
        print("Timer verification failed, reverting to old config")
        copyfile(CONF_FILE + ".orig", CONF_FILE)
    os.remove(CONF_FILE + ".new")
    os.remove(CONF_FILE + ".orig")


# Print current policy
def print_policy():
    config = read_config()
    p = get_policy(config)
    if p == 'auto':
        hwid = get_hwid()
        if not hwid:
            p = 'auto (not setup yet)'
        else:
            try:
                target_repo = subprocess.check_output(["curl", "-L", "-o", "/dev/null", "-s", "-w", "%{url_effective}", "https://rk.virtuozzo.com/" + hwid + "/vz-update/x86_64"])
                if 'stable' in target_repo:
                    p = 'auto (stable)'
                elif 'slow' in target_repo:
                    p = 'auto (slow)'
                elif hwid in target_repo:
                    p = 'auto (not setup on server or key is invalid)'
                else:
                    p = 'auto (fast)'
            except:
                p = 'auto (not setup yet)'
    print(p)

def set_policy():
    global command_line

    config = read_config()
    if 'Default' not in config:
        config['Default'] = {}
    policy = command_line.policy
    config['Default']['Policy'] = policy
    with open(CONFIG_NAME, 'w') as configfile:
       config.write(configfile)

# Get node HWID from vzlicview
def get_hwid():
    try:
        hwid = subprocess.check_output(['vzlicview'])
    except:
        return None
    for l in hwid.split('\n'):
        if 'hwid=' in l:
            hwid_value = l.replace('hwid="', '').replace('"','').strip()
            return hwid_value
    return None

# Set 'hwid' yum variable
def save_hwid_to_yum(hwid):
    with open("/etc/yum/vars/hwid", "w") as f:
        f.write(hwid)

# Get repo IDs corresponding to a given policy 
def get_repo_urls(policy, hwid):
    if policy == 'fast':
        REPO_OS = 'vz-auto-update-fast'
        REPO_UP = 'vz-auto-update-fast-updates'
        REPO_RK = 'vz-auto-update-rk-fast'
    elif policy == 'stable':
        REPO_OS = 'vz-auto-update-stable'
        REPO_UP = 'vz-auto-update-stable-updates'
        REPO_RK = 'vz-auto-update-rk-stable'
    elif policy == 'slow':
        REPO_OS = 'vz-auto-update-slow'
        REPO_UP = 'vz-auto-update-slow-updates'
        REPO_RK = 'vz-auto-update-rk-stable'
    else:
        # Inform the server that we are going to use auto-controlled repo
        # TODO: Currently nothing to check here, but probably we can teach server
        # to provide some valuable response - e.g., if HWID is not found so no sense to continue
        register = subprocess.check_output(["curl", "-s", "https://rk.virtuozzo.com/node.register?HWID=" + hwid])
        # Just in case, give server some time to prepare our URL
        time.sleep(1)
        REPO_OS = 'vz-auto-update-os'
        REPO_UP = 'vz-auto-update-updates'
        REPO_RK = 'vz-auto-update-rk'

    return (REPO_OS, REPO_UP, REPO_RK)


# Send mail to local root about a need to reboot after update
def send_reboot_mail():
    sender = "vzautoupdate@localhost"
    recpts = ["root"]

    msg = MIMEText("After the latest automated update, it is suggested to reboot the node. Please check 'needs-restarting' command for details")
    msg['Subject'] = 'VZ Auto Update: Reboot required'
    msg['From'] = sender
    msg['To'] = "root"

    try:
        s = smtplib.SMTP('localhost')
        s.sendmail(sender, recpts, msg.as_string())
    except:
        logging.info("Failed to send an email to the local admin that a reboot is needed. Sendmail service inactive.")


# Install updates from Vz auto update repos
def update_vz(REPO_OS, REPO_UP, REPO_RK):
    up_out = subprocess.check_output(["yum", "-y", "--disablerepo=*", "--enablerepo=" + REPO_OS, \
                                       "--enablerepo=" + REPO_RK, \
                                       "--enablerepo=" + REPO_UP, "--enablerepo=virtuozzolinux-base", \
                                       "--enablerepo=virtuozzolinux-updates", "update"])
    logging.info(up_out)

# Install RK patches for existing kernels
def update_rk(REPO_RK):
    # The behavior is taken from readykernel tool
    # Easier to reproduce it here then to teach that tool
    # to look into autoupdate repo only
    for KERNELVER in next(os.walk('/lib/modules'))[1]:
        if not os.path.isdir("/lib/modules/" + KERNELVER + "/kernel"):
            continue

        logging.info("Checking for patch updates for kernel %s..." % KERNELVER)
        yum_check = subprocess.check_output(["yum", "--disablerepo=*", "--enablerepo=" + REPO_RK, "provides", "-q", PATCHBASENAME + "-" + KERNELVER])
        if not yum_check:
            logging.info("No ReadyKernel patches are available for the kernel " + KERNELVER)
            continue

        up_out = subprocess.check_output(["yum", "-y", "--disablerepo=*", "--enablerepo=" + REPO_RK, "install", PATCHBASENAME + "-" + KERNELVER])
        logging.info(up_out)


# Check if auto update can be launched
def check_precond():
    if os.path.isdir('/etc/vstorage/clusters') and len(os.listdir('/etc/vstorage/clusters')) > 0:
        logging.info('Storage cluster detected, will not run auto update')
        return False

    if psutil.cpu_percent() >= CPU_THRESHOLD:
        logging.info('CPU Load is too high, will not run auto update')
        return False

    if psutil.virtual_memory().percent >= RAM_THRESHOLD:
        logging.info('RAM Usage is too high, will not run auto update')
        return False


    return True

# Launch update using auto-update repos only
def update():
    logging.basicConfig(filename=LOG_FILE, level=logging.INFO)
    logging.info(str(datetime.datetime.today()) + ": Starting update")

    hwid = get_hwid()
    if not hwid:
        logging.info("Can't get license HWID, rejecting auto update")
        sys.exit(0)

    if not check_precond():
        sys.exit(0)

    save_hwid_to_yum(hwid)
    config = read_config()
    policy = get_policy(config)
    (REPO_OS, REPO_UP, REPO_RK) = get_repo_urls(policy, hwid)

    logging.info("Launching RK auto-update for HWID=%s with policy=%s" % (hwid, policy))
    update_rk(REPO_RK)

    logging.info("Launching VZ auto-update for HWID=%s with policy=%s" % (hwid, policy))
    update_vz(REPO_OS, REPO_UP, REPO_RK)

    try:
        c = subprocess.check_call(["needs-restarting", "-r"])
    except subprocess.CalledProcessError as e:
        if e.returncode == 1:
            logging.info("Reboot required after the update")
            send_reboot_mail()


# Get build numbers in different repos
def get_builds():
    URLS = {'fast':   {'release': 'http://repo.virtuozzo.com/vz/releases/7.0/', 'updates': 'http://repo.virtuozzo.com/vz/updates/7.0/'},
            'slow':   {'release': 'http://repo.virtuozzo.com/vz/releases/7.0-slow/', 'updates': 'http://repo.virtuozzo.com/vz/updates/7.0-slow/'},
            'stable': {'release': 'http://repo.virtuozzo.com/vz/releases/7.0-stable/', 'updates': 'http://repo.virtuozzo.com/vz/updates/7.0-stable/'}
            }

    for policy in URLS:
        try:
            r = requests.get(URLS[policy]['release'] + '/x86_64/os/repodata/repomd.xml')
            root =  ET.fromstring(r.content)
            tags = root.find('{http://linux.duke.edu/metadata/repo}tags')
            distro = tags.find('{http://linux.duke.edu/metadata/repo}distro')
            release_build_num = distro.text
        except Exception as e:
            print("Failed to get build info for %s" % policy)
            release_build_num = "Unknown"

        try:
            r = requests.get(URLS[policy]['updates'] + '/x86_64/os/repodata/repomd.xml')
            root =  ET.fromstring(r.content)
            tags = root.find('{http://linux.duke.edu/metadata/repo}tags')
            distro = tags.find('{http://linux.duke.edu/metadata/repo}distro')
            updates_build_num = distro.text
            print("Got " + updates_build_num)
        except Exception as e:
            print(str(e))
            updates_build_num = release_build_num

        if updates_build_num == release_build_num:
            print("Policy '%s': Build %s" % (policy, release_build_num))
        else:
            # Assume that updates build can't be less in number than a release one
            print("Policy '%s': Build %s" % (policy, updates_build_num))

if __name__ == '__main__':
    parse_command_line()
    command_line.func()
