#!/usr/bin/python3 -u
# -*- coding: utf-8 -*-
#
##########################################################################
# Maj-Auto - Manage automatique update of EOLE server
# Copyright © 2015 Pôle de compétences EOLE <eole@ac-dijon.fr>
#
# License CeCILL:
#  * in french: http://www.cecill.info/licences/Licence_CeCILL_V2-fr.html
#  * in english http://www.cecill.info/licences/Licence_CeCILL_V2-en.html
##########################################################################

#add this in post-upgrade
#systemctl mask netplan-wait-online.service networkd-dispatcher.service netplan-apply.service

import sys

import warnings
import apt

import atexit
from collections import OrderedDict

from argparse import ArgumentParser
import locale
from pyeole.i18n import i18n
from pyeole import lock
from pyeole.process import system_code
from pyeole.ansiprint import print_title, print_red, print_orange, print_green
from pyeole.ihm import question_ouinon
from creole.config import configeoldir

from os import system
from os import access
from os import R_OK
from os import mkdir
from os import unlink
from os.path import join, isfile, isdir, basename, dirname

import time
import re
from shutil import copytree, copy, rmtree
from pyeole.pkg import EolePkg, _configure_sources_mirror
from creole.fonctionseole import controle_kernel, zephir
from creole.config import eoledir
from creole.config import templatedir
from creole.config import vareole
from creole.template import CreoleTemplateEngine
from creole.client import CreoleClient

from glob import glob
import hashlib
from urllib.request import urlopen
from urllib.request import ProxyHandler
from urllib.request import build_opener
from html.parser import HTMLParser
from shutil import disk_usage

warnings.filterwarnings("ignore", "apt API not stable yet", FutureWarning)

_ = i18n('update-manager')

if isfile(join(configeoldir, '.release_available')):
    print("")
    print_red(_(u'No newer major version available'))
    print("")
    print(_(u'Please, use Maj-Release script first to upgrade the minor version'))
    print("")
    sys.exit(1)

if not isfile(join(configeoldir, '.upgrade_available')):
    print("")
    print_red(_(u'No newer major version available'))
    print("")
    sys.exit(1)

try:
    sys.path.append('/usr/lib/python3/dist-packages')
    from UpdateManager.Core.MetaRelease import MetaReleaseCore
    from DistUpgrade.DistUpgradeFetcherCore import DistUpgradeFetcherCore
    from DistUpgrade.utils import init_proxy
    sys.path.pop()
except:
    print("")
    print_red(_(u'No newer major version available'))
    print("")
    sys.exit(1)

quit_re = re.compile(r'q|(quit)|(exit)|(abort)', re.I)

tmp_dir = '/tmp/Upgrade-Auto'


MATRIX = OrderedDict()
#hapy : cf. https://dev-eole.ac-dijon.fr/issues/31354
MATRIX['2.8.1'] = ['amon', 'eolebase', 'scribe', 'seshat', 'seth', 'sphynx', 'thot', 'zephir']
MATRIX['2.8.0'] = ['eolebase', 'scribe', 'seshat', 'seth', 'sphynx', 'thot', 'zephir']


def verify_module_not_allowed(release):
    client = CreoleClient()
    module = client.get_creole('eole_module')
    if module not in MATRIX[release]:
        confirmation_msg = _(u"The upgrade of this module is not officially supported, do you to want to continue?")
        if question_ouinon(confirmation_msg) != 'oui':
            sys.exit(1)

def release_lock():
    if lock.is_locked('upgrade-auto', level='system'):
        lock.release('upgrade-auto', level='system')


class EOLEDistUpgradeFetcherCore(DistUpgradeFetcherCore):
    def verifyDistUprader(self, *args, **kwargs):
        copy(join(tmp_dir, 'DistUpgradeViewEOLE.py'), self.tmpdir)
        copy(join(self.tmpdir, 'utils.py'), join(self.tmpdir, 'old_utils.py'))
        copy(join(tmp_dir, 'utils.py'), self.tmpdir)
        return super(EOLEDistUpgradeFetcherCore, self).verifyDistUprader(*args, **kwargs)


