#!/usr/bin/env python3
# vim: fileencoding=utf-8 sts=4 sw=4 et
import argparse
import collections
import datetime
import dateutil.parser
import gzip
import json
import os
import packaging.version
import requests
import shlex
import subprocess
import sys
import urllib.parse

from gi.repository import GLib

INTERNAL_BASE_URL = "http://images.endlessm-sf.com"
PUBLIC_BASE_URL = "https://images-dl.endlessm.com"
KEYRING_FILENAME = "eos-image-keyring.gpg"
KEYRING_URL = PUBLIC_BASE_URL + "/" + KEYRING_FILENAME
KEYRING_SYSTEM_PATH = "/usr/share/keyrings/eos-image-keyring.gpg"

# CloudFront limits the maximum file size to 20 GB
# http://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/cloudfront-limits.html
# Our current CDN, Fastly, does not have this limitation, but fragmenting the
# download is harmless.
CDN_MAX_SIZE = 20 * 2 ** 30

TIMEOUT = (
    30,  # connect timeout, in seconds
    30,  # read timeout, in seconds, between bytes received
)


def log(*args):
    sys.stderr.write('\n')
    print(*args, file=sys.stderr)
    sys.stderr.flush()


def check_call(*args):
    if hasattr(shlex, 'quote'):
        qargs = map(shlex.quote, args)
    else:
        qargs = args
    log('$ ' + ' '.join(qargs))
    subprocess.check_call(args, stdout=sys.stderr)


def url_dirname(url):
    raw_path = urllib.parse.urlparse(url).path
    path = urllib.parse.unquote(raw_path)
    return os.path.dirname(path.lstrip('/'))


# Simple download progress output, but not when sending to a log file
if os.isatty(sys.stderr.fileno()):
    def progress(path, received, total):
        percent = int(100 * received / total) if total else 100
        if percent >= 100:
            size = GLib.format_size(total)
            end = "\n"
        else:
            size = f"{GLib.format_size(received)} / {GLib.format_size(total)}"
            end = ""
        erase_line = "\r\033[K"
        sys.stderr.write(f"{erase_line}{path} ({size}, {percent}%){end}")
else:
    def progress(path, received, total):
        pass


