#!/usr/bin/env python3
"""
Creates a KeyLime Hashlist for an image.

See: https://keylime.dev/
"""

import argparse
import datetime
import json
import os
import shutil
import subprocess
import sys
import tempfile

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, cosa_dir)

from cosalib import meta
from cosalib.builds import Builds
from cosalib.cmdlib import (
    ensure_glob,
    get_basearch,
    sha256sum_file,
    import_ostree_commit)


class HashListV1(dict):
    """
    Abstraction of a HashList in the version 1 series

    See: https://github.com/keylime/enhancements/blob/master/16_remote_allowlist_retrieval.md
    """

    def __init__(self, metadata, release=None, generator=None, timestamp=None):
        """
        Initialize HashListV1 instance.

        :param metadata: Loaded metadata instance
        :type metadata: GenericBuildMeta
        :param release: Optional release override
        :type release: str
        :param generator: Optional generator override
        :type generator: str
        :param timestamp: Optional timestamp override
        :type timestamp: str
        :raises: KeyError
        """
        super()
        self._metadata = metadata
        if release is None:
            release = self._metadata['buildid']
        self['release'] = release

        self['meta'] = {}
        if generator is None:
            generator = 'coreos-assembler-' + self._metadata['coreos-assembler.container-image-git']['commit']  # noqa
        self['meta']['generator'] = generator
        if timestamp is None:
            timestamp = datetime.datetime.utcnow().isoformat()
        self['meta']['timestamp'] = timestamp
        self['meta']['ostree'] = self._metadata['images']['ostree']['path']
        self['meta']['arch'] = self._metadata['coreos-assembler.basearch']

        self['hashes'] = {}

    def populate_hash_list(self):
        """
        Populates the hashlist instance with ostree and initramfs hashes.

        :param checksum: The ostree checksum from the metadata
        :type checksum: str
        :param commit: The ostree commit
        :type commit: str
        :param hashlist: Initialized hash list instance
        :type hashlist: HashListV1
        :returns: Nothing
        :rtype: None
        :raises: IndexError
        """
        checkout = 'tmp/repo/tmp/keylime-checkout'
        if os.path.isdir(checkout):
            shutil.rmtree(checkout)

        import_ostree_commit(
            os.getcwd(),
            self._metadata.build_dir,
            self._metadata)
        subprocess.check_call([
            'ostree', 'checkout',
            '--repo=tmp/repo', '-U',
            self._metadata['ostree-commit'], checkout])
        #  Make all dirs in 'tmp' checkout readable.
        #  'find' can't recurse into tmp dir files, we add 'x' so that it can recurse into it.
        #  With '+', 'find' accumulates paths and runs 'chmod' once (or batched by the arg limit).
        #  To 'chmod' the dir before it tries to recurse into it, we use ';'.
        subprocess.check_call(['find', checkout, '-type', 'd', '-exec',
                               'chmod', 'u+rwx', '{}', ';'])
        self.hash_from_path(checkout)

        # Extract initramfs contents
        initramfs_path = ensure_glob(
            os.path.join(
                checkout, 'usr/lib/modules/*/initramfs.img'))[0]
        initramfs_path = os.path.realpath(initramfs_path)

        with tempfile.TemporaryDirectory() as tmpdir:
            subprocess.check_call(['lsinitrd', '--unpack', initramfs_path],
                                  cwd=tmpdir)
            self.hash_from_path(tmpdir)

        shutil.rmtree(checkout)

    def hash_from_path(self, toppath):
        """
        Create hashes from files starting at a specific path.

        .. note:: If a file can not be read it is skipped

        :param toppath: The starting path to go through
        :type toppath: str
        :returns: Nothing
        :rtype: None
        """
        for dirpath, _, filenames in os.walk(toppath):
            for fname in filenames:
                filepath = os.path.join(dirpath, fname)
                # Skip symlinks
                if os.path.islink(filepath):
                    continue
                try:
                    filehash = sha256sum_file(filepath)
                    relpath = filepath.replace(toppath, '')
                    self['hashes'][relpath] = [filehash]
                except (PermissionError, FileNotFoundError) as err:
                    print(f'Unable to hash {filepath}: {err}')

    def write(self, output):
        """
        Renders the HashList structure to disk.

        :param output: Where to write the file
        :type output: str
        :raises: FileNotFoundError, IsADirectoryError, PermissionError
        """
        with open(output, 'w') as out:
            out.write(json.dumps(self, indent=4))
        with open(f'{output}-CHECKSUM', 'w') as out:
            subprocess.check_call(['sha256sum', output], stdout=out)


def main():
    """
    Main entry point.
    """
    parser = argparse.ArgumentParser(description=__doc__)
    # parser.add_argument(
    #     '-s', '--sign',
    #     help='If set to a key a signature will be created as well')
    parser.add_argument(
        '-g', '--generator', help='Override the generator name')
    parser.add_argument(
        '-t', '--timestamp', help='Override the timestamp')
    parser.add_argument(
        '-r', '--release', help='Override the release')
    parser.add_argument(
        '-a', '--arch', default=get_basearch(),
        help='Target Architecture')
    parser.add_argument('-b', '--build', default='latest', help='Target build')

    # Parse arguments and ignore anything we didn't explicitly request
    args = parser.parse_args()

    builds = Builds()
    builddir = builds.get_build_dir(build_id=args.build, basearch=args.arch)

    # it's not in the schema nor in meta.json yet
    output = os.path.join(builddir, "exp-hash.json")

    # Load metadata
    metadata = meta.GenericBuildMeta(schema=None, basearch=args.arch)

    # Step 0: Initialize the HashList instance
    hashlist = HashListV1(
        metadata, args.release, args.generator, args.timestamp)

    # Step 1: Populate the hashlist
    hashlist.populate_hash_list()

    # Step 2: Write the hashlist to disk
    hashlist.write(output)
    print(f'Hash list created at {output}')


if __name__ == '__main__':
    main()