def cli_choice(alternatives, prompt, title=None, guess=False):
    """
    Display choices in terminal and return chosen one
    :param alternatives: choices proposed to user
    :type alternatives: list
    :param guess: wether to guess choice
    :type guess: boolean
    """
    def default_input(alt_mapping, prompt, title=None, guess=False):
        choices = "\n".join(["[{0}] {1}".format(alt[0], alt[1])
                            for alt in alt_mapping.items()])
        choices = _(u"Available choices:\n{}\n").format(choices)
        if title is not None:
            print_green(title)
        print(choices)
        if guess is True and len(alt_mapping) < 2:
            print_green(_(u"Automatically selected first choice: {}\n").format(alt_mapping[1]))
            choice = "1"
        else:
            try:
                prompt = prompt + _(u"\n[1]: ")
                choice = input(prompt)
            except(KeyboardInterrupt, EOFError):
                print(_("\nUpgrade aborted by user"))
                sys.exit(0)
            if choice == '':
                choice = "1"
        return choice

    alt_mapping = {num + 1: choice for num, choice in enumerate(alternatives)}
    choice = default_input(alt_mapping, prompt, title=title, guess=guess)
    if choice not in alt_mapping.values():
        try:
            choice = alt_mapping[int(choice)]
        except KeyError:
            print(_("Choice {} not available\n").format(choice))
            choice = cli_choice(alternatives, prompt, guess=guess)
        except ValueError:
            if quit_re.match(choice):
                print(_("Upgrade cancelled by user"))
                sys.exit(0)
            else:
                print(_("Invalid input: {}\n").format(choice))
                choice = cli_choice(alternatives, prompt, guess=guess)
    return choice


def upgrade_container_source(container):
    """Edit source list in container
    :param container: container name
    :type container: str
    """
    print('changement des sources.list pour le conteneur {}'.format(container))
    source_list = '/opt/lxc/{}/rootfs/etc/apt/sources.list'.format(container)
    with open(source_list) as old_source:
        sources = old_source.read()
    with open(source_list, 'w') as new_source:
        new_source.write(sources.replace('precise', 'trusty'))
    cmd = ['apt-get', 'update']
    code = system_code(cmd, container=container)
    return code


# set the rc version here, otherwise ''
RC_VERSION = ''
#RC_VERSION = 'rc1'
ALTERNATIVES = tuple(MATRIX.keys())
RUNPARTS_CMD = u'/bin/run-parts --exit-on-error -v {directory}'
PRE_SOURCE = 'pre_source'
PRE_UPGRADE = 'pre_upgrade'
UPGRADE_DIR = join(eoledir, 'upgrade')
ISO_DIR = join(vareole, 'iso')
ISO_URL_BASE = 'http://eole.ac-dijon.fr/pub/iso'
DEFAULT_RATE = '120k'
z_proc = "UPGRADE"

class ExtractEOLEVersions(HTMLParser):
    """Extrat stable EOLE versions from HTML page

    Gathered versions are stored in ``self.versions``.

    """

    def __init__(self, version):
        HTMLParser.__init__(self)
        self.version = version
        self.versions = []
        self.process_a = False

    def handle_starttag(self, tag, attrs):
        if tag != 'a':
            self.process_a = False
            return
        self.process_a = True


    def handle_data(self, data):
        if not self.process_a:
            return

        # Strip not matching and pre stable versions if it's not a pre stable versions
        if data.lower().startswith(self.version) and ('-' in self.version or '-' not in data):
            self.versions.append(data.rstrip('/'))


def get_most_recent_version(match, url, proxy_handler):
    """Get the most recent matching version

    """
    opener = build_opener(proxy_handler)
    http_request = opener.open(url)
    html_parser = ExtractEOLEVersions(match)
    html_parser.feed(http_request.read())
    html_parser.versions.sort()
    try:
        return html_parser.versions[-1]
    except IndexError:
        print_red(_("Cannot find image for {} in {}").format(match, url))
        exit(1)


def build_iso_name(version):
    """Build ISO name for version

    """
    arch = apt.apt_pkg.config.get('APT::Architecture')
    iso_name = 'eole-{version}-alternate-{arch}.iso'.format(version=version,
                                                            arch=arch)
    return iso_name


