# Copyright Amazon.com, Inc. and its affiliates. All Rights Reserved.
#
# 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 2
# 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/gpl-2.0.html>.

from functools import singledispatch
import logging
from xml.etree import ElementTree
import tempfile
import os
import io
from collections import namedtuple

import dnf
import libdnf
import gpg


DEFAULT_DISTRO_TITLE = 'your OS'

UPDATE_HEADER = """WARNING:
  A newer release of {distro_title} is available.

  Available Versions:
"""

UPDATE_INSTANCE_TEXT = """  Version {name}:
    Run the following command to upgrade to {version_string}:

      dnf upgrade --releasever={version_string}

    Release notes:
     {release_notes}
"""

UPGRADE_HEADER = """NOTE:
  A new version of {distro_title} is available!

  Available Upgrades:
"""

UPGRADE_INSTANCE_TEXT = """  Version {name}:
    See {release_notes}
    To upgrade, run:
      dnf upgrade --releasever={version_string}
"""

RPM_GPG_KEY_PATH = '/etc/pki/rpm-gpg'

ReleaseVersion = namedtuple('UpdateVersion', ('name', 'version_string', 'release_notes'))

class ReleaseNotification(dnf.Plugin):
    """Populates the release metadata and informs the user of available upgrade paths.
    """
    name = "release-notification"
    singleton = None

    def __init__(self, base, cli):
        """Sets up new instance and registers the os release command if were executing in the cli.
        """
        self.base = base
        self.cli = cli
        self.current_release = self.base.conf.releasever

        self.logger = logging.getLogger('dnf')
        self.cfg = None
        self.release_info = None

        if ReleaseNotification.singleton is None:
            ReleaseNotification.singleton = self

        if self.cli is not None:
            self.cli.register_command(CheckReleaseUpdateCommand)

    def _get_os_metadata(self, os_release_file_path="/etc/os-release", cached=True):
        """Retrieve the OS metadata from the /etc/os-release file.  Returns a dictionary of metadata values.  Returns
        an empty dictionary if no data was found.
        """
        if getattr(self, '_os_metadata', None) is None or cached is False:
            try:
                os_release_file = open(os_release_file_path, 'r')
            except OSError:
                self.logger.error(f"Unable to open os release file: {os_release_file_path}")
                return {}

            metadata = {}

            for line in os_release_file.readlines():
                name, value = line.split('=', 1)
                metadata[name] = value.strip()

            os_release_file.close()

            self._os_metadata = metadata

        return self._os_metadata
        
    @property
    def os_release_title(self):
        return self._get_os_metadata().get('NAME', DEFAULT_DISTRO_TITLE)

    def pre_config(self):
        """DNF hook for loading plugin config variables.
        """
        self.cfg = self.read_config(self.base.conf)

        if self.cfg and 'Main' in self.cfg.sections():
            release_info_url = self.cfg.get('Main', 'release_info_url', None)
            release_info_url = libdnf.conf.ConfigParser.substitute(release_info_url, self.base.conf.substitutions)
            try:
                gpgcheck = self.cfg.get('Main', 'gpgcheck', None)
                gpgcheck = gpgcheck.lower() in ['true', 'yes', '1']
            except KeyError as err:
                gpgcheck = True
            try:
                release_info_asc = self.cfg.get('Main', 'release_info_asc', None)
            except KeyError as err:
                release_info_asc = release_info_url + ".asc"
            release_info_asc = libdnf.conf.ConfigParser.substitute(release_info_asc, self.base.conf.substitutions)

            self.release_info = ReleaseUpdateInfo(self.base, release_info_url, gpgcheck, release_info_asc)

    def transaction(self):
        """Reads the plugin config and populates the release metadata. If execution is under the cli, output available
        releases to user.
        """
        if self.cli is None or 'osrelease' == self.cli.command.aliases[0]:
            return

        self.run_release_check()

    def run_release_check(self, output_separators=True, output_versions_only=False, output_latest_version_only=False):
        """Performs check on release info metadata and determines if an update is available.
        """
        if self.release_info is None:
            return

        try:
            release_data = self.release_info.retrieve_release_info(self.current_release)
        except RuntimeError as exp:
            self.logger.error(f'Error encountered while trying to retrieve release update information: {exp.args[0]}')
            return
        
        if not release_data:
            self.logger.error('Returned release info was invalid.')
            return

        if self.current_release not in release_data.get('releases', []):
            return

        relevant_release = release_data['releases'][self.current_release]

        update_releases, upgrade_releases = self.generate_release_versions(relevant_release)

        if output_latest_version_only:
            update_releases = update_releases[-1:]
            upgrade_releases = []

        output = []

        if output_versions_only is True:
            output_separators = False
            output = [release.version_string for release in update_releases + upgrade_releases]
        else:
            if update_releases:
                output.append(UPDATE_HEADER.format(distro_title=self.os_release_title))
                for update in update_releases:
                    output.append(UPDATE_INSTANCE_TEXT.format(**update._asdict()))

            if upgrade_releases:
                output.append(UPGRADE_HEADER.format(distro_title=self.os_release_title))
                for upgrade in upgrade_releases:
                    output.append(UPGRADE_INSTANCE_TEXT.format(**upgrade._asdict()))

        if output:
            if output_separators: self.logger.warning('=' * self.base.output.term.columns)

            self.logger.warning('\n'.join(output))

            if output_separators: self.logger.warning('=' * self.base.output.term.columns)

    def generate_release_versions(self, relevant_release):
        updates = []
        upgrades = []

        for update in relevant_release['update']:
                updates.append(
                    ReleaseVersion(update['name'], update['version_string'], update['release_notes'])
                    )

        for upgrade in relevant_release['upgrade']:
                upgrades.append(
                    ReleaseVersion(upgrade['name'], upgrade['version_string'], upgrade['release_notes'])
                    )

        return (
                sorted(updates, key=lambda r: tuple(map(int, r.version_string.split('.')))),
                sorted(upgrades, key=lambda r: tuple(map(int, r.version_string.split('.'))))
            )

