#!/usr/bin/python3
#
# apt-get install python3-requests-oauthlib
#
# based on script from popey
# oauth stuff from jamestait:
# http://bazaar.launchpad.net/~jamestait/+junk/click-support-tools/files
#
# Basic usage: see ./store-fetch --help

from oauthlib.oauth1.rfc5849 import Client

import argparse
import getpass
import glob
import json
import os
import requests
import stat
import sys
import logging

# Configuration
logging.basicConfig(level=logging.INFO)
# logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)
# logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger()
log.name = 'store-sync'

# search_all_pkgs_uri = 'https://search.apps.ubuntu.com/api/v1/search?q=architecture:armhf,price:0&size=1000&page=1'
# search_all_pkgs_uri = 'https://search.apps.ubuntu.com/api/v1/search?q=price:0&size=2000&page=1'
# give me all of the free ones
search_all_pkgs_uri = 'https://search.apps.ubuntu.com/api/v1/search?q=price:0&page=1&size=100000'
# search_all_pkgs_uri = 'https://search.apps.ubuntu.com/api/v1/search?q=price:0&page=1&size=5'

download_pkg_uri = 'https://search.apps.ubuntu.com/api/v1/package/'


class OAuthClient(object):

    @classmethod
    def credentials_from_file(cls, credentials_file):
        if not os.path.exists(credentials_file):
            logging.error("'%s' does not exist. Please " % credentials_file + \
                          "specify --update-credentials")
            sys.exit(1)
        with open(credentials_file, 'r') as f:
            tokens = json.loads(f.read())
        return cls(**tokens)

    @classmethod
    def credentials_prompt(cls, credentials_file=None):
        print("Please enter your credentials (will not echo)")
        email = getpass.getpass("Username> ")
        password = getpass.getpass("Password> ")
        otp = getpass.getpass("OTP (blank to skip)> ")

        data = dict()
        data["email"] = email
        data["password"] = password
        data["token_name"] = "store-fetch"
        if otp != '':
            data["otp"] = otp

        resp = requests.post('https://login.ubuntu.com/api/v2/tokens/oauth',
                             data=json.dumps(data),
                             headers={'Content-Type': 'application/json',
                                      'accept': 'application/json'})
        if resp.status_code not in (200, 201):
            logging.error("Error {} logging in.".format(resp.status_code))
            exit(1)
        logging.info("Login OK.")
        auth = resp.json()
        tokens = dict()
        tokens['consumer_key'] = auth['consumer_key']
        tokens['consumer_secret'] = auth['consumer_secret']
        tokens['token_key'] = auth['token_key']
        tokens['token_secret'] = auth['token_secret']

        if credentials_file is not None:
            s = json.dumps(tokens, sort_keys=True, indent=4,
                           separators=(',',':'))
            mode = stat.S_IRUSR | stat.S_IWUSR
            orig = os.umask(0)
            try:
                fd = os.fdopen(os.open(credentials_file,
                                       os.O_WRONLY | os.O_CREAT, mode), 'w')
            finally:
                os.umask(orig)
            fd.write(s)
            fd.close()
        return cls(**tokens)

    def __init__(self, consumer_key, consumer_secret, token_key, token_secret):
        self.consumer_key = consumer_key
        self.consumer_secret = consumer_secret
        self.token_key = token_key
        self.token_secret = token_secret

    def do_request(self, url):
        client = Client(self.consumer_key,
                        self.consumer_secret,
                        self.token_key,
                        self.token_secret)
        uri, headers, body = client.sign(url)
        resp = requests.get(uri, headers=headers)
        return resp


def get_json_for_uri(store_uri, arch=None):
    headers = None
    if arch is not None:
        architecture_header = ('X-Ubuntu-Architecture', arch)
        headers = dict([architecture_header])

    log.debug(('Getting %s' % store_uri))
    resp = requests.get(store_uri, headers=headers)
    if resp.status_code != 200:
        logging.error("Error {} retrieving json.".format(resp.status_code))
        sys.exit(1)
    json_content = resp.json()
    return json_content


def get_store_app(download_dir, pkgname, auth):
    """
    Get single store app from the store.
    """
    detail_url = download_pkg_uri + pkgname
    json_content = get_json_for_uri(detail_url)
    store_file = json_content['download_url'].split('/')[-1]
    target_file = os.path.join(download_dir, store_file)
    download_url = json_content['download_url']

    if os.path.exists(target_file):
        log.warn("Skipping '%s' (already present)" % store_file)
        sys.exit(1)

    log.info(download_url)
    resp = auth.do_request(download_url)
    if resp.status_code != 200:
        logging.error("Error {} retrieving package info.".format(
            resp.status_code))
        sys.exit(1)

    with open(target_file, 'wb') as f:
        f.write(resp.content)

    download_list.append(store_file)