class EosDownloadImage(object):
    def download(self, src, dest, src_size=None, src_mtime=None):
        """Download a file using a Requests Session

        This acts like wget --continue or rsync where the sizes and
        modification times of the source and destination files are
        considered.

        This handles the case where the file is larger than the CDN allows
        by using a provided size and splitting up the requests across
        ranges.
        """
        # By default, requests sets Accept-Encoding to gzip, deflate, but we
        # want raw bytes
        headers = requests.utils.default_headers()
        headers['Accept-Encoding'] = 'identity'

        # Make a HEAD request to get the size and mtime
        resp = self.session.head(src, headers=headers, timeout=TIMEOUT,
                                 allow_redirects=True)
        if resp.status_code == 200:
            if src != resp.url:
                log("Followed redirect to", resp.url)
                src = resp.url

            src_mtime = dateutil.parser.parse(resp.headers['Last-Modified'])
            src_size = int(resp.headers['Content-Length'])
        elif resp.status_code == 403:
            # Assume this is CloudFront returning a 403 for a file exceeding
            # the maximum limit. In that case, the caller must supply the
            # size. If this is a real error, it will be hit below
            # when doing the real download request.
            if src_size is None:
                resp.raise_for_status()
        else:
            resp.raise_for_status()

        # Check the status of an existing download
        dest_size = 0
        if os.path.exists(dest):
            dest_stat = os.stat(dest)
            dest_size = dest_stat.st_size

            if src_mtime is None:
                # Assume it's a partial download; the signature may prove that wrong
                outdated = False
            else:
                dest_mtime = datetime.datetime.fromtimestamp(dest_stat.st_mtime,
                                                             datetime.timezone.utc)

                # Is the existing file outdated? Strip off microseconds since
                # one or the other is likely to not support it
                src_mtime = src_mtime.replace(microsecond=0)
                dest_mtime = dest_mtime.replace(microsecond=0)
                outdated = dest_mtime < src_mtime

            # Is the times and sizes match, the file is already downloaded.
            # Delete an outdated file or one that's larger then the source.
            # Assume that a smaller file is a partial download.
            if not outdated and dest_size == src_size:
                return
            elif outdated or dest_size > src_size:
                os.unlink(dest)
                dest_size = 0

        # Special case - empty file. Don't bother downloading since it
        # breaks the progress calculations below.
        if src_size == 0:
            with open(dest, 'w'):
                pass
            progress(dest, 0, 0)
            return

        # Break the download up into ranges if it exceeds the CDN maximum
        # size.
        ranges = collections.deque()
        start_byte = dest_size
        remaining = src_size - dest_size
        while remaining > 0:
            if remaining > CDN_MAX_SIZE:
                req_size = CDN_MAX_SIZE
            else:
                req_size = remaining
            ranges.append((start_byte, start_byte + req_size - 1))
            remaining -= req_size
            start_byte += req_size

        # Start the actual downloading
        with open(dest, 'ab') as f:
            # Make sure that the file size still matches what was used for
            # the calculations
            if f.tell() != dest_size:
                log('File', dest, 'modified outside', __file__)
                exit(1)

            received = dest_size
            progress(dest, received, src_size)
            while ranges:
                start, stop = ranges.popleft()
                headers['Range'] = 'bytes={}-{}'.format(start, stop)
                resp = self.session.get(src, headers=headers, stream=True,
                                        timeout=TIMEOUT)
                if resp.status_code not in (200, 206):
                    resp.raise_for_status()
                for chunk in resp.iter_content(chunk_size=(1024 ** 2)):
                    received += len(chunk)
                    f.write(chunk)
                    progress(dest, received, src_size)

                if received != stop + 1:
                    ranges.appendleft((received, stop))

        if received != src_size:
            log('Received up to', received, 'but target size was', src_size)
            exit(1)

    def ensure_keyring(self):
        if os.path.isfile(KEYRING_SYSTEM_PATH):
            return KEYRING_SYSTEM_PATH

        # The absolute path is needed so that gpgv doesn't try to use a
        # keyring from the gnupg homedir
        keyring = os.path.abspath(os.path.join(
            self.args.outdir, KEYRING_FILENAME))

        log("Fetching Endless keyring")
        self.download(KEYRING_URL, keyring)

        return keyring

    def download_and_verify(self, meta, what):
        destdir = self.args.outdir
        image_url = '{}/{}'.format(self.base, meta['file'])
        image_size = meta['compressed_size']
        if meta.get('last_modified'):
            image_mtime = dateutil.parser.parse(meta['last_modified'])
        else:
            image_mtime = None
        signature_url = '{}/{}'.format(self.base, meta['signature'])

        image_dir = signature_dir = ''
        if self.args.mirror:
            image_dir = url_dirname(image_url)
            signature_dir = url_dirname(signature_url)
        image = os.path.join(destdir, image_dir, os.path.basename(image_url))
        signature = os.path.join(destdir, signature_dir,
                                 os.path.basename(signature_url))

        log("Downloading", what, "image from", image_url)
        full_image_dir = os.path.dirname(image)
        if full_image_dir:
            os.makedirs(full_image_dir, exist_ok=True)
        self.download(image_url, image, src_size=image_size,
                      src_mtime=image_mtime)

        if meta.get('extracted_size'):
            with open(image + '.size', 'w') as f:
                f.write(str(meta['extracted_size']))

        log("Downloading", what, "image signature from", signature_url)
        full_signature_dir = os.path.dirname(signature)
        if full_signature_dir:
            os.makedirs(full_signature_dir, exist_ok=True)
        self.download(signature_url, signature)

        if meta.get('extracted_signature'):
            url = '{}/{}'.format(self.base, meta['extracted_signature'])
            ext_sig_dir = ''
            if self.args.mirror:
                ext_sig_dir = url_dirname(url)
            ext_signature = os.path.join(destdir, ext_sig_dir,
                                         os.path.basename(url))
            log("Downloading uncompressed", what, "image signature from", url)
            full_ext_sig_dir = os.path.dirname(ext_signature)
            if full_ext_sig_dir:
                os.makedirs(full_ext_sig_dir, exist_ok=True)
            self.download(url, ext_signature)

        keyring = self.ensure_keyring()

        log("Verifying")
        check_call("gpgv", "--keyring", keyring, signature, image)

        log("Download okay.")
        return image

    def __init__(self, session):
        self.session = session

        p = argparse.ArgumentParser(
            description='Fetches and verifies an Endless OS image (default); '
                        'fetches the Endless Installer for Windows; '
                        'or mirrors an entire Endless OS product.',
            epilog='In the default and --windows-tool modes, the path to the '
                   'downloaded file (and nothing else) is printed to stdout, '
                   'for consumption by other scripts.')
        p.add_argument('-o', '--outdir',
                       help='Output directory (default: current directory)',
                       default='')
        p.add_argument('-r', '--product',
                       help='Product (eg: eosinstaller; default: eos)',
                       default='eos')
        p.add_argument('-i', '--internal',
                       action='store_true',
                       help='Fetch images from the Endless internal network')
        p.add_argument('-s', '--size', type=int,
                       help='Expected size in bytes of compressed image (for '
                            'use with --url)')

        # Mode
        m = p.add_mutually_exclusive_group()
        m.add_argument('-m', '--mirror',
                       action='store_true',
                       help='Mirror all versions and personalities')
        m.add_argument('-w', '--windows-tool',
                       action='store_true',
                       help='Fetch the Windows USB creator/installer, '
                            'not an Endless OS image')
        m.add_argument('-u', '--url',
                       help='Fetch compressed image by URL')

        # Which image? NB. product is not part of this group, because it *is*
        # valid together with --mirror
        g = p.add_argument_group(title='image selection (ignored with '
                                       '--mirror or --windows-tool)')
        g.add_argument('-p', '--personality',
                       help='Image personality (default: base)',
                       default='base')
        g.add_argument('-v', '--version',
                       help='Image version (default: newest)')
        g.add_argument('-I', '--iso', action='store_true',
                       help='Fetch an ISO image, not a raw disk image')

        self.args = args = p.parse_args()

        if args.url is not None and args.url.startswith(INTERNAL_BASE_URL):
            args.internal = True

        self.base = INTERNAL_BASE_URL if args.internal else PUBLIC_BASE_URL

    def run(self):
        # Create the output directory if necessary
        if self.args.outdir:
            os.makedirs(self.args.outdir, exist_ok=True)

        if self.args.windows_tool:
            self.fetch_windows_tool()
        elif self.args.mirror:
            self.mirror_images()
        elif self.args.url:
            self.fetch_image_url()
        else:
            self.fetch_image()

    def fetch_windows_tool(self):
        '''Downloads the latest release of the Endless Installer for Windows
        (endless-installer.exe), replacing an existing file in outdir if needed.

        Prints the path to the downloaded file to stdout, to be consumed by
        eos-write-live-image.'''

        filename = 'endless-installer.exe'
        url = '{}/endless-installer/{}'.format(self.base, filename)
        dest = os.path.join(self.args.outdir, filename)
        log("Fetching Windows Installer")
        self.download(url, dest)
        print(dest)

    def fetch_manifest(self):
        manifest_url = '{base}/releases-{product}-3.json'.format(
            base=self.base, product=self.args.product)
        log("Fetching manifest from", manifest_url)
        response = self.session.get(manifest_url, timeout=TIMEOUT)
        if response.url != manifest_url:
            log("Followed redirect to", response.url)

        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            log("Couldn't fetch manifest for product", self.args.product)
            log(e)
            sys.exit(1)

        manifest = response.json()
        images = list(manifest['images'].values())
        images.sort(key=lambda i: packaging.version.Version(i['version']))
        return manifest, images

    def fetch_image_url(self):
        '''Downloads an image using the specified URL and size.'''
        args = self.args

        what = 'iso' if args.url.endswith('.iso') else 'full'
        # Create a fake meta object for download_and_verify()
        meta = {}
        meta['file'] = '/'.join(args.url.split('/')[3:])
        meta['signature'] = meta['file'] + '.asc'
        meta['compressed_size'] = args.size
        print(self.download_and_verify(meta, what=what))

    def fetch_image(self):
        '''Downloads a single OS image, and its signatures. If args.iso is True,
        fetches the ISO; if not, fetches the full raw image, plus the
        corresponding bootloader bundle.

        Prints the path to the downloaded OS image to stdout, to be consumed by
        eos-write-live-image.'''

        _, images = self.fetch_manifest()
        args = self.args

        if args.version is not None:
            for img in images:
                if img['version'] == args.version:
                    break
            else:
                log("Image version", repr(args.version), "not found")
                log("Known versions:", *[i['version'] for i in images])
                sys.exit(1)
        else:
            img = images[-1]

        personalities = img['personality_images']
        if args.personality not in personalities:
            log("Personality", repr(args.personality), "not found")
            log("Known personalities:", *sorted(personalities.keys()))
            sys.exit(1)

        personality = personalities[args.personality]
        if args.iso:
            what = 'iso'
        else:
            what = 'full'
            if 'boot' in personality and personality['boot']['file']:
                # We deliberately don't print() the path returned here. This
                # script is run by eos-write-live-image, which expects it to
                # print a single image path to stdout. The boot bundle is
                # found by replacing the file extension.
                self.download_and_verify(personality['boot'], what='boot')

        if what not in personality:
            log(what, "image not found for", args.product, img['version'],
                args.personality)
            sys.exit(1)

        meta = personality[what]
        print(self.download_and_verify(meta, what=what))

    def mirror_images(self):
        '''Mirror all images for the given product.'''
        manifest, images = self.fetch_manifest()

        for img in images:
            personalities = img['personality_images']

            for _name, personality in sorted(personalities.items()):
                for meta in sorted(personality):
                    print(self.download_and_verify(personality[meta],
                                                   what=meta))

        # Write out the manifest when mirroring after all images downloaded
        out_manifest = os.path.join(self.args.outdir,
                                    'releases-{}-3.json'.format(
                                        self.args.product))
        out_manifest_gz = out_manifest + '.gz'
        with open(out_manifest, 'w') as f:
            json.dump(manifest, f, sort_keys=True)
        with gzip.open(out_manifest_gz, 'wt') as f:
            json.dump(manifest, f, sort_keys=True, indent=4)


def main():
    with requests.sessions.session() as session:
        app = EosDownloadImage(session)
        app.run()


if __name__ == '__main__':
    main()