class CheckReleaseUpdateCommand(dnf.cli.Command):
    """Provides a manual command for checking for release updates.
    """
    aliases = ('check-release-update', 'relupdate')
    summary = ('check for updates or upgrades to the current release')

    @staticmethod
    def set_argparser(parser):
        parser.add_argument(
            '--version-only',
            action='store_true',
            help='List only version strings for OS updates.'
            )
        parser.add_argument(
            '--latest-only',
            action='store_true',
            help='Show only the latest OS update version.'
            )

    def configure(self):
        self.logger = logging.getLogger('dnf')
        
        demands = self.cli.demands

    def run(self):
        if(ReleaseNotification.singleton is None):
            self.logger.error(
                'Cannot check for release updates. The ReleaseNotification dnf plugin has not been properly initialized'
                )
            return

        ReleaseNotification.singleton.run_release_check(
            output_separators=False,
            output_versions_only=self.opts.version_only,
            output_latest_version_only=self.opts.latest_only,
            )


class ReleaseUpdateInfo():
    """Manages retrieving and validating release update information.
    """

    def __init__(self, dnf_base, release_info_url, gpgcheck, release_info_asc):
        self.dnf_base = dnf_base
        self.release_info_url = release_info_url
        self.release_info_asc = release_info_asc
        self.gpgcheck = gpgcheck

        self.logger = logging.getLogger('dnf')

    def _validate_release_info_signature(self, release_info_data, release_info_sig):
        """Validate the release info data against the provided signature.  Utilize the keys provided by the os rpm-gpg
        key store under /etc/pki/rpm-gpg.

            Returns True if the signature matches a key.
            Returns False otherwise.
        """
        try:
            data_io = io.BytesIO(release_info_data)
            sig_io = io.BytesIO(release_info_sig)
        except TypeError as exp:
            data_io = io.BytesIO(release_info_data.encode())
            sig_io = io.BytesIO(release_info_sig.encode())

        try:
            key_dir_entries = os.scandir(RPM_GPG_KEY_PATH)
        except FileNotFoundError as exp:
            return False

        # use a temp dir for gpg home to ensure no unexpected keys are loaded
        with tempfile.TemporaryDirectory(prefix='tmp.{}'.format(os.getpid())) as gpg_home_dir:
            gpg_context = gpg.Context()
            gpg_context.home_dir = gpg_home_dir

            for dir_entry in key_dir_entries:
                if not dir_entry.is_file() or not os.access(dir_entry.path, os.R_OK):
                    continue

                gpg_context.key_import(open(dir_entry.path))

            try:
                _, result = gpg_context.verify(data_io, sig_io)
            except gpg.errors.BadSignatures as exp:
                self.logger.error('Unable to verify release info signature.')
                return False

        return True

    def _process_release_info(self, release_info_data):
        """Process xml release info data.
            1. Validate data against accepted schema (TODO)
            2. Translate validated data to dictionary.
        
            Returns dictionary of release info data.
            Returns empty dictionary on failure.
        """

        # assume correctedly formatted xml data; should be validated via schema before now

        try:
            tree = ElementTree.fromstring(release_info_data)
        except ElementTree.ParseError as exp:
            self.logger.error('Unable to parse release info data.')
            return {}

        release_data = {}

        for release_node in tree.findall('./releases/release'):
            release_version = release_node.attrib['version']
            if not release_version in release_data:
                release_data[release_version] = {'update': [], 'upgrade': []}

            for update_child in release_node.findall('./update') + release_node.findall('./upgrade'):
                release_data[release_version][update_child.tag.lower()].append(
                    {
                        'name': update_child.findtext('./name'),
                        'version_string': update_child.findtext('./version_string'),
                        'release_notes': update_child.findtext('./release_notes'),
                    }
                )

        # we pack the release data into another dictionary to leave room to return other types of data in the future
        return {'releases': release_data}

    def retrieve_release_info(self, target_release):
        """Retrieve the OS release metadata from the configured locations.
        """
        if self.release_info_url is None or self.release_info_asc is None:
            raise RuntimeError('Release info location is invalid.') 

        try:
            with self.dnf_base.urlopen(self.release_info_url.format(releasever=target_release)) as fh:
                release_info_data = fh.read()
        except (RuntimeError, OSError) as exp:
            raise RuntimeError(f'Unable to retrieve release info data. {exp.args[0]}')

        if self.gpgcheck:
            try:
                with self.dnf_base.urlopen(self.release_info_asc.format(releasever=target_release)) as fh:
                    release_info_asc = fh.read()
            except (RuntimeError, OSError) as exp:
                raise RuntimeError(f'Unable to retrieve release info signature. {exp.args[0]}')

        if not release_info_data:
            raise RuntimeError('Retrieved release info data was invalid.')

        if self.gpgcheck and not release_info_asc:
            raise RuntimeError('Retrieved release info signature was invalid.')

        if self.gpgcheck and not self._validate_release_info_signature(release_info_data, release_info_asc):
            raise RuntimeError('Unable to validate release info data against the provided signatures.')

        return self._process_release_info(release_info_data)