<reponame>uktrade/ftps-s3
import asyncio
import contextlib
import itertools


class ExpiringDict:

    def __init__(self, loop, seconds):
        self._loop = loop
        self._seconds = seconds
        self._store = {}

    def __getitem__(self, key):
        return self._store[key][0]

    def __setitem__(self, key, value):
        def delete():
            del self._store[key]

        if key in self._store:
            self._store[key][1].cancel()
            del self._store[key]

        delete_handle = self._loop.call_later(self._seconds, delete)
        self._store[key] = (value, delete_handle)

    def __contains__(self, key):
        return key in self._store


class ExpiringSet:

    def __init__(self, loop, seconds):
        self._loop = loop
        self._store = ExpiringDict(loop, seconds)

    def add(self, item):
        self._store[item] = True

    def __contains__(self, item):
        return item in self._store


def constant_time_compare(bytes_a, bytes_b):
    if len(bytes_a) != len(bytes_b):
        return False

    diffs = 0

    for char_a, char_b in zip(bytes_a, bytes_b):
        diffs |= char_a ^ char_b

    return diffs == 0


class Timeout(Exception):
    pass


@contextlib.asynccontextmanager
async def timeout(loop, max_time):

    cancelling_due_to_timeout = False
    current_task = asyncio.current_task()

    def cancel():
        nonlocal cancelling_due_to_timeout
        cancelling_due_to_timeout = True
        current_task.cancel()

    handle = loop.call_later(max_time, cancel)

    try:
        yield
    except asyncio.CancelledError:
        if not cancelling_due_to_timeout:
            raise
        raise Timeout()
    finally:
        handle.cancel()


def normalise_environment(key_values):
    ''' Converts denormalised dict of (string -> string) pairs, where the first string
        is treated as a path into a nested list/dictionary structure

        {
            "FOO__1__BAR": "setting-1",
            "FOO__1__BAZ": "setting-2",
            "FOO__2__FOO": "setting-3",
            "FOO__2__BAR": "setting-4",
            "FIZZ": "setting-5",
        }

        to the nested structure that this represents

        {
            "FOO": [{
                "BAR": "setting-1",
                "BAZ": "setting-2",
            }, {
                "BAR": "setting-3",
                "BAZ": "setting-4",
            }],
            "FIZZ": "setting-5",
        }

        If all the keys for that level parse as integers, then it's treated as a list
        with the actual keys only used for sorting

        This function is recursive, but it would be extremely difficult to hit a stack
        limit, and this function would typically by called once at the start of a
        program, so efficiency isn't too much of a concern.
    '''

    # Separator is chosen to
    # - show the structure of variables fairly easily;
    # - avoid problems, since underscores are usual in environment variables
    separator = '__'

    def get_first_component(key):
        return key.split(separator)[0]

    def get_later_components(key):
        return separator.join(key.split(separator)[1:])

    without_more_components = {
        key: value
        for key, value in key_values.items()
        if not get_later_components(key)
    }

    with_more_components = {
        key: value
        for key, value in key_values.items()
        if get_later_components(key)
    }

    def grouped_by_first_component(items):
        def by_first_component(item):
            return get_first_component(item[0])

        # groupby requires the items to be sorted by the grouping key
        return itertools.groupby(
            sorted(items, key=by_first_component),
            by_first_component,
        )

    def items_with_first_component(items, first_component):
        return {
            get_later_components(key): value
            for key, value in items
            if get_first_component(key) == first_component
        }

    nested_structured_dict = {
        **without_more_components, **{
            first_component: normalise_environment(
                items_with_first_component(items, first_component))
            for first_component, items in grouped_by_first_component(with_more_components.items())
        }}

    def all_keys_are_ints():
        def is_int(string_to_test):
            try:
                int(string_to_test)
                return True
            except ValueError:
                return False

        return all([is_int(key) for key, value in nested_structured_dict.items()])

    def list_sorted_by_int_key():
        return [
            value
            for key, value in sorted(
                nested_structured_dict.items(),
                key=lambda key_value: int(key_value[0])
            )
        ]

    return \
        list_sorted_by_int_key() if all_keys_are_ints() else \
        nested_structured_dict

import asyncio
import contextlib
import itertools


class ExpiringDict:

    def __init__(self, loop, seconds):
        self._loop = loop
        self._seconds = seconds
        self._store = {}

    def __getitem__(self, key):
        return self._store[key][0]

    def __setitem__(self, key, value):
        def delete():
            del self._store[key]

        if key in self._store:
            self._store[key][1].cancel()
            del self._store[key]

        delete_handle = self._loop.call_later(self._seconds, delete)
        self._store[key] = (value, delete_handle)

    def __contains__(self, key):
        return key in self._store


