Asset-Compression-Manager/acm.py

#!/usr/bin/env python
import asyncio
import hashlib
import io
import json
import os
import platform
import sys
import tempfile
from typing import List, Dict, Callable

import click
from minio import Minio, ResponseError
from minio.error import NoSuchKey

# Size of the buffer to read files with
BUF_SIZE = 4096


###########
# AsyncIO #
###########


async def run_command_shell(
        command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, on_success: Callable = ()):
    """Run command in subprocess (shell).

    Note:
        This can be used if you wish to execute e.g. "copy"
        on Windows, which can only be executed in the shell.
    """
    process = await asyncio.create_subprocess_shell(
        command, stdout=stdout, stderr=stderr
    )

    process_stdout, process_stderr = await process.communicate()

    if process.returncode == 0:
        on_success()

    if stdout != asyncio.subprocess.DEVNULL:
        result = process_stdout.decode().strip()
        return result
    else:
        return None


def make_chunks(tasks, chunk_size):
    """Yield successive chunk_size-sized chunks from tasks.

    Note:
        Taken from https://stackoverflow.com/a/312464
        modified for python 3 only
    """
    for i in range(0, len(tasks), chunk_size):
        yield tasks[i: i + chunk_size]


def run_asyncio_commands(tasks, max_concurrent_tasks=0):
    """Run tasks asynchronously using asyncio and return results.

    If max_concurrent_tasks are set to 0, no limit is applied.

    Note:
        By default, Windows uses SelectorEventLoop, which does not support
        subprocesses. Therefore ProactorEventLoop is used on Windows.
        https://docs.python.org/3/library/asyncio-eventloops.html#windows
    """
    all_results = []

    if max_concurrent_tasks == 0:
        chunks = [tasks]
        num_chunks = len(chunks)
    else:
        chunks = make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)
        num_chunks = len(list(make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)))

    if asyncio.get_event_loop().is_closed():
        asyncio.set_event_loop(asyncio.new_event_loop())
    if platform.system() == "Windows":
        asyncio.set_event_loop(asyncio.ProactorEventLoop())
    loop = asyncio.get_event_loop()

    chunk = 1
    for tasks_in_chunk in chunks:
        commands = asyncio.gather(*tasks_in_chunk)
        results = loop.run_until_complete(commands)
        all_results += results
        chunk += 1

    loop.close()
    return all_results


###########
# Helpers #
###########


def get_metadata_name(key):
    return METADATA_PREFIX + 'SHA256SUM'.capitalize()


def get_clean_stdin_iterator(stdin_stream):
    return (line for line in [line.strip() for line in stdin_stream if line.strip() != ''])


def strip_prefix(prefix: str, file: str) -> str:
    if file.startswith(prefix):
        return file.replace(prefix, '')
    return file


def get_file_identity(ctx_obj, file):
    if 'REMOVE_PREFIX' in ctx_obj and ctx_obj['REMOVE_PREFIX'] is not None:
        path = strip_prefix(ctx_obj['REMOVE_PREFIX'], file)
    else:
        path = file

    if os.pathsep != '/':
        path = '/'.join(path.split(os.pathsep))

    return path


def list_s3_dir(s3: Minio, bucket: str, prefix: str) -> List[str]:
    found_files = []
    for obj in s3.list_objects_v2(bucket, prefix=prefix):
        if obj.is_dir:
            found_files.extend(list_s3_dir(s3, bucket, obj.object_name))
        else:
            found_files.append(obj.object_name)
    return found_files


def get_s3_client(config: any) -> Minio:
    host = config['host']
    secure = config['secure']
    access_key = config['access']
    secret_key = config['secret']
    return Minio(host, secure=secure, access_key=access_key, secret_key=secret_key)


def prep_s3(ctx):
    s3_config = ctx.obj['CONFIG']['s3']
    s3_bucket = ctx.obj['CONTEXT']

    s3 = get_s3_client(s3_config)

    if not s3.bucket_exists(s3_bucket):
        s3.make_bucket(s3_bucket)

    return s3_bucket, s3