def build_release_url(version_url, latest_version):
    """Build the URL of latest release of a version

    """
    release_url = '{base}/{release}'.format(base=version_url,
                                            release=latest_version)
    return release_url


def check_iso(iso_name, path, release_url, proxy, proxy_handler):
    """Verify checksum and signature

    """
    if not access(path, R_OK):
        err_msg = _("Unreadable file: {iso}").format(iso=path)
    sha256_file = join(ISO_DIR, 'SHA256SUMS')
    sha256_url = release_url + '/SHA256SUMS'

    sha256_gpg_file = join(ISO_DIR, 'SHA256SUMS.gpg')
    sha256_gpg_url = sha256_url + '.gpg'

    sha256 = hashlib.sha256()

    iso_ok = False

    print(_(u"Verifying ISO image {iso}").format(iso=path))

    if not isfile(sha256_file):
        print(_(u"Download SHA256SUMS file"))
        download_with_wget(sha256_file, sha256_url, proxy, proxy_handler)

    if not isfile(sha256_gpg_file):
        print(_(u"Download SHA256SUMS.gpg file"))
        download_with_wget(sha256_gpg_file, sha256_gpg_url, proxy, proxy_handler)

    print(_(u"Check SHA256SUMS file signature"))

    gpg_cmd = ['gpgv', '-q', '--keyring',
               '/etc/apt/trusted.gpg.d/eole-archive-keyring.gpg',
               sha256_gpg_file, sha256_file]

    ret = system_code(gpg_cmd)
    if ret == 0:
        print(_(u"Check ISO SHA256..."))
        sha_fh = open(sha256_file, 'r')
        for line in sha_fh:
            sha, filename = line.split()
            if filename != '*{iso_name}'.format(iso_name=iso_name):
                continue

            with open(path, 'rb') as iso_fh:
                while True:
                    block = iso_fh.read(2**10)
                    if not block:
                        break

                    sha256.update(block)

                iso_ok = sha == sha256.hexdigest()
                if iso_ok:
                    print(_(u'OK'))
                    return True

        print(_(u'Error'))


    # Default
    return False


def download_with_wget(out_file, url, proxy, proxy_handler, limit_rate="0"):
    """Use wget to download a file with a progress bar and rate limit

    """
    opener = build_opener(proxy_handler)
    http_request = opener.open(url)
    dl_file_size = int(http_request.info()['content-length'])
    dusage = disk_usage(dirname(out_file))
    if dusage.free <= dl_file_size:
        err_msg = _("The {url} file is too big to be saved in {file}, you can download it to another location and use --iso to set destination filename")
        err_msg = err_msg.format(file=out_file, url=url)
        zephir("ERR", err_msg, z_proc)
        raise SystemError(err_msg)
    wcmd = ['wget', '-c', '--progress', 'dot:giga']
    if limit_rate != '0':
        wcmd.extend(['--limit-rate', limit_rate])

    wcmd.extend(['-O', out_file, url])

    if proxy is not None:
        env = {'http_proxy': proxy,
               'https_proxy': proxy}
    else:
        env = {}

    if not system_code(wcmd, env=env) == 0:
        err_msg = _("Error downloading {file} with wget from {url}")
        err_msg = err_msg.format(file=out_file, url=url)
        zephir("ERR", err_msg, z_proc)
        raise SystemError(err_msg)

def download_iso_with_zsync(iso_file, iso_url):
    """Use zsync to download the ISO

    """
    zcmd = ['zsync', '-o', iso_file, iso_url + '.zsync']
    if not system_code(zcmd) == 0:
        err_msg = _("Error downloading the image with zsync")
        zephir("ERR", err_msg, z_proc)
        raise SystemError(err_msg)

def clean_iso_dir(iso_file=None):
    """Clean ISO directory

    If :data:`iso_file` is not `None`, it's keept

    """
    # Remove any file other than targeted release ISO
    # This permit to resume download
    for filename in glob('{iso_dir}/*'.format(iso_dir=ISO_DIR)):
        if filename == iso_file:
            continue

        unlink(filename)