class ExpiringSet:

    def __init__(self, loop, seconds):
        self._loop = loop
        self._store = ExpiringDict(loop, seconds)

    def add(self, item):
        self._store[item] = True

    def __contains__(self, item):
        return item in self._store


def constant_time_compare(bytes_a, bytes_b):
    if len(bytes_a) != len(bytes_b):
        return False

    diffs = 0

    for char_a, char_b in zip(bytes_a, bytes_b):
        diffs |= char_a ^ char_b

    return diffs == 0


class Timeout(Exception):
    pass


@contextlib.asynccontextmanager
async def timeout(loop, max_time):

    cancelling_due_to_timeout = False
    current_task = asyncio.current_task()

    def cancel():
        nonlocal cancelling_due_to_timeout
        cancelling_due_to_timeout = True
        current_task.cancel()

    handle = loop.call_later(max_time, cancel)

    try:
        yield
    except asyncio.CancelledError:
        if not cancelling_due_to_timeout:
            raise
        raise Timeout()
    finally:
        handle.cancel()


def normalise_environment(key_values):
    ''' Converts denormalised dict of (string -> string) pairs, where the first string
        is treated as a path into a nested list/dictionary structure

        {
            "FOO__1__BAR": "setting-1",
            "FOO__1__BAZ": "setting-2",
            "FOO__2__FOO": "setting-3",
            "FOO__2__BAR": "setting-4",
            "FIZZ": "setting-5",
        }

        to the nested structure that this represents

        {
            "FOO": [{
                "BAR": "setting-1",
                "BAZ": "setting-2",
            }, {
                "BAR": "setting-3",
                "BAZ": "setting-4",
            }],
            "FIZZ": "setting-5",
        }

        If all the keys for that level parse as integers, then it's treated as a list
        with the actual keys only used for sorting

        This function is recursive, but it would be extremely difficult to hit a stack
        limit, and this function would typically by called once at the start of a
        program, so efficiency isn't too much of a concern.
    '''

    # Separator is chosen to
    # - show the structure of variables fairly easily;
    # - avoid problems, since underscores are usual in environment variables
    separator = '__'

    def get_first_component(key):
        return key.split(separator)[0]

    def get_later_components(key):
        return separator.join(key.split(separator)[1:])

    without_more_components = {
        key: value
        for key, value in key_values.items()
        if not get_later_components(key)
    }

    with_more_components = {
        key: value
        for key, value in key_values.items()
        if get_later_components(key)
    }

    def grouped_by_first_component(items):
        def by_first_component(item):
            return get_first_component(item[0])

        # groupby requires the items to be sorted by the grouping key
        return itertools.groupby(
            sorted(items, key=by_first_component),
            by_first_component,
        )

    def items_with_first_component(items, first_component):
        return {
            get_later_components(key): value
            for key, value in items
            if get_first_component(key) == first_component
        }

    nested_structured_dict = {
        **without_more_components, **{
            first_component: normalise_environment(
                items_with_first_component(items, first_component))
            for first_component, items in grouped_by_first_component(with_more_components.items())
        }}

    def all_keys_are_ints():
        def is_int(string_to_test):
            try:
                int(string_to_test)
                return True
            except ValueError:
                return False

        return all([is_int(key) for key, value in nested_structured_dict.items()])

    def list_sorted_by_int_key():
        return [
            value
            for key, value in sorted(
                nested_structured_dict.items(),
                key=lambda key_value: int(key_value[0])
            )
        ]

    return \
        list_sorted_by_int_key() if all_keys_are_ints() else \
        nested_structured_dict
Converts denormalised dict of (string -> string) pairs, where the first string         is treated as a path into a nested list/dictionary structure          {             "FOO__1__BAR": "setting-1",             "FOO__1__BAZ": "setting-2",             "FOO__2__FOO": "setting-3",             "FOO__2__BAR": "setting-4",             "FIZZ": "setting-5",         }          to the nested structure that this represents          {             "FOO": [{                 "BAR": "setting-1",                 "BAZ": "setting-2",             }, {                 "BAR": "setting-3",                 "BAZ": "setting-4",             }],             "FIZZ": "setting-5",         }          If all the keys for that level parse as integers, then it's treated as a list         with the actual keys only used for sorting          This function is recursive, but it would be extremely difficult to hit a stack         limit, and this function would typically by called once at the start of a         program, so efficiency isn't too much of a concern. # Separator is chosen to # - show the structure of variables fairly easily; # - avoid problems, since underscores are usual in environment variables # groupby requires the items to be sorted by the grouping key