def get_file_sha256sum(stored_data, file):
    stored_file_hash = stored_data['sha256sum']
    stored_profile_hash = stored_data['profilesHash']
    sha256sum = hashlib.sha256()
    with open(file, 'rb') as f:
        for byte_block in iter(lambda: f.read(BUF_SIZE), b""):
            sha256sum.update(byte_block)
    calculated_file_hash = sha256sum.hexdigest()
    return stored_profile_hash, stored_file_hash, calculated_file_hash


def load_config(path: str) -> any:
    with open(path, 'r') as config_file:
        config = json.load(config_file)

        # Setup S3 Settings
        config['s3']['access'] = os.getenv('ACM_S3_ACCESS')
        config['s3']['secret'] = os.getenv('ACM_S3_SECRET')

        # Setup concurrency
        if 'concurrency' in config:
            config['concurrency'] = abs(int(config['concurrency']))
        else:
            config['concurrency'] = 0

        # Calculate profiles hash
        sha256sum = hashlib.sha256()
        with io.BytesIO(json.dumps(config['profiles']).encode('utf-8')) as c:
            for byte_block in iter(lambda: c.read(BUF_SIZE), b''):
                sha256sum.update(byte_block)
        profiles_hash = sha256sum.hexdigest()
        config['profilesHash'] = profiles_hash

        return config


@click.group()
@click.option('-d', '--debug/--no-debug', default=False)
@click.option('-c', '--config', default=lambda: os.path.join(os.getcwd(), 'acm-config.json'), show_default=True)
@click.option('-s', '--stdin/--no-stdin', default=False)
@click.option('--remove-prefix', default=None)
@click.option('--add-prefix', default=None)
@click.pass_context
def cli(ctx, debug, config, stdin, remove_prefix, add_prefix):
    ctx.ensure_object(dict)
    ctx.obj['DEBUG'] = debug
    ctx.obj['CONFIG'] = load_config(config)
    ctx.obj['READ_STDIN'] = stdin
    ctx.obj['REMOVE_PREFIX'] = remove_prefix
    ctx.obj['ADD_PREFIX'] = add_prefix


###############################
# S3 Storage Focused Commands #
###############################


@cli.command(name="list")
@click.option('--sha256sum/--no-sha256sum', default=False)
@click.option('--suffix', default=None)
@click.option('-x', '--context', required=True)
@click.pass_context
def list_files(ctx, context, sha256sum, suffix):
    ctx.obj['CONTEXT'] = context
    s3_config = ctx.obj['CONFIG']['s3']
    s3_bucket = ctx.obj['CONTEXT']

    s3 = get_s3_client(s3_config)

    if not s3.bucket_exists(s3_bucket):
        s3.make_bucket(s3_bucket)

    found_files: List[str] = []
    found_objects: List[str] = []

    for obj in s3.list_objects_v2(s3_bucket, recursive=False):
        if obj.is_dir:
            found_objects.extend(list_s3_dir(s3, s3_bucket, obj.object_name))
        else:
            found_objects.append(obj.object_name)

    for obj in found_objects:
        file = obj
        if 'REMOVE_PREFIX' in ctx.obj and ctx.obj['REMOVE_PREFIX'] is not None:
            file = os.path.join(ctx.obj['REMOVE_PREFIX'], file)

        if suffix is not None and suffix in file:
            file = file.replace(suffix, '')

        file = file.strip()

        if sha256sum:
            stat = s3.stat_object(s3_bucket, obj)
            sha256sum_value = stat.metadata[get_metadata_name("SHA256SUM")]
            file = f'{sha256sum_value}  {file}'

        found_files.append(file)

    print(os.linesep.join(found_files))


@cli.command(name="match")
@click.option('-x', '--context', required=True)
@click.option('--print-identity/--no-print-identity', default=False)
@click.argument('files', nargs=-1)
@click.pass_context
def check_matched_files_hashes(ctx, context, print_identity, files):
    """
    List all files that have matching stored sha256sum and profilesHash
    """
    ctx.obj['CONTEXT'] = context
    s3_bucket, s3 = prep_s3(ctx)
    matching_files: List[str] = []

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    for file in files:
        file_identity = f'{get_file_identity(ctx.obj, file)}.json'
        try:
            file_object = s3.get_object(s3_bucket, file_identity)
            stored_data = json.load(file_object)
            stored_profile_hash, stored_file_hash, calculated_file_hash = get_file_sha256sum(stored_data, file)
            if calculated_file_hash == stored_file_hash \
                    and ctx.obj['CONFIG']['profilesHash'] == stored_profile_hash:
                if print_identity:
                    matching_files.append(stored_data['storedAssetIdentity'])
                else:
                    matching_files.append(file)
        except NoSuchKey as e:
            continue
        except ValueError or ResponseError as e:
            print(f'ERROR: {file} {e}')

    print(os.linesep.join(matching_files))