def get_cdrom_device():
    """
    Manage CDROM device
    """
    device = None
    client = CreoleClient()
    mount_point = '/media/cdrom'
    mounted = False
    if not isdir(mount_point):
        mkdir(mount_point)

    for cdrom in client.get_creole('cdrom_devices'):
        cmd = ['/bin/mount', cdrom, mount_point, '-o', 'ro']
        if system_code(cmd) != 0:
            continue

        mounted = True

        if isdir('{0}/dists'.format(mount_point)):
            device = cdrom

        cmd = ['/bin/umount', mount_point]
        if system_code(cmd) != 0:
            err_msg = _("Unable to umount {cdrom}").format(cdrom=cdrom)
            zephir("ERR", err_msg, z_proc)
            raise SystemError(err_msg)

        if device is not None:
            break

    return device


def download_iso(args, limit_rate):
    """Download ISO image

    Download an ISO image from internet or copy one from :data:`path`.

    :parameter path: path of an existing ISO image
    :type path: `str`

    """
    proxy = None
    client = CreoleClient()

    if client.get_creole('activer_proxy_client') == u'oui':
        address = client.get_creole('proxy_client_adresse')
        port = client.get_creole('proxy_client_port')
        proxy = 'http://{address}:{port}'.format(address=address, port=port)

    release = args.release
    if RC_VERSION:
        release += '-' + RC_VERSION
    version_url = ISO_URL_BASE + '/EOLE-' + '.'.join(release.split('.')[0:2])
    if proxy is not None:
        proxy_handler = ProxyHandler({'http': proxy,
                                      'https': proxy})
    else:
        proxy_handler = ProxyHandler({})
    if args.iso is not None:
        release_url, iso_name = args.iso.rsplit('/', 1)
    else:
        version_url = ISO_URL_BASE + '/EOLE-' + '.'.join(args.release.split('.')[0:2])
        latest_version = get_most_recent_version(args.release, version_url, proxy_handler)

        release_url = build_release_url(version_url, latest_version)
        iso_name = build_iso_name(latest_version)
    iso_file = join(ISO_DIR, iso_name)
    iso_url = '{url}/{iso}'.format(url=release_url, iso=iso_name)

    print_title(_(u"Downloading ISO image for {release}").format(release=release))

    if not isdir(ISO_DIR):
        mkdir(ISO_DIR)

    if not args.cdrom and isfile(iso_file) and check_iso(iso_name, iso_file, release_url, proxy, proxy_handler):
        # ISO is downloaded and verified
        return iso_file

    # Remove SHA265SUMS* files
    # Keep ISO to resume download if possible
    clean_iso_dir(iso_file)
    err_msg = None

    if args.cdrom or is_local_iso(args):
        if args.cdrom:
            path = get_cdrom_device()
            if path is None:
                err_msg = _("No CDROM found")
        else:
            path = args.iso
            if not isfile(path):
                err_msg = _("No such file: {iso}").format(iso=path)
            elif not access(path, R_OK):
                err_msg = _("Unreadable file: {iso}").format(iso=path)
        if err_msg is None:
            # Should we also check source image before copying?
            # elif not check_iso(iso_name, path, release_url):
            #     raise SystemError("ISO image is not valid for {iso}".format(iso=path))

            print(_("Copying {source} to {iso}").format(source=path, iso=iso_file))

            copy(path, iso_file)

            if not check_iso(iso_name, iso_file, release_url, proxy, proxy_handler):
                # Copy was OK but check fails
                # Remove copied ISO as it may be corrupted and prevent
                # futur download
                clean_iso_dir()
                msg = _("Error checking ISO after copy, remove {iso}")
                err_msg = msg.format(iso=iso_file)

    else:
        # Try resuming download
        download_with_wget(iso_file, iso_url, proxy, proxy_handler, limit_rate)
        # download_iso_with_zsync(iso_file, iso_url)
        msg = _(u"Error checking ISO image after download, remove {iso}")

        if not check_iso(iso_name, iso_file, release_url, proxy, proxy_handler):
            # Download was OK but check fails
            # Remove downloaded ISO as it may be fully downloaded but
            # corrupted and prevent futur download
            clean_iso_dir()
            msg = _(u"Error checking ISO image after download, remove {iso}")
            err_msg = msg.format(iso=iso_file)

    if err_msg:
        zephir('ERR', err_msg, z_proc)
        raise SystemError(err_msg)
    if args.cdrom:
        return ''
    return iso_file


