rougail/upgrade/Upgrade-Auto

#!/usr/bin/python -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
##########################################################################

import sys

import warnings
import apt

import atexit

from argparse import ArgumentParser
import locale
from pyeole.i18n import i18n
from pyeole import lock
from pyeole.process import system_code
from pyeole.ihm import print_title, print_red, 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
from os.path import isfile
from os.path import isdir

import time
import re
from shutil import copytree
from shutil import copy
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 urllib2 import urlopen
from urllib2 import ProxyHandler
from urllib2 import build_opener
from HTMLParser import HTMLParser

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

_ = i18n('update-manager')

if not isfile(join(configeoldir, '.upgrade_available')):
    print ""
    print_red(_(u'Upgrade to a newer major version is not available'))
    print _(u'Use Maj-Release script to upgrade to next minoir version')
    print ""
    sys.exit(1)

from UpdateManager.Core.MetaRelease import MetaReleaseCore
from UpdateManager.Core.DistUpgradeFetcherCore import DistUpgradeFetcherCore
from DistUpgrade.utils import init_proxy

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

tmp_dir = '/tmp/Upgrade-Auto'

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)
        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 = raw_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 ', 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


ALTERNATIVES = ("2.5.1", )
RUNPARTS_CMD = u'/bin/run-parts --exit-on-error -v {directory}'
PRE_DOWNLOAD = 'pre_download'
UPGRADE_DIR = join(eoledir, 'upgrade')
ISO_DIR = join(vareole, 'iso')
ISO_URL_BASE = 'http://eole.ac-dijon.fr/pub/iso'
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 data.lower().startswith(self.version) and '-' not in data:
            self.versions.append(data.rstrip('/'))


def get_most_recent_version(match, url, proxy=None):
    """Get the most recent matching version

    """
    if proxy is not None:
        proxy_handler = ProxyHandler({'http': proxy,
                                      'https': proxy})
    else:
        proxy_handler = ProxyHandler({})

    opener = build_opener(proxy_handler)
    http_request = opener.open(url)
    html_parser = ExtractEOLEVersions(match)
    html_parser.feed(http_request.read())
    html_parser.versions.sort()
    return html_parser.versions[-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, proxy=None):
    """Build the URL of latest release of a version

    """
    version_url = ISO_URL_BASE + '/EOLE-' + '.'.join(version.split('.')[0:2])
    latest_version = get_most_recent_version(version, version_url, proxy)
    release_url = '{base}/{release}'.format(base=version_url,
                                            release=latest_version)
    return release_url


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

    """
    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)

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

    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=None, limit_rate="0"):
    """Use wget to download a file with a progress bar and rate limit

    """
    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():
    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):
    """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_url = build_release_url(args.release, proxy)
    iso_name = build_iso_name(args.release)
    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=args.release))

    if not isdir(ISO_DIR):
        mkdir(ISO_DIR)

    if isfile(iso_file) and check_iso(iso_name, iso_file, release_url, proxy):
        # ISO is downloaded and verified
        return True

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

    if args.cdrom or args.iso is not None:

        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):
                # 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, args.limit_rate)
        # download_iso_with_zsync(iso_file, iso_url)

        if not check_iso(iso_name, iso_file, release_url, proxy):
            # 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)

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', default='120k',
                             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)

    # Ask for release if none provided on command line
    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 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 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 end_msg
            sys.exit(0)

    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')
        _configure_sources_mirror(PKGMGR.pkgmgr)
        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'))

    download_iso(args)

    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)
        sys.exit(1)

    copytree(UPGRADE_DIR, tmp_dir)
    pre_download = join(tmp_dir, PRE_DOWNLOAD)

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

    file_info = {'name': '/etc/update-manager/release-upgrades.d/eole.cfg',
                 'source': join(templatedir, 'eole.cfg'),
                 'full_name': '/etc/update-manager/release-upgrades.d/eole.cfg',
                 'activate' : True,
                 'del_comment': u'',
                 'mkdir' : False,
                 'rm' : False}
    engine.prepare_template(file_info)
    engine.process(file_info, rootctx)

    _configure_sources_mirror(PKGMGR.pkgmgr, eole_release=args.release)

    print_title(_("Module specific commands"))
    code = system(RUNPARTS_CMD.format(directory=pre_download))
    if code != 0:
        err_msg = _(u'Error {0}').format(pre_download)
        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)

    m = MetaReleaseCore(useDevelopmentRelease=False,
                        useProposed=False)
    # this will timeout eventually
    while m.downloading:
        time.sleep(0.5)
    progress = apt.progress.text.AcquireProgress()
    fetcher = EOLEDistUpgradeFetcherCore(new_dist=m.new_dist,
                                     progress=progress)
    fetcher.run_options += ["--mode=server",
                            "--frontend=DistUpgradeViewEOLE",
                            ]
    print_title(_("Upgrading server"))
    fetcher.run()
    # all lines below will not be executed

if __name__ == "__main__":
    try:
        main()
    except (KeyboardInterrupt, EOFError):
        print_red(_("\nUpgrade aborted by user"))
        exit(0)