@cli.command(name="check")
@click.option('-x', '--context', required=True)
@click.argument('files', nargs=-1)
@click.pass_context
def check_changed_files_hashes(ctx, context, files):
    """
    List all files that do not have a matching sha256sum or profilesHash
    """
    ctx.obj['CONTEXT'] = context
    s3_bucket, s3 = prep_s3(ctx)
    changed_files: List[str] = []

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    for file in files:
        file_identity = f'{get_file_identity(ctx.obj, file)}.json'
        try:
            file_object = s3.get_object(s3_bucket, file_identity)
            stored_data = json.load(file_object)
            stored_profile_hash, stored_file_hash, calculated_file_hash = get_file_sha256sum(stored_data, file)
            if calculated_file_hash != stored_file_hash \
                    or ctx.obj['CONFIG']['profilesHash'] != stored_profile_hash:
                changed_files.append(file)
        except NoSuchKey as e:
            changed_files.append(file)
        except ValueError or ResponseError as e:
            print(f'ERROR: {file} {e}')

    print(os.linesep.join(changed_files))


@cli.command(name="update")
@click.option('-x', '--context', required=True)
@click.option('--input-and-identity/--no-input-and-identity', default=False)
@click.argument('files', nargs=-1)
@click.pass_context
def update_changed_files_hashes(ctx, context, input_and_identity, files):
    """
    Store new data objects for the provided files
    """
    ctx.obj['CONTEXT'] = context
    s3_bucket, s3 = prep_s3(ctx)
    updated_files: List[str] = []

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    for file in files:
        identity = None
        if input_and_identity:
            file, identity = file.split('\t')
        file_identity = f'{get_file_identity(ctx.obj, file)}.json'
        try:
            sha256sum = hashlib.sha256()
            with open(file, 'rb') as f:
                for byte_block in iter(lambda: f.read(BUF_SIZE), b''):
                    sha256sum.update(byte_block)
            calculated_file_hash = sha256sum.hexdigest()

            object_data = {
                "sourcePath": file,
                "storedAssetIdentity": identity,
                "identity": file_identity,
                "sha256sum": calculated_file_hash,
                "profilesHash": ctx.obj['CONFIG']['profilesHash']
            }

            with io.BytesIO(json.dumps(object_data, sort_keys=True, indent=None).encode('utf-8')) as data:
                data.seek(0, os.SEEK_END)
                data_length = data.tell()
                data.seek(0)
                s3.put_object(
                    s3_bucket,
                    file_identity,
                    data,
                    data_length,
                    content_type="application/json",
                    metadata={}
                )
                updated_files.append(file)
        except ValueError or ResponseError as e:
            print(f'ERROR: {file} {e}')

    print(os.linesep.join(updated_files))


@cli.command(name="store")
@click.option('-x', '--context', required=True)
@click.argument('files', nargs=-1)
@click.pass_context
def store_files(ctx, context, files):
    """
    Store specified files in a <context> bucket for retrieval.
    """
    ctx.obj['CONTEXT'] = context
    s3_bucket, s3 = prep_s3(ctx)
    stored_files: List[str] = []

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    for file in files:
        file_identity = get_file_identity(ctx.obj, file)
        try:
            s3.fput_object(
                s3_bucket,
                file_identity,
                file,
                content_type="application/octet-stream"
            )
            if 'ADD_PREFIX' in ctx.obj and ctx.obj['ADD_PREFIX'] is not None:
                stored_files.append(os.path.join(ctx.obj['ADD_PREFIX'], file_identity))
            else:
                stored_files.append(file)
        except ResponseError as e:
            print(f'ERROR: {file} {e}', file=sys.stderr)

    print(os.linesep.join(stored_files))


