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

# Copyright (C) 2016 Canonical Ltd.
# Author: Łukasz 'sil2100' Zemczak <lukasz.zemczak@ubuntu.com>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import logging
import os
import sys

sys.path.insert(0, os.path.join(sys.path[0], os.pardir, "lib"))

from systemimage import config, tree, gpg


def resign_and_replace(conf, path, sign_path):
    '''Helper function for re-signing files.'''
    new_path = "%s.new" % sign_path
    gpg.sign_file(
        conf, "image-signing", path, new_path)
    os.rename(new_path, sign_path)


def main():
    parser = argparse.ArgumentParser(
        description="Helper script for re-signing all indexes and their "
                    "latest images with the new image-signing key in the "
                    "keyring.")
    parser.add_argument("--verbose", "-v", action="count", default=0)

    args = parser.parse_args()

    # Setup logging
    formatter = logging.Formatter(
        "%(asctime)s %(levelname)s %(message)s")

    levels = {1: logging.ERROR,
              2: logging.WARNING,
              3: logging.INFO,
              4: logging.DEBUG}

    if args.verbose > 0:
        stdoutlogger = logging.StreamHandler(sys.stdout)
        stdoutlogger.setFormatter(formatter)
        logging.root.setLevel(levels[min(4, args.verbose)])
        logging.root.addHandler(stdoutlogger)
    else:
        logging.root.addHandler(logging.NullHandler())

    # Load the configuration
    conf = config.Config()

    # Load the tree
    pub = tree.Tree(conf)

    # List of files that we have re-signed already
    done = set()

    # Also re-sign the channel.json file
    logging.info(
        "Re-signing the channels index file: %s..." % pub.indexpath)
    channels_sign_path = "%s.asc" % pub.indexpath
    resign_and_replace(conf, pub.indexpath, channels_sign_path)

    channels = pub.list_channels()
    for channel in channels:
        # Ignore redirects as the target channels will already be handled
        if "redirect" in channels[channel]:
            logging.info("Skipping %s, channel redirect." % channel)
            continue

        for device in pub.list_devices(channel):
            logging.info("Looking at %s/%s." % (channel, device))
            device = pub.get_device(channel, device)
            images = device.list_images()

            # Re-sign the index.json file as well
            logging.info(
                "Re-signing the device index file: %s..." % device.indexpath)
            index_sign_path = "%s.asc" % device.indexpath
            resign_and_replace(conf, device.indexpath, index_sign_path)

            if len(images) > 0:
                latest_version = images[-1]['version']
                for image in reversed(images):
                    # Only re-sign full files for the latest image
                    # but resign all deltas
                    if image['version'] != latest_version and image['type'] != "delta":
                        continue

                    for f in image['files']:
                        # Re-sign the tarball and .json files with the new keys
                        tar_path = "%s%s" % (conf.publish_path, f['path'])
                        tar_sign_path = "%s%s" % (conf.publish_path,
                                                  f['signature'])
                        json_path = tar_path.replace(".tar.xz", ".json")
                        json_sign_path = tar_sign_path.replace(
                            ".tar.xz", ".json")

                        if tar_path not in done:
                            if os.path.exists(tar_path):
                                logging.info("Re-signing %s..." % tar_path)
                                resign_and_replace(
                                    conf, tar_path, tar_sign_path)

                            if os.path.exists(json_path):
                                logging.info("Re-signing %s..." % json_path)
                                resign_and_replace(
                                    conf, json_path, json_sign_path)

                            # Mark the file as done in case more than one index
                            # points to it
                            done.add(tar_path)
            else:
                logging.info("No images in channel.")


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