#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
vyos_upgrade.py

Compatible with:
  - Python 2.6.6
  - Python 3.x

Usage:
  ./vyos_upgrade.py [-u URL] [-s] [-n N] [-y] [-h]
"""

from __future__ import print_function

import hashlib
import hashlib
import optparse
import os
import re
import shutil
import subprocess
import sys

try:
    from urllib.parse import urljoin, urlparse
    from urllib.request import urlopen
except ImportError:
    from urlparse import urljoin, urlparse
    from urllib2 import urlopen

try:
    from html.parser import HTMLParser
except ImportError:
    from HTMLParser import HTMLParser


DEFAULT_PAGE_URL = "https://vyos.net/get/nightly-builds/"
MINISIGN_PUBLIC_KEY = (
    "RWSIhkR/dkM2DSaBRniv/bbbAf8hmDqdbOEmgXkf1RxRoxzodgKcDyGq"
)


def log(message):
    print("[*] " + message)


def error(message):
    print("[!] " + message, file=sys.stderr)


def text_value(value):
    """
    Convert bytes/unicode to text consistently on Python 2 and 3.
    """
    if isinstance(value, bytes):
        try:
            return value.decode("utf-8")
        except UnicodeDecodeError:
            return value.decode("latin-1")
    return value


class BuildParser(HTMLParser):
    """
    Extracts the first two links from each <li> inside:

        <div id="rolling-current">
    """

    def __init__(self):
        HTMLParser.__init__(self)
        self.in_container = False
        self.container_depth = 0
        self.current_li = None
        self.li_depth = 0
        self.current_link = None
        self.builds = []

    def handle_starttag(self, tag, attrs):
        attrs = dict(attrs)

        if tag == "div" and attrs.get("id") == "rolling-current":
            self.in_container = True
            self.container_depth = 1
            return

        if not self.in_container:
            return

        if tag == "div":
            self.container_depth += 1

        if tag == "li":
            self.current_li = {
                "links": []
            }
            self.li_depth = 1
            return

        if self.current_li is not None and tag == "a":
            href = attrs.get("href")
            if href:
                self.current_link = {
                    "href": href,
                    "text": []
                }

    def handle_endtag(self, tag):
        if not self.in_container:
            return

        if tag == "a" and self.current_link is not None:
            text = "".join(self.current_link["text"]).strip()
            self.current_li["links"].append(
                (self.current_link["href"], text)
            )
            self.current_link = None
            return

        if tag == "li" and self.current_li is not None:
            links = self.current_li["links"]

            if len(links) >= 2:
                iso_url, iso_text = links[0]
                sig_url, sig_text = links[1]

                iso_name = iso_text.strip()
                if not iso_name:
                    iso_name = os.path.basename(
                        urlparse(iso_url).path
                    )

                iso_name = iso_name.split("/")[-1]

                self.builds.append({
                    "name": iso_name,
                    "iso_url": iso_url,
                    "sig_url": sig_url
                })

            self.current_li = None
            self.current_link = None
            return

        if tag == "div":
            self.container_depth -= 1
            if self.container_depth <= 0:
                self.in_container = False

    def handle_data(self, data):
        if self.current_link is not None:
            self.current_link["text"].append(data)


def read_url(url):
    response = urlopen(url)
    try:
        data = response.read()
    finally:
        response.close()

    return text_value(data)


def download_file(url, filename):
    response = urlopen(url)
    try:
        output = open(filename, "wb")
        try:
            shutil.copyfileobj(response, output)
        finally:
            output.close()
    finally:
        response.close()


def sha256_file(filename):
    digest = hashlib.sha256()
    handle = open(filename, "rb")

    try:
        while True:
            block = handle.read(1024 * 1024)
            if not block:
                break
            digest.update(block)
    finally:
        handle.close()

    return digest.hexdigest()


def verify_sha256(iso_file, checksum_file):
    """
    Supports common sha256sum formats, for example:

      abcdef...  filename.iso
      abcdef... *filename.iso
    """

    handle = open(checksum_file, "rb")
    try:
        content = handle.read()
    finally:
        handle.close()

    content = text_value(content)

    expected = None

    for line in content.splitlines():
        line = line.strip()

        if not line or line.startswith("#"):
            continue

        match = re.match(
            r"^([0-9a-fA-F]{64})\s+[* ]?(.*)$",
            line
        )

        if match:
            expected = match.group(1).lower()
            break

    if expected is None:
        raise RuntimeError(
            "Could not read a SHA-256 checksum from %s" %
            checksum_file
        )

    actual = sha256_file(iso_file).lower()
    return actual == expected


def run_minisign(iso_file, signature_file):
    """
    Verify using minisign.

    The original shell script uses:
      minisign -Vm ISO -P PUBLIC_KEY

    This assumes the signature file is named in a way that minisign
    automatically discovers, or that the signature has been saved beside
    the ISO using the expected .minisig suffix.
    """

    command = [
        "minisign",
        "-Vm",
        iso_file,
        "-P",
        MINISIGN_PUBLIC_KEY
    ]

    process = subprocess.Popen(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    stdout, stderr = process.communicate()

    if process.returncode != 0:
        if stderr:
            error(text_value(stderr).strip())
        return False

    if stdout:
        print(text_value(stdout).strip())

    return True


def verify_iso(iso_file, signature_file):
    checksum_file = iso_file + ".sha256"

    if os.path.isfile(checksum_file):
        log("Checking SHA-256 checksum ...")

        if not verify_sha256(iso_file, checksum_file):
            error("SHA-256 mismatch!")
            return False

        log("SHA-256 OK")
        return True

    if os.path.isfile(signature_file):
        log("Running minisign verification ...")

        if not run_minisign(iso_file, signature_file):
            error("minisign verification failed!")
            return False

        log("minisign verification OK")
        return True

    error("No checksum or signature file found to verify the ISO.")
    return False


def ask_confirmation():
    try:
        answer = raw_input(
            "Do you want to install this image now? (y/N) "
        )
    except NameError:
        answer = input(
            "Do you want to install this image now? (y/N) "
        )

    return re.match(r"^[Yy]$", answer.strip()) is not None


def main():
    parser = optparse.OptionParser(
        usage="%prog [-u URL] [-s] [-n N] [-y] [-h]"
    )

    parser.add_option(
        "-u",
        "--url",
        dest="page_url",
        default=DEFAULT_PAGE_URL,
        help="URL containing the build list"
    )

    parser.add_option(
        "-s",
        "--install",
        dest="do_install",
        action="store_true",
        default=False,
        help="Print commands for installing the selected image"
    )

    parser.add_option(
        "-n",
        "--number",
        dest="select_num",
        type="int",
        default=1,
        help="Select the N-th build; default is 1"
    )

    parser.add_option(
        "-y",
        "--yes",
        dest="auto_confirm",
        action="store_true",
        default=False,
        help="Skip the installation confirmation prompt"
    )

    options, unused_args = parser.parse_args()

    if options.select_num < 1:
        error("Selection must be at least 1.")
        return 1

    log("Downloading page %s ..." % options.page_url)

    try:
        html = read_url(options.page_url)
    except Exception as exc:
        error("Failed to download page: %s" % exc)
        return 1

    parser_html = BuildParser()

    try:
        parser_html.feed(html)
        parser_html.close()
    except Exception as exc:
        error("Failed to parse page: %s" % exc)
        return 1

    builds = parser_html.builds

    if not builds:
        error("No ISO entries were found on the page.")
        return 1

    # Resolve relative URLs against the page URL.
    for build in builds:
        build["iso_url"] = urljoin(
            options.page_url,
            build["iso_url"]
        )
        build["sig_url"] = urljoin(
            options.page_url,
            build["sig_url"]
        )

    log("Found %d nightly builds:" % len(builds))

    for index, build in enumerate(builds):
        print("  %d) %s" % (index + 1, build["name"]))

    if options.select_num > len(builds):
        error(
            "Selection %d is out of range (1-%d)" %
            (options.select_num, len(builds))
        )
        return 1

    selected = builds[options.select_num - 1]

    iso_name = selected["name"]
    iso_url = selected["iso_url"]
    sig_url = selected["sig_url"]

    log("You selected: %s" % iso_name)
    log("ISO URL   : %s" % iso_url)
    log("SIG URL   : %s" % sig_url)

    iso_basename = os.path.basename(urlparse(iso_url).path)
    sig_basename = os.path.basename(urlparse(sig_url).path)

    iso_file = "vyos-" + iso_basename
    sig_file = "vyos-" + sig_basename

    log("Downloading ISO ...")

    try:
        download_file(iso_url, iso_file)
    except Exception as exc:
        error("ISO download failed: %s" % exc)
        return 1

    log("Downloading signature ...")

    try:
        download_file(sig_url, sig_file)
    except Exception as exc:
        error("Signature download failed: %s" % exc)
        return 1

    log("Verifying ISO ...")

    if not verify_iso(iso_file, sig_file):
        return 1

    log("ISO verification succeeded.")

    if options.do_install:
        if not options.auto_confirm:
            if not ask_confirmation():
                log("Installation aborted by user.")
                return 0

        absolute_iso_path = os.path.abspath(iso_file)

        log("Commands to run inside the VyOS CLI:")
        print("add system image %s" % absolute_iso_path)
        print("set system image default-boot %s" % iso_name)
        print("commit")
        print("save")
        print("---------------------------------------------------")

        if options.auto_confirm:
            log(
                "Commands were printed only; run them inside the VyOS CLI."
            )

    log(
        "All done. Keep the ISO and its signature file; "
        "they can be used again later."
    )

    return 0


if __name__ == "__main__":
    sys.exit(main())