@cli.command(name="retrieve")
@click.option('-x', '--context', required=True)
@click.option('-d', '--destination', default=None)
@click.argument('files', nargs=-1)
@click.pass_context
def retrieve_files(ctx, context, destination, files):
    """
    Retrieve specified files from a <context> bucket
    """
    ctx.obj['CONTEXT'] = context
    s3_bucket, s3 = prep_s3(ctx)
    retrieved_files: List[str] = []

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    for file in files:
        file_identity = get_file_identity(ctx.obj, file)
        file_destination = file
        if destination is not None:
            file_destination = os.path.join(destination, file_identity)
        try:
            s3.fget_object(
                s3_bucket,
                file_identity,
                file_destination
            )
            retrieved_files.append(file_destination)
        except NoSuchKey as e:
            print(f'ERROR: {file_identity} {file_destination} {e}', file=sys.stderr)
        except ResponseError as e:
            print(f'ERROR: {file_destination} {e}', file=sys.stderr)

    print(os.linesep.join(retrieved_files))


######################################
# Asset Compression Focused Commands #
######################################


@cli.command(name="compress")
@click.option('-p', '--profile', default='default')
@click.option('-c', '--content', default='all')
@click.option('-d', '--destination', default=None)
@click.option('--print-input-and-identity/--no-print-input-and-identity', default=False)
@click.argument('files', nargs=-1)
@click.pass_context
def compress_assets(ctx, profile, content, destination, print_input_and_identity, files):
    profiles = ctx.obj['CONFIG']['profiles']

    if profile not in profiles:
        raise ValueError(f'Unrecognized profile: {profile}')

    default_profile: Dict[str, any] = profiles['default']
    profile: Dict[str, any] = profiles[profile]

    if content != 'all':
        if content not in profile and content not in default_profile:
            raise ValueError(f'Unrecognized content: {content}')

    content_configurations = []

    if content == 'all':
        content_names: set = set()

        for content_name in profile.keys():
            content_names.add(content_name)
            content_configurations.append(profile[content_name])

        for content_name in default_profile.keys():
            if content_name not in content_names:
                content_names.add(content_name)
                content_configurations.append(default_profile[content_name])
    else:
        if content in profile:
            content_configurations.append(profile[content])
        else:
            content_configurations.append(default_profile[content])

    if ctx.obj['READ_STDIN']:
        files = get_clean_stdin_iterator(click.get_text_stream('stdin'))

    if destination is None:
        destination = tempfile.mkdtemp()

    compressed_files = []
    tasks = []

    def store_filename(storage_list: List[str], filename: str):
        """
        A simple lambda wrapper to asynchronously add processed files to the list

        :param storage_list:
        :param filename:
        :return:
        """
        return lambda: storage_list.append(filename)

    for input_file in files:
        for content_configuration in content_configurations:
            if any([input_file.endswith(extension) for extension in content_configuration['extensions']]):
                file = input_file
                if 'REMOVE_PREFIX' in ctx.obj and ctx.obj['REMOVE_PREFIX'] is not None:
                    file = strip_prefix(ctx.obj['REMOVE_PREFIX'], input_file)

                if 'preserveInputExtension' in content_configuration \
                        and content_configuration['preserveInputExtension']:
                    output_file = os.path.join(destination, file)
                else:
                    output_file_without_ext = os.path.splitext(os.path.join(destination, file))[0]
                    output_file = f'{output_file_without_ext}.{content_configuration["outputExtension"]}'

                output_file_identity = get_file_identity({'REMOVE_PREFIX': destination}, output_file)

                output_file_dir = os.path.dirname(output_file)
                os.makedirs(output_file_dir, exist_ok=True)

                command: str = content_configuration['command'] \
                    .replace('{{input_file}}', f'\'{input_file}\'') \
                    .replace('{{output_file}}', f'\'{output_file}\'')

                tasks.append(
                    run_command_shell(
                        command,
                        stdout=asyncio.subprocess.DEVNULL,
                        stderr=asyncio.subprocess.DEVNULL,
                        on_success=store_filename(
                            compressed_files,
                            f'{input_file}\t{output_file_identity}' if print_input_and_identity else output_file
                        )
                    )
                )

    results = run_asyncio_commands(
        tasks, max_concurrent_tasks=ctx.obj['CONFIG']['concurrency']
    )

    print(os.linesep.join(compressed_files))


if __name__ == '__main__':
    cli(obj={})