def is_local_iso(args):
    return args.iso and args.iso.split('://')[0] not in ['http', 'https']


def main():
    args_parser = ArgumentParser(description=_("EOLE distribution upgrade tool."))

    args_parser.add_argument('--release',
                             help=_(u"Target release number"))

    args_parser.add_argument('--download', action='store_true',
                             help=_(u"Only download the ISO image"))

    args_parser.add_argument('--iso', metavar=u'PATH',
                             help=_(u"Path to an ISO image"))

    args_parser.add_argument('--cdrom', action='store_true',
                             help=_(u"Use CDROM device instead of downloading ISO image"))

    args_parser.add_argument('--limit-rate', metavar=u'BANDWIDTH',
                             help=_(u"Pass limit rate to wget. “0” to disable."))

    args_parser.add_argument('-f', '--force', action='store_true',
                             help=_(u"Do not ask confirmation"))

    args = args_parser.parse_args()


    try:
        locale.setlocale(locale.LC_ALL, "")
    except:
        pass

    print_red(_(u"This script will upgrade this server to a new release"))
    print_red(_(u"Modifications will be irreversible."))
    init_proxy()
    zephir("INIT", _(u"Starting Upgrade-Auto ({})").format(" ".join(sys.argv[1:])), z_proc)

    if is_local_iso(args) and not isfile(args.iso):
        print_red(_(u'The "{}" file does not exist').format(args.iso))
        sys.exit(1)
    # Ask for release if none provided on command line
    if args.iso:
        filename = basename(args.iso)
        match = re.search(r'eole-(?P<version>[0-9\.]+)-alternate-amd64.iso', filename)
        if match and match.group('version') in ALTERNATIVES:
            args.release = match.group('version')
        else:
            msg = _(u'Unable to find version with the iso file name')
            print_orange(msg)
    if args.release is not None:
        if args.release not in ALTERNATIVES:
            msg = _(u"Invalid release {version} use: {values}")
            err_msg = msg.format(version=args.release,
                                 values=', '.join(ALTERNATIVES))
            zephir("ERR", err_msg, z_proc)
            print_red(err_msg)
            sys.exit(1)
    else:
        title = _(u"Choose which version you want to upgrade to\n")
        prompt = _(u"Which version do you want to upgrade to (or 'q' to quit)?")
        args.release = cli_choice(ALTERNATIVES, prompt, title=title, guess=False)


    if is_local_iso(args):
        limit_rate = None
    elif args.limit_rate is None:
        if args.force:
            limit_rate = DEFAULT_RATE
        else:
            print(_(u"Bandwidth limit to use for download (“0” disables the limitation)?"))
            limit_rate = raw_input("[{}] : ".format(DEFAULT_RATE))
            if limit_rate == '':
                limit_rate = DEFAULT_RATE
    else:
        limit_rate = args.limit_rate

    if not args.force:
        confirmation_msg = _(u"Do you really want to upgrade to version {}?")
        if question_ouinon(confirmation_msg.format(args.release)) != 'oui':
            end_msg = _(u'Upgrade cancelled by user')
            zephir("FIN", end_msg, z_proc)
            print_red(end_msg)
            sys.exit(0)

    verify_module_not_allowed(args.release)

    lock.acquire('upgrade-auto', valid=False, level='system')
    atexit.register(release_lock)

    if not args.download:
        print_title(_("Check update status"))
        PKGMGR = EolePkg('apt', ignore=False)
        PKGMGR.set_option('APT::Get::Simulate', 'true')
        if RC_VERSION:
            eole_level = 'proposed'
            envole_level = 'proposed'
        else:
            eole_level = 'stable'
            envole_level = 'stable'
        _configure_sources_mirror(PKGMGR.pkgmgr,
                                  eole_level=eole_level,
                                  envole_level=envole_level)
        PKGMGR.update(silent=True)
        upgrades = PKGMGR.get_upgradable_list()

        for container, packages in upgrades.items():
            if packages:
                err_msg = _(u"Some packages are not up-to-date!")
                zephir("ERR", err_msg, z_proc)
                print_red(err_msg)
                print_red(_(u"Update this server (Maj-Auto) before another attempt to upgrade"))
                sys.exit(1)

        print_green(_(u'Server is up-to-date'))

        if controle_kernel():
            err_msg = _(u"In order to upgrade, most recent kernel endorsed for this release must be used")
            zephir("ERR", err_msg, z_proc)
            print_red(err_msg)
            sys.exit(1)

        print_green(_(u'This server uses most recent kernel'))

    if is_local_iso(args):
        iso_file = args.iso
    else:
        iso_file = download_iso(args, limit_rate)

    if args.download:
        end_msg = _(u"Download only detected, stop")
        print_green(end_msg)
        zephir("FIN", end_msg, z_proc)
        sys.exit(0)

    print_title(_("Copying upgrade scripts"))
    if isdir(tmp_dir):
        err_msg = _(u"Directory {0} already exists").format(tmp_dir)
        zephir("ERR", err_msg, z_proc)
        print_red(err_msg)
        if args.force:
            sys.exit(1)
        confirmation_msg = _(u"Do you want to delete the directory and continue?")
        if question_ouinon(confirmation_msg) != 'oui':
            sys.exit(1)
        rmtree(tmp_dir)

    copytree(UPGRADE_DIR, tmp_dir)
    with open(join(tmp_dir, 'iso.file'), 'w') as fh:
        fh.write(iso_file)

    print_title(_("Module specific commands"))
    pre_source = join(UPGRADE_DIR, PRE_SOURCE)
    code = system(RUNPARTS_CMD.format(directory=pre_source))
    if code != 0:
        err_msg = _(u'Error {0}').format(pre_source)
        print_red(err_msg)
        zephir("ERR", err_msg, z_proc)
        sys.exit(1)

    print_title(_("Configuring upgrade"))
    engine = CreoleTemplateEngine()
    rootctx = {u'name': u'root', u'path': u''}

    # order is arbitrary and eole.cfg must be the last loaded
    # so remove extra .cfg
    for filename in glob('/etc/update-manager/release-upgrades.d/*.cfg'):
        unlink(filename)
    for filename in ['/etc/update-manager/release-upgrades.d/eole.cfg', '/etc/update-manager/mirrors_eole.cfg', '/etc/update-manager/meta-release']:
        file_info = {'name': filename,
                     'source': join(templatedir, basename(filename)),
                     'full_name': filename,
                     'activate' : True,
                     'del_comment': u'',
                     'mkdir' : False,
                     'rm' : False}
        engine.prepare_template(file_info)
        engine.process(file_info, rootctx)

    release = args.release
    #overwrite envole_version
    envole_release = '8'
    level = u'stable'
    if RC_VERSION:
        level = u'proposed'
    _configure_sources_mirror(PKGMGR.pkgmgr, release=release, envole_release=envole_release, eole_level=level)

    print_title(_("Module specific commands"))
    pre_upgrade = join(tmp_dir, PRE_UPGRADE)
    code = system(RUNPARTS_CMD.format(directory=pre_upgrade))
    if code != 0:
        err_msg = _(u'Error {0}').format(pre_upgrade)
        print_red(err_msg)
        zephir("ERR", err_msg, z_proc)
        sys.exit(1)

    for container in (cont for cont in upgrades if cont != 'root'):
        upgrade_container_source(container)
        PKGMGR.dist_upgrade(container=container, silent=False)
    if system_code([join(tmp_dir, 'Upgrade-Auto3')]):
        print_red(_('The upgrade finished with an error'))
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except (KeyboardInterrupt, EOFError):
        print_red(_("\nUpgrade aborted by user"))
    except SystemError as err:
        print_red(str(err))
        sys.exit(1)