def get_store_apps(download_dir, auth, architecture=None, ext=None):
    """
    Gets store apps from the store.
    """
    architectures = ['amd64', 'i386', 'armhf']
    if architecture is not None:
        architectures = [architecture]

    extension = ext
    if ext is not None:
        extension = ext.split('.')[-1]

    download_list = []
    existing = set(glob.glob("%s/*" % download_dir))
    for arch in architectures:
        log.debug(('Getting %s for arch=%s' % (search_all_pkgs_uri, arch)))
        all_pkgs_json = get_json_for_uri(search_all_pkgs_uri, arch)
        packages = all_pkgs_json.get('_embedded', {}).get('clickindex:package',
                                                          [])
        packages = sorted(packages, key=lambda k: k['name'])
        log.debug("Found %d packages for arch=%s" % (len(packages), arch))
        # log.debug(json.dumps(packages, sort_keys=True,
        #                      indent=4, separators=(',',':')))
        for p in packages:
            # Don't hit the download_url if we already have it on disk
            pkgname_version = "%s_%s_" % (p.get('name'), p.get('version'))
            pkgname_version_abs = os.path.join(download_dir, pkgname_version)
            if any(item.startswith(pkgname_version) for item in download_list):
                # log.info("Skipping '%s' (already downloaded this session)" %
                #          pkgname_version.rstrip('_'))
                continue
            elif any(item.startswith(pkgname_version_abs) for item in existing):
                log.warn("Skipping '%s' (already present)" %
                         pkgname_version.rstrip('_'))
                continue
            elif extension == "snap" and "release" not in p or \
                    extension == "click" and "release" in p:
                log.warn("Skipping '%s' (not '%s')" %
                         (pkgname_version.rstrip('_'), extension))
                continue

            detail_url = p.get('_links', {}).get('self', {}).get('href')
            if not detail_url:
                continue
            json_content = get_json_for_uri(detail_url)
            store_file = json_content['download_url'].split('/')[-1]
            target_file = os.path.join(download_dir, store_file)
            download_url = json_content['download_url']

            # This shouldn't normally be hit due to the above, but have it just
            # in case something other than 'snap' or 'click' is used
            if extension is not None and \
                    not store_file.endswith("." + extension):
                log.warn("Skipping '%s' (not '%s')" % (store_file, extension))
                continue

            log.info(download_url)
            resp = auth.do_request(download_url)
            if resp.status_code != 200:
                logging.error("Error {} retrieving package info.".format(
                    resp.status_code))
                continue

            with open(target_file, 'wb') as f:
                f.write(resp.content)

            download_list.append(store_file)

    return download_list


def parse_arguments():
    parser = argparse.ArgumentParser(description='store package sync tool')
    parser.add_argument('download_dir',
                        help='''Path to store the downloaded store packages.
                                The path must exist''')
    parser.add_argument('--package', default=None, metavar='PKG',
                        help='''Package to fetch. If omitted, download all''')
    parser.add_argument('--credentials-file', default=None,
                        help='''json dictionary with tokens for oauth. Eg
                             --credentials-file=~/.config/store-fetch.conf''')
    parser.add_argument('--update-credentials', default=False,
                        action='store_true',
                        help='''Create/update --credentials-file''')
    parser.add_argument('--arch', default=None,
                        help='''Limit to specified architecture''')
    parser.add_argument('--extension', default=None,
                        help='''Limit to the specified extension''')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()

    if not os.path.exists(args.download_dir):
        logging.error("'%s' does not exist. Please create")
        sys.exit(1)

    creds_file = None
    if args.credentials_file:
        creds_file = os.path.expanduser(args.credentials_file)

    if not args.update_credentials and not args.credentials_file:
        logging.error("Must specify --credentials-file with " +
                      "--update-credentials")
        sys.exit(1)
    elif args.credentials_file and not args.update_credentials:
        auth = OAuthClient.credentials_from_file(creds_file)
    else:
        auth = OAuthClient.credentials_prompt(creds_file)

    if args.package is not None:
        download_list = get_store_app(args.download_dir,
                                      args.package,
                                      auth)
    else:
        download_list = get_store_apps(args.download_dir,
                                       auth,
                                       args.arch,
                                       args.extension)

    store_list = os.path.join(args.download_dir, 'store_list')
    with open(store_list, 'w') as f:
        for i in download_list:
            f.write('%s\n' % i)