From b06fa85c545957801629a9d63f0ac4beb71f2c85 Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Mon, 9 May 2022 15:28:56 -0400 Subject: [PATCH 1/7] WIP adding in basic env resolver. --- spock/backend/field_handlers.py | 4 +- spock/backend/resolvers.py | 90 +++++++++++++++++++++++++++++++++ spock/backend/typed.py | 10 +++- spock/exceptions.py | 6 +++ 4 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 spock/backend/resolvers.py diff --git a/spock/backend/field_handlers.py b/spock/backend/field_handlers.py index d0f9d085..23455856 100644 --- a/spock/backend/field_handlers.py +++ b/spock/backend/field_handlers.py @@ -13,6 +13,7 @@ from attr import NOTHING, Attribute +from spock.backend.resolvers import parse_env_variables from spock.backend.spaces import AttributeSpace, BuilderSpace, ConfigSpace from spock.backend.utils import ( _get_name_py_version, @@ -132,7 +133,8 @@ def handle_optional_attribute_value( Returns: """ - attr_space.field = attr_space.attribute.default + cleaned_attribute = parse_env_variables(attr_space.attribute.default, attr_space.attribute.type) + attr_space.field = cleaned_attribute @abstractmethod def handle_optional_attribute_type( diff --git a/spock/backend/resolvers.py b/spock/backend/resolvers.py new file mode 100644 index 00000000..da536eef --- /dev/null +++ b/spock/backend/resolvers.py @@ -0,0 +1,90 @@ +# -*- coding: utf-8 -*- + +# Copyright FMR LLC +# SPDX-License-Identifier: Apache-2.0 + +"""Resolver functions for Spock""" + +import re +import os +from typing import Any + +from spock.exceptions import _SpockEnvResolverError +from spock.utils import _T + + +# ENV Resolver -- full regex is ^\${spock\.env:.*}$ +CLIP_ENV_PATTERN = r"^\${spock\.env:" +CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) +END_ENV_PATTERN = r"}$" +END_REGEX_OP = re.compile(END_ENV_PATTERN) +FULL_ENV_PATTERN = CLIP_ENV_PATTERN + r".*" + END_ENV_PATTERN +FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) + + +def parse_env_variables(value: Any, value_type: _T) -> Any: + # Check if it matches the regex + # if so then split and replace with the value from the env -- have to check here if the env variable actually + # exists -- if it doesn't then we need to raise an exception -- allow for None? + + # If it's a string we can check the regex + if isinstance(value, str): + # Check the regex + match_obj = FULL_REGEX_OP.fullmatch(value) + # if the object exists then we've matched a pattern and need to handle it + if match_obj is not None: + return _get_env_value(value, value_type) + # Regex doesn't match so just passthrough + else: + return value + # If it's not a string we can't resolve anything so just passthrough and let spock handle the value + else: + return value + + +def _handle_default(value: str): + env_value, default_value = value.split(',') + default_value = default_value.strip() + # Swap string None to type None + if default_value == "None": + default_value = None + return env_value, default_value + + +def _get_env_value(value: str, value_type: _T): + # Based on the start and end regex ops find the value the user set + env_str = END_REGEX_OP.split(CLIP_REGEX_OP.split(value)[1])[0] + # Attempt to split on a comma for a default value + split_len = len(env_str.split(',')) + # Default found if the len is 2 + if split_len == 2: + env_value, default_value = _handle_default(env_str) + # If the length is larger than two then the syntax is messed up + elif split_len > 2: + raise _SpockEnvResolverError( + f"Issue with environment variable syntax -- currently `{value}` has more than one `,` which means the " + f"optional default value cannot be resolved -- please use only one `,` separator within the syntax" + ) + else: + env_value = env_str + default_value = "None" + # Attempt to get the env variable + if default_value == "None": + maybe_env = os.getenv(env_value) + else: + maybe_env = os.getenv(env_value, default_value) + if maybe_env is None and default_value == "None": + raise _SpockEnvResolverError( + f"Attempted to get `{env_value}` from environment variables but it is not set -- please set this " + f"variable or provide a default via the following syntax ${{spock.env:{env_value},DEFAULT}}" + ) + else: + # Attempt to cast in a try to be able to catch the failed type casts with an exception + try: + typed_env = value_type(maybe_env) if maybe_env is not None else None + except Exception as e: + raise _SpockEnvResolverError( + f"Failed attempting to cast environment variable (name: {env_value}, value: `{maybe_env}`) " + f"into Spock specified type `{value_type.__name__}`" + ) + return typed_env diff --git a/spock/backend/typed.py b/spock/backend/typed.py index 5e8ac1c1..9f683e5f 100644 --- a/spock/backend/typed.py +++ b/spock/backend/typed.py @@ -389,7 +389,15 @@ def _type_katra(typed, default=None, optional=False): optional = True special_key = name typed = str - if default is not None: + if default is not None and optional: + # if a default is provided, that takes precedence + x = attr.ib( + validator=attr.validators.optional(attr.validators.instance_of(typed)), + default=default, + type=typed, + metadata={"optional": True, "base": name, "special_key": special_key}, + ) + elif default is not None: # if a default is provided, that takes precedence x = attr.ib( validator=attr.validators.instance_of(typed), diff --git a/spock/exceptions.py b/spock/exceptions.py index e060490b..19c1a50d 100644 --- a/spock/exceptions.py +++ b/spock/exceptions.py @@ -38,3 +38,9 @@ class _SpockValueError(Exception): """Custom exception for throwing value errors""" pass + + +class _SpockEnvResolverError(Exception): + """Custom exception for environment resolver""" + + pass \ No newline at end of file From 16477be299df9f37a268da091401e1cdabd60596 Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Wed, 11 May 2022 16:42:54 -0400 Subject: [PATCH 2/7] WIP: basic functionality for env resolver and crypto functionality. also added all currently installed packages to the info dump (in comments) such that a minimal python env should be able to be re-constructed --- NOTICE.txt | 2 + REQUIREMENTS.txt | 2 + spock/backend/builder.py | 20 ++- spock/backend/field_handlers.py | 135 +++++++++++++----- spock/backend/resolvers.py | 237 ++++++++++++++++++++++---------- spock/backend/saver.py | 12 +- spock/backend/spaces.py | 9 ++ spock/backend/typed.py | 43 +++++- spock/backend/utils.py | 19 +++ spock/backend/wrappers.py | 4 + spock/builder.py | 69 +++++++++- spock/exceptions.py | 10 +- spock/handlers.py | 82 +++++++++-- spock/utils.py | 23 ++++ 14 files changed, 538 insertions(+), 129 deletions(-) diff --git a/NOTICE.txt b/NOTICE.txt index 5b21e0b1..9f10e80c 100644 --- a/NOTICE.txt +++ b/NOTICE.txt @@ -11,9 +11,11 @@ FMR LLC (https://www.fidelity.com/). This product relies on the following works (and the dependencies thereof), installed separately: - attrs | https://github.com/python-attrs/attrs | MIT License +- cryptography | https://github.com/pyca/cryptography | Apache License 2.0 + BSD License - GitPython | https://github.com/gitpython-developers/GitPython | BSD 3-Clause License - pytomlpp | https://github.com/bobfang1992/pytomlpp | MIT License - PyYAML | https://github.com/yaml/pyyaml | MIT License +- setuptools | https://github.com/pypa/setuptools | MIT License Optional extensions rely on the following works (and the dependencies thereof), installed separately: diff --git a/REQUIREMENTS.txt b/REQUIREMENTS.txt index b037bfae..9b67642a 100644 --- a/REQUIREMENTS.txt +++ b/REQUIREMENTS.txt @@ -1,4 +1,6 @@ attrs~=21.4 +cryptography~=37.0 GitPython~=3.1 pytomlpp~=1.0 pyYAML~=5.4 +setuptools~=59.6 diff --git a/spock/backend/builder.py b/spock/backend/builder.py index f6fe4d24..73ca8391 100644 --- a/spock/backend/builder.py +++ b/spock/backend/builder.py @@ -7,7 +7,7 @@ import argparse from abc import ABC, abstractmethod from enum import EnumMeta -from typing import Dict, List +from typing import ByteString, Dict, List import attr @@ -34,11 +34,20 @@ class BaseBuilder(ABC): # pylint: disable=too-few-public-methods _module_name: module name to register in the spock module space save_path: list of path(s) to save the configs to _lazy: attempts to lazily find @spock decorated classes registered within sys.modules["spock"].backend.config + _salt: salt use for crypto purposes + _key: key used for crypto purposes """ def __init__( - self, *args, max_indent: int = 4, module_name: str, lazy: bool, **kwargs + self, + *args, + max_indent: int = 4, + module_name: str, + lazy: bool, + salt: str, + key: ByteString, + **kwargs, ): """Init call for BaseBuilder @@ -46,10 +55,15 @@ def __init__( *args: iterable of @spock decorated classes max_indent: max indent for pretty print of help module_name: module name to register in the spock module space + lazy: lazily find @spock decorated classes + salt: cryptographic salt + key: cryptographic key **kwargs: keyword args """ self._input_classes = args self._lazy = lazy + self._salt = salt + self._key = key self._graph = Graph(input_classes=self.input_classes, lazy=self._lazy) # Make sure the input classes are updated -- lazy evaluation self._input_classes = self._graph.nodes @@ -144,7 +158,7 @@ def resolve_spock_space_kwargs(self, graph: Graph, dict_args: Dict) -> Dict: for spock_cls in graph.roots: # Initial call to the RegisterSpockCls generate function (which will handle recursing if needed) spock_instance, special_keys = RegisterSpockCls.recurse_generate( - spock_cls, builder_space + spock_cls, builder_space, self._salt, self._key ) builder_space.spock_space[spock_cls.__name__] = spock_instance diff --git a/spock/backend/field_handlers.py b/spock/backend/field_handlers.py index 23455856..481e0b75 100644 --- a/spock/backend/field_handlers.py +++ b/spock/backend/field_handlers.py @@ -9,16 +9,17 @@ import sys from abc import ABC, abstractmethod from enum import EnumMeta -from typing import Callable, Dict, List, Tuple, Type +from typing import Any, ByteString, Callable, Dict, List, Tuple, Type from attr import NOTHING, Attribute -from spock.backend.resolvers import parse_env_variables +from spock.backend.resolvers import CryptoResolver, EnvResolver from spock.backend.spaces import AttributeSpace, BuilderSpace, ConfigSpace from spock.backend.utils import ( _get_name_py_version, _recurse_callables, _str_2_callable, + encrypt_value, ) from spock.exceptions import _SpockInstantiationError, _SpockNotOptionalError from spock.utils import ( @@ -45,12 +46,16 @@ class RegisterFieldTemplate(ABC): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call for RegisterFieldTemplate class Args: """ self.special_keys = {} + self._salt = salt + self._key = key + self._env_resolver = EnvResolver() + self._crypto_resolver = CryptoResolver(self._salt, self._key) def __call__(self, attr_space: AttributeSpace, builder_space: BuilderSpace): """Call method for RegisterFieldTemplate @@ -93,8 +98,10 @@ def _is_attribute_in_config_arguments( _is_spock_instance(attr_space.attribute.type) and attr_space.attribute.default is not None ): - attr_space.field, special_keys = RegisterSpockCls().recurse_generate( - attr_space.attribute.type, builder_space + attr_space.field, special_keys = RegisterSpockCls( + self._salt, self._key + ).recurse_generate( + attr_space.attribute.type, builder_space, self._salt, self._key ) attr_space.attribute = attr_space.attribute.evolve(default=attr_space.field) builder_space.spock_space[ @@ -133,8 +140,46 @@ def handle_optional_attribute_value( Returns: """ - cleaned_attribute = parse_env_variables(attr_space.attribute.default, attr_space.attribute.type) - attr_space.field = cleaned_attribute + + value, env_annotation = self._env_resolver.resolve( + attr_space.attribute.default, attr_space.attribute.type + ) + if env_annotation is not None: + self._handle_env_annotations( + attr_space, env_annotation, value, attr_space.attribute.default + ) + value, crypto_annotation = self._crypto_resolver.resolve( + value, attr_space.attribute.type + ) + if crypto_annotation is not None: + self._handle_crypto_annotations( + attr_space, crypto_annotation, attr_space.attribute.default + ) + attr_space.field = value + + def _handle_env_annotations( + self, attr_space: AttributeSpace, annotation: str, value: Any, og_value: Any + ): + if annotation == "crypto": + # Take the current value to string and then encrypt + attr_space.annotations = ( + f"${{spock.crypto:{encrypt_value(str(value), self._key, self._salt)}}}" + ) + elif annotation == "inject": + attr_space.annotations = og_value + else: + raise _SpockInstantiationError(f"Got unknown env annotation `{annotation}`") + + @staticmethod + def _handle_crypto_annotations( + attr_space: AttributeSpace, annotation: str, og_value: str + ): + if annotation == "crypto": + attr_space.annotations = og_value + else: + raise _SpockInstantiationError( + f"Got unknown crypto annotation `{annotation}`" + ) @abstractmethod def handle_optional_attribute_type( @@ -157,12 +202,12 @@ class RegisterList(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterList Args: """ - super(RegisterList, self).__init__() + super(RegisterList, self).__init__(salt, key) def handle_attribute_from_config( self, attr_space: AttributeSpace, builder_space: BuilderSpace @@ -255,12 +300,12 @@ class RegisterEnum(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterEnum Args: """ - super(RegisterEnum, self).__init__() + super(RegisterEnum, self).__init__(salt, key) def handle_attribute_from_config( self, attr_space: AttributeSpace, builder_space: BuilderSpace @@ -327,7 +372,7 @@ def _handle_and_register_enum( Returns: """ attr_space.field, special_keys = RegisterSpockCls.recurse_generate( - enum_cls, builder_space + enum_cls, builder_space, self._salt, self._key ) self.special_keys.update(special_keys) builder_space.spock_space[enum_cls.__name__] = attr_space.field @@ -341,12 +386,12 @@ class RegisterCallableField(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterSimpleField Args: """ - super(RegisterCallableField, self).__init__() + super(RegisterCallableField, self).__init__(salt, key) def handle_attribute_from_config( self, attr_space: AttributeSpace, builder_space: BuilderSpace @@ -394,12 +439,12 @@ class RegisterGenericAliasCallableField(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterSimpleField Args: """ - super(RegisterGenericAliasCallableField, self).__init__() + super(RegisterGenericAliasCallableField, self).__init__(salt, key) def handle_attribute_from_config( self, attr_space: AttributeSpace, builder_space: BuilderSpace @@ -461,17 +506,17 @@ class RegisterSimpleField(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterSimpleField Args: """ - super(RegisterSimpleField, self).__init__() + super(RegisterSimpleField, self).__init__(salt, key) def handle_attribute_from_config( self, attr_space: AttributeSpace, builder_space: BuilderSpace ): - """Handles setting a simple attribute when it is a spock class type + """Handles setting a simple attribute from a config file Args: attr_space: holds information about a single attribute that is mapped to a ConfigSpace @@ -479,9 +524,20 @@ def handle_attribute_from_config( Returns: """ - attr_space.field = builder_space.arguments[attr_space.config_space.name][ + og_value = builder_space.arguments[attr_space.config_space.name][ attr_space.attribute.name ] + value, env_annotation = self._env_resolver.resolve( + og_value, attr_space.attribute.type + ) + if env_annotation is not None: + self._handle_env_annotations(attr_space, env_annotation, value, og_value) + value, crypto_annotation = self._crypto_resolver.resolve( + value, attr_space.attribute.type + ) + if crypto_annotation is not None: + self._handle_crypto_annotations(attr_space, crypto_annotation, og_value) + attr_space.field = value self.register_special_key(attr_space) def handle_optional_attribute_type( @@ -543,12 +599,12 @@ class RegisterTuneCls(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterTuneCls Args: """ - super(RegisterTuneCls, self).__init__() + super(RegisterTuneCls, self).__init__(salt, key) @staticmethod def _attr_type(attr_space: AttributeSpace): @@ -622,12 +678,12 @@ class RegisterSpockCls(RegisterFieldTemplate): """ - def __init__(self): + def __init__(self, salt: str, key: ByteString): """Init call to RegisterSpockCls Args: """ - super(RegisterSpockCls, self).__init__() + super(RegisterSpockCls, self).__init__(salt, key) @staticmethod def _attr_type(attr_space: AttributeSpace): @@ -656,7 +712,9 @@ def handle_attribute_from_config( Returns: """ attr_type = self._attr_type(attr_space) - attr_space.field, special_keys = self.recurse_generate(attr_type, builder_space) + attr_space.field, special_keys = self.recurse_generate( + attr_type, builder_space, self._salt, self._key + ) builder_space.spock_space[attr_type.__name__] = attr_space.field self.special_keys.update(special_keys) @@ -694,7 +752,7 @@ def handle_optional_attribute_type( Returns: """ attr_space.field, special_keys = RegisterSpockCls.recurse_generate( - self._attr_type(attr_space), builder_space + self._attr_type(attr_space), builder_space, self._salt, self._key ) self.special_keys.update(special_keys) @@ -738,7 +796,9 @@ def _find_callables(cls, typed: _T): return out @classmethod - def recurse_generate(cls, spock_cls: _C, builder_space: BuilderSpace): + def recurse_generate( + cls, spock_cls: _C, builder_space: BuilderSpace, salt: str, key: ByteString + ): """Call on a spock classes to iterate through the attrs attributes and handle each based on type and optionality Triggers a recursive call when an attribute refers to another spock classes @@ -754,6 +814,7 @@ def recurse_generate(cls, spock_cls: _C, builder_space: BuilderSpace): # Empty dits for storing info special_keys = {} fields = {} + annotations = {} # Init the ConfigSpace for this spock class config_space = ConfigSpace(spock_cls, fields) # Iterate through the attrs within the spock class @@ -764,7 +825,7 @@ def recurse_generate(cls, spock_cls: _C, builder_space: BuilderSpace): if ( (attribute.type is list) or (attribute.type is List) ) and _is_spock_instance(attribute.metadata["type"].__args__[0]): - handler = RegisterList() + handler = RegisterList(salt, key) # Dict/List of Callables elif ( (attribute.type is list) @@ -775,31 +836,37 @@ def recurse_generate(cls, spock_cls: _C, builder_space: BuilderSpace): or (attribute.type is Tuple) ) and cls._find_callables(attribute.metadata["type"]): # handler = RegisterListCallableField() - handler = RegisterGenericAliasCallableField() + handler = RegisterGenericAliasCallableField(salt, key) # Enums elif isinstance(attribute.type, EnumMeta) and _check_iterable( attribute.type ): - handler = RegisterEnum() + handler = RegisterEnum(salt, key) # References to other spock classes elif _is_spock_instance(attribute.type): - handler = RegisterSpockCls() + handler = RegisterSpockCls(salt, key) # References to tuner classes elif _is_spock_tune_instance(attribute.type): - handler = RegisterTuneCls() + handler = RegisterTuneCls(salt, key) # References to callables elif isinstance(attribute.type, _SpockVariadicGenericAlias): - handler = RegisterCallableField() + handler = RegisterCallableField(salt, key) # Basic field else: - handler = RegisterSimpleField() + handler = RegisterSimpleField(salt, key) handler(attr_space, builder_space) special_keys.update(handler.special_keys) + # Handle annotations by attaching them to a dictionary + if attr_space.annotations is not None: + annotations.update({attr_space.attribute.name: attr_space.annotations}) # Try except on the class since it might not be successful -- throw the attrs message as it will know the # error on instantiation try: + # If there are annotations attach them to the spock class in the __resolver__ attribute + if len(annotations) > 0: + spock_cls.__resolver__ = annotations spock_instance = spock_cls(**fields) except Exception as e: raise _SpockInstantiationError( diff --git a/spock/backend/resolvers.py b/spock/backend/resolvers.py index da536eef..ee5a367e 100644 --- a/spock/backend/resolvers.py +++ b/spock/backend/resolvers.py @@ -4,87 +4,184 @@ # SPDX-License-Identifier: Apache-2.0 """Resolver functions for Spock""" - -import re import os -from typing import Any +import re +from abc import ABC, abstractmethod +from typing import Any, ByteString, Optional, Tuple -from spock.exceptions import _SpockEnvResolverError +from spock.backend.utils import decrypt_value +from spock.exceptions import _SpockResolverError from spock.utils import _T -# ENV Resolver -- full regex is ^\${spock\.env:.*}$ -CLIP_ENV_PATTERN = r"^\${spock\.env:" -CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) -END_ENV_PATTERN = r"}$" -END_REGEX_OP = re.compile(END_ENV_PATTERN) -FULL_ENV_PATTERN = CLIP_ENV_PATTERN + r".*" + END_ENV_PATTERN -FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) - - -def parse_env_variables(value: Any, value_type: _T) -> Any: - # Check if it matches the regex - # if so then split and replace with the value from the env -- have to check here if the env variable actually - # exists -- if it doesn't then we need to raise an exception -- allow for None? - - # If it's a string we can check the regex - if isinstance(value, str): - # Check the regex - match_obj = FULL_REGEX_OP.fullmatch(value) - # if the object exists then we've matched a pattern and need to handle it - if match_obj is not None: - return _get_env_value(value, value_type) - # Regex doesn't match so just passthrough +class BaseResolver(ABC): + def __init__(self): + self._annotation_set = {"crypto", "inject"} + + @abstractmethod + def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: + pass + + @staticmethod + def _handle_default(value: str): + env_value, default_value = value.split(",") + default_value = default_value.strip() + # Swap string None to type None + if default_value == "None": + default_value = None + return env_value, default_value + + @staticmethod + def _check_base_regex( + full_regex_op: re.Pattern, + value: Any, + ) -> bool: + # If it's a string we can check the regex + if isinstance(value, str): + # Check the regex and return non None status + return full_regex_op.fullmatch(value) is not None + # If it's not a string we can't resolve anything so just passthrough and let spock handle the value else: - return value - # If it's not a string we can't resolve anything so just passthrough and let spock handle the value - else: - return value - - -def _handle_default(value: str): - env_value, default_value = value.split(',') - default_value = default_value.strip() - # Swap string None to type None - if default_value == "None": - default_value = None - return env_value, default_value - - -def _get_env_value(value: str, value_type: _T): - # Based on the start and end regex ops find the value the user set - env_str = END_REGEX_OP.split(CLIP_REGEX_OP.split(value)[1])[0] - # Attempt to split on a comma for a default value - split_len = len(env_str.split(',')) - # Default found if the len is 2 - if split_len == 2: - env_value, default_value = _handle_default(env_str) - # If the length is larger than two then the syntax is messed up - elif split_len > 2: - raise _SpockEnvResolverError( - f"Issue with environment variable syntax -- currently `{value}` has more than one `,` which means the " - f"optional default value cannot be resolved -- please use only one `,` separator within the syntax" - ) - else: - env_value = env_str - default_value = "None" - # Attempt to get the env variable - if default_value == "None": - maybe_env = os.getenv(env_value) - else: - maybe_env = os.getenv(env_value, default_value) - if maybe_env is None and default_value == "None": - raise _SpockEnvResolverError( - f"Attempted to get `{env_value}` from environment variables but it is not set -- please set this " - f"variable or provide a default via the following syntax ${{spock.env:{env_value},DEFAULT}}" - ) - else: + return False + + @staticmethod + def _attempt_cast(maybe_env: Optional[str], value_type: _T, env_value: str): # Attempt to cast in a try to be able to catch the failed type casts with an exception try: typed_env = value_type(maybe_env) if maybe_env is not None else None except Exception as e: - raise _SpockEnvResolverError( + raise _SpockResolverError( f"Failed attempting to cast environment variable (name: {env_value}, value: `{maybe_env}`) " f"into Spock specified type `{value_type.__name__}`" ) return typed_env + + def _apply_regex( + self, + end_regex_op: re.Pattern, + clip_regex_op: re.Pattern, + value: str, + allow_default: bool, + allow_annotation: bool, + ): + # Based on the start and end regex ops find the value the user set + env_str = end_regex_op.split(clip_regex_op.split(value)[-1])[0] + if allow_annotation and len(clip_regex_op.split(value)) > 2 and clip_regex_op.split(value)[1] != "": + annotation = clip_regex_op.split(value)[1] + if annotation not in self._annotation_set: + raise _SpockResolverError( + f"Environment variable annotation must be within {self._annotation_set} -- got `{annotation}`" + ) + elif not allow_annotation and len(clip_regex_op.split(value)) > 2: + raise _SpockResolverError( + f"Found annotation style format however `{value}` does not support annotations" + ) + else: + annotation = None + # Attempt to split on a comma for a default value + split_len = len(env_str.split(",")) + # Default found if the len is 2 + if split_len == 2 and allow_default: + env_value, default_value = self._handle_default(env_str) + # If the length is larger than two then the syntax is messed up + elif split_len > 2 and allow_default: + raise _SpockResolverError( + f"Issue with environment variable syntax -- currently `{value}` has more than one `,` which means the " + f"optional default value cannot be resolved -- please use only one `,` separator within the syntax" + ) + elif split_len > 1 and not allow_default: + raise _SpockResolverError( + f"Syntax does not support default values -- currently `{value}` contains the separator `,` which " + f"id used to indicate default values" + ) + else: + env_value = env_str + default_value = "None" + return env_value, default_value, annotation + + +class EnvResolver(BaseResolver): + # ENV Resolver -- full regex is ^\${spock\.env:.*}$ + CLIP_ENV_PATTERN = r"^\${spock\.env\.?([a-z]*?):" + CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) + END_ENV_PATTERN = r"}$" + END_REGEX_OP = re.compile(END_ENV_PATTERN) + FULL_ENV_PATTERN = CLIP_ENV_PATTERN + r".*" + END_ENV_PATTERN + FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) + + def __init__(self): + super(EnvResolver, self).__init__() + + def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: + # Check the full regex for a match + regex_match = self._check_base_regex(self.FULL_REGEX_OP, value) + # if there is a regex match it needs to be handled by the underlying resolver ops + if regex_match: + # Apply the regex + env_value, default_value, annotation = self._apply_regex( + self.END_REGEX_OP, + self.CLIP_REGEX_OP, + value, + allow_default=True, + allow_annotation=True, + ) + # Get the value from the env + maybe_env = self._get_from_env(default_value, env_value) + # Attempt to cast the value to its underlying type + typed_env = self._attempt_cast(maybe_env, value_type, env_value) + # Else just pass through + else: + typed_env = value + annotation = None + return typed_env, annotation + + @staticmethod + def _get_from_env(default_value: Optional[str], env_value: str): + # Attempt to get the env variable + if default_value == "None": + maybe_env = os.getenv(env_value) + else: + maybe_env = os.getenv(env_value, default_value) + if maybe_env is None and default_value == "None": + raise _SpockResolverError( + f"Attempted to get `{env_value}` from environment variables but it is not set -- please set this " + f"variable or provide a default via the following syntax ${{spock.env:{env_value},DEFAULT}}" + ) + return maybe_env + + +class CryptoResolver(BaseResolver): + # ENV Resolver -- full regex is ^\${spock\.crypto:.*}$ + CLIP_ENV_PATTERN = r"^\${spock\.crypto:" + CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) + END_ENV_PATTERN = r"}$" + END_REGEX_OP = re.compile(END_ENV_PATTERN) + FULL_ENV_PATTERN = CLIP_ENV_PATTERN + r".*" + END_ENV_PATTERN + FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) + + def __init__(self, salt: str, key: ByteString): + super(CryptoResolver, self).__init__() + self._salt = salt + self._key = key + + def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: + regex_match = self._check_base_regex(self.FULL_REGEX_OP, value) + if regex_match: + crypto_value, default_value, annotation = self._apply_regex( + self.END_REGEX_OP, + self.CLIP_REGEX_OP, + value, + allow_default=False, + allow_annotation=False, + ) + decrypted_value = decrypt_value(crypto_value, self._key, self._salt) + typed_decrypted = self._attempt_cast( + decrypted_value, value_type, crypto_value + ) + annotation = "crypto" + # Pass through + else: + typed_decrypted = value + annotation = None + # Crypto in --> crypto out annotation wise or else this exposes the value in plaintext + return typed_decrypted, annotation diff --git a/spock/backend/saver.py b/spock/backend/saver.py index 34cfa985..df18ab81 100644 --- a/spock/backend/saver.py +++ b/spock/backend/saver.py @@ -13,7 +13,7 @@ from spock.backend.handler import BaseHandler from spock.backend.utils import _callable_2_str, _get_iter, _recurse_callables from spock.backend.wrappers import Spockspace -from spock.utils import add_info +from spock.utils import add_info, get_packages class BaseSaver(BaseHandler): # pylint: disable=too-few-public-methods @@ -76,12 +76,18 @@ def save( """ # Check extension self._check_extension(file_extension=file_extension) - # Make the filename -- always append a uuid for unique-ness + # Make the filename -- always append uuid for unique-ness uuid_str = str(uuid4()) if fixed_uuid is None else fixed_uuid fname = "" if file_name is None else f"{file_name}." name = f"{fname}{uuid_str}.spock.cfg{file_extension}" # Fix up values -- parameters out_dict = self._clean_up_values(payload) + # Handle any env annotations that are present + # Just stuff them into the dictionary + for k, v in payload: + if hasattr(v, "__resolver__"): + for key, val in v.__resolver__.items(): + out_dict[k][key] = val # Fix up the tuner values if present tuner_dict = ( self._clean_tuner_values(tuner_payload) @@ -92,10 +98,12 @@ def save( out_dict.update(tuner_dict) # Get extra info extra_dict = add_info() if extra_info else None + library_dict = get_packages() if extra_info else None try: self._supported_extensions.get(file_extension)().save( out_dict=out_dict, info_dict=extra_dict, + library_dict=library_dict, path=path, name=name, create_path=create_save_path, diff --git a/spock/backend/spaces.py b/spock/backend/spaces.py index 1eccb38b..e6fed8a6 100644 --- a/spock/backend/spaces.py +++ b/spock/backend/spaces.py @@ -55,6 +55,15 @@ def __init__(self, attribute: Type[Attribute], config_space: ConfigSpace): """ self.config_space = config_space self.attribute = attribute + self._annotations = None + + @property + def annotations(self): + return self._annotations + + @annotations.setter + def annotations(self, x): + self._annotations = x @property def field(self): diff --git a/spock/backend/typed.py b/spock/backend/typed.py index 9f683e5f..816873c4 100644 --- a/spock/backend/typed.py +++ b/spock/backend/typed.py @@ -138,7 +138,19 @@ def _generic_alias_katra(typed, default=None, optional=False): """ # base python class from which a GenericAlias is derived base_typed = typed.__origin__ - if default is not None: + if default is not None and optional: + # if there's no default, but marked as optional, then set the default to None + x = attr.ib( + validator=attr.validators.optional(_recursive_generic_validator(typed)), + type=base_typed, + default=default, + metadata={ + "optional": True, + "base": _extract_base_type(typed), + "type": typed, + }, + ) + elif default is not None: x = attr.ib( validator=_recursive_generic_validator(typed), default=default, @@ -261,7 +273,16 @@ def _enum_base_katra(typed, base_type, allowed, default=None, optional=False): x: Attribute from attrs """ - if default is not None: + if default is not None and optional: + x = attr.ib( + validator=attr.validators.optional( + [attr.validators.instance_of(base_type), attr.validators.in_(allowed)] + ), + default=_cast_enum_default(default), + type=typed, + metadata={"base": typed.__name__, "optional": True}, + ) + elif default is not None: x = attr.ib( validator=[ attr.validators.instance_of(base_type), @@ -330,7 +351,14 @@ def _enum_class_katra(typed, allowed, default=None, optional=False): x: Attribute from attrs """ - if default is not None: + if default is not None and optional: + x = attr.ib( + validator=attr.validators.optional([partial(_in_type, options=allowed)]), + default=_cast_enum_default(default), + type=typed, + metadata={"base": typed.__name__, "optional": True}, + ) + elif default is not None: x = attr.ib( validator=[partial(_in_type, options=allowed)], default=_cast_enum_default(default), @@ -470,7 +498,14 @@ def _callable_katra(typed, default=None, optional=False): x: Attribute from attrs """ - if default is not None: + if default is not None and optional: + x = attr.ib( + validator=attr.validators.optional(attr.validators.is_callable()), + default=default, + type=typed, + metadata={"optional": True, "base": _get_name_py_version(typed)}, + ) + elif default is not None: # if a default is provided, that takes precedence x = attr.ib( validator=attr.validators.is_callable(), diff --git a/spock/backend/utils.py b/spock/backend/utils.py index 5296ff43..b9f458e2 100644 --- a/spock/backend/utils.py +++ b/spock/backend/utils.py @@ -8,10 +8,29 @@ import importlib from typing import Any, Callable, Dict, List, Tuple, Type, Union +from cryptography.fernet import Fernet + from spock.exceptions import _SpockValueError from spock.utils import _C, _T, _SpockVariadicGenericAlias +def encrypt_value(value, key, salt): + # Make the class to encrypt + encrypt = Fernet(key=key) + # Encrypt the plaintext value + salted_password = value + salt + # encode to utf-8 -> encrypt -> decode from utf-8 + return encrypt.encrypt(str.encode(salted_password)).decode() + + +def decrypt_value(value, key, salt): + # Make the class to encrypt + decrypt = Fernet(key=key) + # Decrypt back to plaintext value + salted_password = decrypt.decrypt(str.encode(value)).decode() + return salted_password[: -len(salt)] + + def _str_2_callable(val: str, **kwargs): """Tries to convert a string representation of a module and callable to the reference to the callable diff --git a/spock/backend/wrappers.py b/spock/backend/wrappers.py index ffe2b9c6..55ea77bd 100644 --- a/spock/backend/wrappers.py +++ b/spock/backend/wrappers.py @@ -24,3 +24,7 @@ def __repr__(self): # Remove aliases in YAML print yaml.Dumper.ignore_aliases = lambda *args: True return yaml.dump(self.__dict__, default_flow_style=False) + + def __iter__(self): + for k, v in self.__dict__.items(): + yield k, v diff --git a/spock/builder.py b/spock/builder.py index 0eba4939..74d04e65 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -10,22 +10,26 @@ from collections import Counter from copy import deepcopy from pathlib import Path -from typing import Dict, List, Optional, Tuple, Type, Union +from typing import ByteString, Dict, List, Optional, Tuple, Type, Union from uuid import uuid4 import attr +from cryptography.fernet import Fernet from spock.backend.builder import AttrBuilder from spock.backend.payload import AttrPayload +from spock.backend.resolvers import EnvResolver from spock.backend.saver import AttrSaver from spock.backend.wrappers import Spockspace -from spock.exceptions import _SpockEvolveError, _SpockValueError +from spock.exceptions import _SpockCryptoError, _SpockEvolveError, _SpockValueError +from spock.handlers import YAMLHandler from spock.utils import ( _C, _T, _is_spock_instance, check_payload_overwrite, deep_payload_update, + make_salt, ) @@ -56,6 +60,8 @@ class ConfigArgBuilder: thus alleviating the need to pass all @spock decorated classes to *args _no_cmd_line: turn off cmd line args _desc: description for help + _salt: salt use for crypto purposes + _key: key used for crypto purposes """ @@ -67,6 +73,7 @@ def __init__( lazy: bool = False, no_cmd_line: bool = False, s3_config: Optional[_T] = None, + crypto: Optional[Union[str, Tuple[str, ByteString]]] = None, **kwargs, ): """Init call for ConfigArgBuilder @@ -76,10 +83,11 @@ def __init__( configs: list of config paths desc: description for help lazy: attempts to lazily find @spock decorated classes registered within sys.modules["spock"].backend.config - as well as the parents of any lazily inherited @spock class - thus alleviating the need to pass all @spock decorated classes to *args + as well as the parents of any lazily inherited @spock class thus alleviating the need to pass all + @spock decorated classes to *args no_cmd_line: turn off cmd line args s3_config: s3Config object for S3 support + crypto: either a path to a prior spock saved crypto.yaml file or a tuple of a cryptographic salt and key **kwargs: keyword args """ @@ -89,13 +97,16 @@ def __init__( self._lazy = lazy self._no_cmd_line = no_cmd_line self._desc = desc + self._salt, self._key = self._maybe_crypto(crypto, s3_config) # Build the payload and saver objects self._payload_obj = AttrPayload(s3_config=s3_config) self._saver_obj = AttrSaver(s3_config=s3_config) # Split the fixed parameters from the tuneable ones (if present) fixed_args, tune_args = self._strip_tune_parameters(args) # The fixed parameter builder - self._builder_obj = AttrBuilder(*fixed_args, lazy=lazy, **kwargs) + self._builder_obj = AttrBuilder( + *fixed_args, lazy=lazy, salt=self._salt, key=self._key, **kwargs + ) # The possible tunable parameter builder -- might return None self._tune_obj, self._tune_payload_obj = self._handle_tuner_objects( tune_args, s3_config, kwargs @@ -163,6 +174,14 @@ def best(self) -> Spockspace: """Returns a Spockspace of the best hyper-parameter config and the associated metric value""" return self._tuner_interface.best + @property + def salt(self): + return self._salt + + @property + def key(self): + return self._key + def sample(self) -> Spockspace: """Sample method that constructs a namespace from the fixed parameters and samples from the tuner space to generate a Spockspace derived from both @@ -772,3 +791,43 @@ def _set_matching_attrs_by_name( f"Evolved: Parent = {parent_cls_name}, Child = {current_cls_name}, Value = {v}" ) return new_arg_namespace + + @staticmethod + def _maybe_crypto( + crypto: Optional[Union[str, Tuple[str, ByteString]]], + s3_config: Optional[_T] = None, + salt_len: int = 16, + ): + """Handles setting up the underlying cryptography needs + + Args: + crypto: either a path to a prior spock saved crypto.yaml file or a tuple of a cryptographic salt and key + s3_config: s3Config object for S3 support + salt_len: length of the salt to create + + Returns: + tuple containing a salt and a key that spock can use to hide parameters + + """ + if isinstance(crypto, str): + # Read from the yaml and then split + payload = YAMLHandler().load(Path(crypto), s3_config) + salt = payload["salt"] + key = payload["key"] + elif isinstance(crypto, (tuple, Tuple)): + # order is salt, key + salt, key = crypto + elif crypto is None: + salt = make_salt(salt_len) + key = Fernet.generate_key() + else: + raise _SpockCryptoError( + f"The crypto argument expects a path to a valid prior spock saved crypto.yaml file, a tuple of a " + f"cryptographic salt and key (as raw strings or using env resolver syntax), or None -- got " + f"type {type(crypto)} and value `{crypto}`" + ) + # Parse if defined by env vars + env_resolver = EnvResolver() + salt, _ = env_resolver.resolve(salt, str) + key, _ = env_resolver.resolve(key, str) + return salt, key diff --git a/spock/exceptions.py b/spock/exceptions.py index 19c1a50d..862b6d73 100644 --- a/spock/exceptions.py +++ b/spock/exceptions.py @@ -40,7 +40,13 @@ class _SpockValueError(Exception): pass -class _SpockEnvResolverError(Exception): +class _SpockResolverError(Exception): """Custom exception for environment resolver""" - pass \ No newline at end of file + pass + + +class _SpockCryptoError(Exception): + """Custom exception for dealing with the crypto side of things""" + + pass diff --git a/spock/handlers.py b/spock/handlers.py index 95a690a8..0d2ed8c1 100644 --- a/spock/handlers.py +++ b/spock/handlers.py @@ -72,6 +72,7 @@ def save( self, out_dict: Dict, info_dict: Optional[Dict], + library_dict: Optional[Dict], path: Path, name: str, create_path: bool = False, @@ -85,6 +86,7 @@ def save( Args: out_dict: payload to write info_dict: info payload to write + library_dict: package info to write path: path to write out name: spock generated file name create_path: boolean to create the path if non-existent (for non S3) @@ -95,7 +97,12 @@ def save( write_path, is_s3 = self._handle_possible_s3_save_path( path=path, name=name, create_path=create_path, s3_config=s3_config ) - write_path = self._save(out_dict=out_dict, info_dict=info_dict, path=write_path) + write_path = self._save( + out_dict=out_dict, + info_dict=info_dict, + library_dict=library_dict, + path=write_path, + ) # After write check if it needs to be pushed to S3 if is_s3: try: @@ -111,12 +118,19 @@ def save( print("Error importing spock s3 utils after detecting s3:// save path") @abstractmethod - def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: + def _save( + self, + out_dict: Dict, + info_dict: Optional[Dict], + library_dict: Optional[Dict], + path: str, + ) -> str: """Write function for file type Args: out_dict: payload to write info_dict: info payload to write + library_dict: package info to write path: path to write out Returns: @@ -182,19 +196,35 @@ def _handle_possible_s3_save_path( return write_path, is_s3 @staticmethod - def write_extra_info(path, info_dict): + def write_extra_info( + path: str, + info_dict: Dict, + version: bool = True, + write_mode: str = "w+", + newlines: Optional[int] = None, + header: Optional[str] = None, + ): """Writes extra info to commented newlines Args: path: path to write out info_dict: info payload to write + version: write the spock version string first + write_mode: write mode for the file + newlines: number of new lines to add to start Returns: """ # Write the commented info as new lines - with open(path, "w+") as fid: + with open(path, write_mode) as fid: + if newlines is not None: + for _ in range(newlines): + fid.write("\n") + if header is not None: + fid.write(header) # Write a spock header - fid.write(f"# Spock Version: {__version__}\n") + if version: + fid.write(f"# Spock Version: {__version__}\n") # Write info dict if not None if info_dict is not None: for k, v in info_dict.items(): @@ -242,12 +272,19 @@ def _load(self, path: str) -> Dict: base_payload = yaml.safe_load(file_contents) return base_payload - def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: + def _save( + self, + out_dict: Dict, + info_dict: Optional[Dict], + library_dict: Optional[Dict], + path: str, + ) -> str: """Write function for YAML type Args: out_dict: payload to write info_dict: info payload to write + library_dict: package info to write path: path to write out Returns: @@ -258,6 +295,15 @@ def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: yaml.Dumper.ignore_aliases = lambda *args: True with open(path, "a") as yaml_fid: yaml.safe_dump(out_dict, yaml_fid, default_flow_style=False) + # Write the library info at the bottom + self.write_extra_info( + path=path, + info_dict=library_dict, + version=False, + write_mode="a", + newlines=2, + header="################\n# Package Info #\n################\n", + ) return path @@ -281,12 +327,19 @@ def _load(self, path: str) -> Dict: base_payload = pytomlpp.load(path) return base_payload - def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: + def _save( + self, + out_dict: Dict, + info_dict: Optional[Dict], + library_dict: Optional[Dict], + path: str, + ) -> str: """Write function for TOML type Args: out_dict: payload to write info_dict: info payload to write + library_dict: package info to write path: path to write out Returns: @@ -295,6 +348,10 @@ def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: self.write_extra_info(path=path, info_dict=info_dict) with open(path, "a") as toml_fid: pytomlpp.dump(out_dict, toml_fid) + # Write the library info at the bottom + self.write_extra_info( + path=path, info_dict=library_dict, version=False, write_mode="a", newlines=2 + ) return path @@ -319,17 +376,24 @@ def _load(self, path: str) -> Dict: base_payload = json.load(json_fid) return base_payload - def _save(self, out_dict: Dict, info_dict: Optional[Dict], path: str) -> str: + def _save( + self, + out_dict: Dict, + info_dict: Optional[Dict], + library_dict: Optional[Dict], + path: str, + ) -> str: """Write function for JSON type Args: out_dict: payload to write info_dict: info payload to write + library_dict: package info to write path: path to write out Returns: """ - if info_dict is not None: + if (info_dict is not None) or (library_dict is not None): warn( "JSON does not support comments and thus cannot save extra info to file... removing extra info" ) diff --git a/spock/utils.py b/spock/utils.py index f61e96d5..8cad895b 100644 --- a/spock/utils.py +++ b/spock/utils.py @@ -7,6 +7,7 @@ import ast import os +import random import socket import subprocess import sys @@ -20,12 +21,18 @@ import attr import git +import pkg_resources from spock.exceptions import _SpockValueError minor = sys.version_info.minor +def make_salt(salt_len: int = 16): + alphabet = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + return "".join(random.choice(alphabet) for i in range(salt_len)) + + def _get_alias_type(): if minor < 7: from typing import GenericMeta as _GenericAlias @@ -480,6 +487,22 @@ def _handle_generic_type_args(val: str) -> Any: return ast.literal_eval(val) +def get_packages() -> Dict: + """Gets all currently installed packages and assembles a dictionary of name: version + + Notes: + https://stackoverflow.com/a/50013400 + + Returns: + dictionary of all currently available packages + + """ + named_list = sorted([str(i.key) for i in pkg_resources.working_set]) + return { + f"# {i}": str(pkg_resources.working_set.by_key[i].version) for i in named_list + } + + def add_info() -> Dict: """Adds extra information to the output dictionary From 8bfab66046080a1e202ad746d87607a6fc9b3ebe Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Thu, 12 May 2022 16:58:49 -0400 Subject: [PATCH 3/7] WIP: working ability to save/load key and salt to/from file when necessary --- spock/backend/field_handlers.py | 7 ++ spock/backend/resolvers.py | 6 +- spock/backend/saver.py | 5 ++ spock/backend/spaces.py | 1 + spock/backend/wrappers.py | 6 +- spock/builder.py | 122 +++++++++++++++++++++++++------- spock/handlers.py | 103 ++++++++++++++++++++++----- spock/utils.py | 23 +++++- 8 files changed, 229 insertions(+), 44 deletions(-) diff --git a/spock/backend/field_handlers.py b/spock/backend/field_handlers.py index 481e0b75..583ca555 100644 --- a/spock/backend/field_handlers.py +++ b/spock/backend/field_handlers.py @@ -155,6 +155,7 @@ def handle_optional_attribute_value( self._handle_crypto_annotations( attr_space, crypto_annotation, attr_space.attribute.default ) + attr_space.crypto = True attr_space.field = value def _handle_env_annotations( @@ -537,6 +538,7 @@ def handle_attribute_from_config( ) if crypto_annotation is not None: self._handle_crypto_annotations(attr_space, crypto_annotation, og_value) + attr_space.crypto = True attr_space.field = value self.register_special_key(attr_space) @@ -815,6 +817,7 @@ def recurse_generate( special_keys = {} fields = {} annotations = {} + crypto = False # Init the ConfigSpace for this spock class config_space = ConfigSpace(spock_cls, fields) # Iterate through the attrs within the spock class @@ -860,6 +863,8 @@ def recurse_generate( # Handle annotations by attaching them to a dictionary if attr_space.annotations is not None: annotations.update({attr_space.attribute.name: attr_space.annotations}) + if attr_space.crypto: + crypto = True # Try except on the class since it might not be successful -- throw the attrs message as it will know the # error on instantiation @@ -867,6 +872,8 @@ def recurse_generate( # If there are annotations attach them to the spock class in the __resolver__ attribute if len(annotations) > 0: spock_cls.__resolver__ = annotations + if crypto: + spock_cls.__crypto__ = True spock_instance = spock_cls(**fields) except Exception as e: raise _SpockInstantiationError( diff --git a/spock/backend/resolvers.py b/spock/backend/resolvers.py index ee5a367e..0cca403e 100644 --- a/spock/backend/resolvers.py +++ b/spock/backend/resolvers.py @@ -66,7 +66,11 @@ def _apply_regex( ): # Based on the start and end regex ops find the value the user set env_str = end_regex_op.split(clip_regex_op.split(value)[-1])[0] - if allow_annotation and len(clip_regex_op.split(value)) > 2 and clip_regex_op.split(value)[1] != "": + if ( + allow_annotation + and len(clip_regex_op.split(value)) > 2 + and clip_regex_op.split(value)[1] != "" + ): annotation = clip_regex_op.split(value)[1] if annotation not in self._annotation_set: raise _SpockResolverError( diff --git a/spock/backend/saver.py b/spock/backend/saver.py index df18ab81..5a8ede1a 100644 --- a/spock/backend/saver.py +++ b/spock/backend/saver.py @@ -84,10 +84,13 @@ def save( out_dict = self._clean_up_values(payload) # Handle any env annotations that are present # Just stuff them into the dictionary + crypto_flag = False for k, v in payload: if hasattr(v, "__resolver__"): for key, val in v.__resolver__.items(): out_dict[k][key] = val + if hasattr(v, "__crypto__"): + crypto_flag = True # Fix up the tuner values if present tuner_dict = ( self._clean_tuner_values(tuner_payload) @@ -108,6 +111,8 @@ def save( name=name, create_path=create_save_path, s3_config=self._s3_config, + salt=payload.__salt__ if crypto_flag else None, + key=payload.__key__ if crypto_flag else None, ) except OSError as e: print(f"Unable to write to given path: {path / name}") diff --git a/spock/backend/spaces.py b/spock/backend/spaces.py index e6fed8a6..a80a415b 100644 --- a/spock/backend/spaces.py +++ b/spock/backend/spaces.py @@ -56,6 +56,7 @@ def __init__(self, attribute: Type[Attribute], config_space: ConfigSpace): self.config_space = config_space self.attribute = attribute self._annotations = None + self.crypto = False @property def annotations(self): diff --git a/spock/backend/wrappers.py b/spock/backend/wrappers.py index 55ea77bd..50a7d76a 100644 --- a/spock/backend/wrappers.py +++ b/spock/backend/wrappers.py @@ -20,10 +20,14 @@ class Spockspace(argparse.Namespace): def __init__(self, **kwargs): super(Spockspace, self).__init__(**kwargs) + @property + def __repr_dict__(self): + return {k: v for k, v in self.__dict__.items() if k not in {"__key__", "__salt__"}} + def __repr__(self): # Remove aliases in YAML print yaml.Dumper.ignore_aliases = lambda *args: True - return yaml.dump(self.__dict__, default_flow_style=False) + return yaml.dump(self.__repr_dict__, default_flow_style=False) def __iter__(self): for k, v in self.__dict__.items(): diff --git a/spock/builder.py b/spock/builder.py index 74d04e65..2427479c 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -6,6 +6,7 @@ """Handles the building/saving of the configurations from the Spock config classes""" import argparse +import os.path import sys from collections import Counter from copy import deepcopy @@ -73,6 +74,8 @@ def __init__( lazy: bool = False, no_cmd_line: bool = False, s3_config: Optional[_T] = None, + key: Optional[Union[str, ByteString]], + salt: Optional[str], crypto: Optional[Union[str, Tuple[str, ByteString]]] = None, **kwargs, ): @@ -87,7 +90,9 @@ def __init__( @spock decorated classes to *args no_cmd_line: turn off cmd line args s3_config: s3Config object for S3 support - crypto: either a path to a prior spock saved crypto.yaml file or a tuple of a cryptographic salt and key + salt: either a path to a prior spock saved salt.yaml file or a string of the salt (can be an env reference) + key: either a path to a prior spock saved key.yaml file, a ByteString of the key, or a str of the key + (can be an env reference) **kwargs: keyword args """ @@ -97,7 +102,7 @@ def __init__( self._lazy = lazy self._no_cmd_line = no_cmd_line self._desc = desc - self._salt, self._key = self._maybe_crypto(crypto, s3_config) + self._salt, self._key = self._maybe_crypto(key, salt, s3_config) # Build the payload and saver objects self._payload_obj = AttrPayload(s3_config=s3_config) self._saver_obj = AttrSaver(s3_config=s3_config) @@ -128,6 +133,9 @@ def __init__( # Build the Spockspace from the payload and the classes # Fixed configs self._arg_namespace = self._builder_obj.generate(self._dict_args) + # Attach the key and salt to the Spockspace + self._arg_namespace.__salt__ = self.salt + self._arg_namespace.__key__ = self.key # Get the payload from the config files -- hyper-parameters -- only if the obj is not None if self._tune_obj is not None: self._tune_args = self._get_payload( @@ -792,16 +800,19 @@ def _set_matching_attrs_by_name( ) return new_arg_namespace - @staticmethod def _maybe_crypto( - crypto: Optional[Union[str, Tuple[str, ByteString]]], + self, + key: Optional[Union[str, ByteString]], + salt: Optional[str], s3_config: Optional[_T] = None, salt_len: int = 16, ): """Handles setting up the underlying cryptography needs Args: - crypto: either a path to a prior spock saved crypto.yaml file or a tuple of a cryptographic salt and key + salt: either a path to a prior spock saved salt.yaml file or a string of the salt (can be an env reference) + key: either a path to a prior spock saved key.yaml file, a ByteString of the key, or a str of the key + (can be an env reference) s3_config: s3Config object for S3 support salt_len: length of the salt to create @@ -809,25 +820,88 @@ def _maybe_crypto( tuple containing a salt and a key that spock can use to hide parameters """ - if isinstance(crypto, str): - # Read from the yaml and then split - payload = YAMLHandler().load(Path(crypto), s3_config) - salt = payload["salt"] - key = payload["key"] - elif isinstance(crypto, (tuple, Tuple)): - # order is salt, key - salt, key = crypto - elif crypto is None: + env_resolver = EnvResolver() + salt = self._get_salt(salt, env_resolver, salt_len, s3_config) + key = self._get_key(key, env_resolver, s3_config) + return salt, key + + def _get_salt( + self, + salt: Optional[str], + env_resolver: EnvResolver, + salt_len: int, + s3_config: Optional[_T] = None, + ): + """ + + Args: + salt: either a path to a prior spock saved salt.yaml file or a string of the salt (can be an env reference) + env_resolver: EnvResolver class to handle env variable resolution if needed + salt_len: length of the salt to create + s3_config: s3Config object for S3 support + + Returns: + salt as a string + + """ + # Byte string is assumed to be a direct key + if salt is None: salt = make_salt(salt_len) + elif os.path.splitext(salt)[1] in {".yaml", ".YAML", ".yml", ".YML"}: + salt = self._handle_yaml_read(salt, access="salt", s3_config=s3_config) + else: + salt, _ = env_resolver.resolve(salt, str) + return salt + + def _get_key( + self, + key: Optional[Union[str, ByteString]], + env_resolver: EnvResolver, + s3_config: Optional[_T] = None, + ): + """ + + Args: + key: either a path to a prior spock saved key.yaml file, a ByteString of the key, or a str of the key + (can be an env reference) + env_resolver: EnvResolver class to handle env variable resolution if needed + s3_config: s3Config object for S3 support + + Returns: + key as ByteString + + """ + if key is None: key = Fernet.generate_key() + # Byte string is assumed to be a direct key + elif os.path.splitext(key)[1] in {".yaml", ".YAML", ".yml", ".YML"}: + key = self._handle_yaml_read(key, access="key", s3_config=s3_config, encode=True) else: - raise _SpockCryptoError( - f"The crypto argument expects a path to a valid prior spock saved crypto.yaml file, a tuple of a " - f"cryptographic salt and key (as raw strings or using env resolver syntax), or None -- got " - f"type {type(crypto)} and value `{crypto}`" - ) - # Parse if defined by env vars - env_resolver = EnvResolver() - salt, _ = env_resolver.resolve(salt, str) - key, _ = env_resolver.resolve(key, str) - return salt, key + # Byte string is assumed to be a direct key + # So only handle the str here + if isinstance(key, str): + key, _ = env_resolver.resolve(key, str) + key = str.encode(key) + return key + + @staticmethod + def _handle_yaml_read(value: str, access: str, s3_config: Optional[_T] = None, encode: bool = False): + """Reads in a salt/key yaml + + Args: + value: path to the key/salt yaml + access: which variable name to use from the yaml + s3_config: s3Config object for S3 support + + Returns: + + """ + # Read from the yaml and then split + try: + payload = YAMLHandler().load(Path(value), s3_config) + read_value = payload[access] + if encode: + read_value = str.encode(read_value) + return read_value + except Exception as e: + _SpockCryptoError(f"Attempted to read from path `{value}` but failed") diff --git a/spock/handlers.py b/spock/handlers.py index 0d2ed8c1..2a1d370d 100644 --- a/spock/handlers.py +++ b/spock/handlers.py @@ -10,14 +10,14 @@ import re from abc import ABC, abstractmethod from pathlib import Path, PurePosixPath -from typing import Dict, Optional, Tuple, Union +from typing import ByteString, Dict, Optional, Tuple, Union from warnings import warn import pytomlpp import yaml from spock._version import get_versions -from spock.utils import check_path_s3, path_object_to_s3path +from spock.utils import _T, check_path_s3, path_object_to_s3path __version__ = get_versions()["version"] @@ -52,7 +52,6 @@ def _post_process_config_paths(payload): """ if (payload is not None) and "config" in payload: payload["config"] = [Path(c) for c in payload["config"]] - return payload @abstractmethod @@ -68,6 +67,45 @@ def _load(self, path: str) -> Dict: """ raise NotImplementedError + def _write_crypto( + self, + value: Union[str, ByteString], + path: Path, + name: str, + crypto_name: str, + create_path: bool, + s3_config: Optional[_T], + ): + """Write values of the underlying cryptography data used to encode some spock values + + Args: + value: current crypto attribute + path: path to write out + name: spock generated file name + create_path: boolean to create the path if non-existent (for non S3) + s3_config: optional s3 config object if using s3 storage + crypto_name: name of the crypto attribute + + Returns: + None + + """ + # Convert ByteString to str + value = value.decode("utf-8") if isinstance(value, ByteString) else value + write_path, is_s3 = self._handle_possible_s3_save_path( + path=path, name=name, create_path=create_path, s3_config=s3_config + ) + # We need to shim in the crypto value name into the name used for S3 + name_root, name_extension = os.path.splitext(name) + name = f"{name_root}.{crypto_name}.yaml" + # Also need to shim the crypto name into the full path + root, extension = os.path.splitext(write_path) + full_name = f"{root}.{crypto_name}.yaml" + YAMLHandler.write({crypto_name: value}, full_name) + # After write check if it needs to be pushed to S3 + if is_s3: + self._check_s3_write(write_path, path, name, s3_config) + def save( self, out_dict: Dict, @@ -76,7 +114,9 @@ def save( path: Path, name: str, create_path: bool = False, - s3_config=None, + s3_config: Optional[_T] = None, + salt: Optional[str] = None, + key: Optional[ByteString] = None, ): """Write function for file type @@ -91,6 +131,8 @@ def save( name: spock generated file name create_path: boolean to create the path if non-existent (for non S3) s3_config: optional s3 config object if using s3 storage + salt: string of the salt used for crypto + key: ByteString of the key used for crypto Returns: """ @@ -105,17 +147,39 @@ def save( ) # After write check if it needs to be pushed to S3 if is_s3: - try: - from spock.addons.s3.utils import handle_s3_save_path + self._check_s3_write(write_path, path, name, s3_config) + # Write the crypto files if needed + if (salt is not None) and (key is not None): + # If the values are not none then write the salt and key into individual files + self._write_crypto(salt, path, name, "salt", create_path, s3_config) + self._write_crypto(key, path, name, "key", create_path, s3_config) - handle_s3_save_path( - temp_path=write_path, - s3_path=str(PurePosixPath(path)), - name=name, - s3_config=s3_config, - ) - except ImportError: - print("Error importing spock s3 utils after detecting s3:// save path") + @staticmethod + def _check_s3_write( + write_path: str, path: Path, name: str, s3_config: Optional[_T] + ): + """Handles writing to S3 if necessary + + Args: + write_path: path the file was written to locally + path: original path specified + name: original file name + s3_config: optional s3 config object if using s3 storage + + Returns: + + """ + try: + from spock.addons.s3.utils import handle_s3_save_path + + handle_s3_save_path( + temp_path=write_path, + s3_path=str(PurePosixPath(path)), + name=name, + s3_config=s3_config, + ) + except ImportError: + print("Error importing spock s3 utils after detecting s3:// save path") @abstractmethod def _save( @@ -286,15 +350,13 @@ def _save( info_dict: info payload to write library_dict: package info to write path: path to write out - Returns: """ # First write the commented info self.write_extra_info(path=path, info_dict=info_dict) # Remove aliases in YAML dump yaml.Dumper.ignore_aliases = lambda *args: True - with open(path, "a") as yaml_fid: - yaml.safe_dump(out_dict, yaml_fid, default_flow_style=False) + self.write(out_dict, path) # Write the library info at the bottom self.write_extra_info( path=path, @@ -306,6 +368,13 @@ def _save( ) return path + @staticmethod + def write(write_dict: Dict, path: str): + # Remove aliases in YAML dump + yaml.Dumper.ignore_aliases = lambda *args: True + with open(path, "a") as yaml_fid: + yaml.safe_dump(write_dict, yaml_fid, default_flow_style=False) + class TOMLHandler(Handler): """TOML class for loading TOML config files diff --git a/spock/utils.py b/spock/utils.py index 8cad895b..441d8744 100644 --- a/spock/utils.py +++ b/spock/utils.py @@ -29,11 +29,26 @@ def make_salt(salt_len: int = 16): + """Make a salt of specific length + + Args: + salt_len: length of the constructed salt + + Returns: + salt string + + """ alphabet = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" - return "".join(random.choice(alphabet) for i in range(salt_len)) + return "".join(random.choice(alphabet) for _ in range(salt_len)) def _get_alias_type(): + """Gets the correct type of GenericAlias for versions less than 3.6 + + Returns: + _GenericAlias type + + """ if minor < 7: from typing import GenericMeta as _GenericAlias else: @@ -43,6 +58,12 @@ def _get_alias_type(): def _get_callable_type(): + """Gets the correct underlying type reference for callable objects depending on the python version + + Returns: + _VariadicGenericAlias type + + """ if minor == 6: from typing import CallableMeta as _VariadicGenericAlias elif (minor > 6) and (minor < 9): From 1a3d8627deaa06e30b2e40ef2fe393816f88d13a Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Mon, 16 May 2022 11:12:09 -0400 Subject: [PATCH 4/7] WIP: docs for new resolver functionality --- spock/addons/tune/payload.py | 6 +- spock/backend/resolvers.py | 133 ++++++++++++++++++++++++++++++++--- spock/backend/saver.py | 30 ++++++-- spock/backend/wrappers.py | 3 + spock/builder.py | 13 ++-- spock/handlers.py | 6 +- tests/base/test_state.py | 4 ++ 7 files changed, 167 insertions(+), 28 deletions(-) diff --git a/spock/addons/tune/payload.py b/spock/addons/tune/payload.py index 7514996d..14fcb486 100644 --- a/spock/addons/tune/payload.py +++ b/spock/addons/tune/payload.py @@ -6,7 +6,9 @@ """Handles the tuner payload backend""" from spock.backend.payload import BasePayload -from spock.backend.utils import get_attr_fields +from spock.backend.utils import get_attr_fields, _T + +from typing import Optional class TunerPayload(BasePayload): @@ -20,7 +22,7 @@ class TunerPayload(BasePayload): """ - def __init__(self, s3_config=None): + def __init__(self, s3_config: Optional[_T] = None): """Init for TunerPayload Args: diff --git a/spock/backend/resolvers.py b/spock/backend/resolvers.py index 0cca403e..3d6d87c5 100644 --- a/spock/backend/resolvers.py +++ b/spock/backend/resolvers.py @@ -7,7 +7,7 @@ import os import re from abc import ABC, abstractmethod -from typing import Any, ByteString, Optional, Tuple +from typing import Any, ByteString, Optional, Tuple, Pattern, Union from spock.backend.utils import decrypt_value from spock.exceptions import _SpockResolverError @@ -15,15 +15,43 @@ class BaseResolver(ABC): + """Base class for resolvers + + Contains base methods for handling resolver syntax + + Attributes: + _annotation_set: current set of supported resolver annotations + + """ def __init__(self): + """Init for BaseResolver class""" self._annotation_set = {"crypto", "inject"} @abstractmethod def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: + """Resolves a variable from a given resolver syntax + + Args: + value: current value to attempt to resolve + value_type: type of the value to cast into + + Returns: + Tuple of correctly typed resolved variable and any annotations + + """ pass @staticmethod - def _handle_default(value: str): + def _handle_default(value: str) -> Tuple[str, Union[str, None]]: + """Handles setting defaults if allowed for a resolver + + Args: + value: current string value + + Returns: + tuple of given value and the default value + + """ env_value, default_value = value.split(",") default_value = default_value.strip() # Swap string None to type None @@ -33,9 +61,19 @@ def _handle_default(value: str): @staticmethod def _check_base_regex( - full_regex_op: re.Pattern, + full_regex_op: Pattern, value: Any, ) -> bool: + """Check if the value passed into the resolver matches the compiled regex op + + Args: + full_regex_op: the full compiled regex + value: the value passed into the resolver + + Returns: + boolean if there is a regex match + + """ # If it's a string we can check the regex if isinstance(value, str): # Check the regex and return non None status @@ -45,7 +83,21 @@ def _check_base_regex( return False @staticmethod - def _attempt_cast(maybe_env: Optional[str], value_type: _T, env_value: str): + def _attempt_cast(maybe_env: Optional[str], value_type: _T, env_value: str) -> Any: + """Attempts to cast the resolved variable into the given type + + Args: + maybe_env: possible resolved variable + value_type: type to cast into + env_value: the reference to the resolved variable + + Returns: + value type cast into the correct type + + Raises: + _SpockResolverError if it cannot be cast into the specified type + + """ # Attempt to cast in a try to be able to catch the failed type casts with an exception try: typed_env = value_type(maybe_env) if maybe_env is not None else None @@ -58,12 +110,29 @@ def _attempt_cast(maybe_env: Optional[str], value_type: _T, env_value: str): def _apply_regex( self, - end_regex_op: re.Pattern, - clip_regex_op: re.Pattern, + end_regex_op: Pattern, + clip_regex_op: Pattern, value: str, allow_default: bool, allow_annotation: bool, - ): + ) -> Tuple[str, str, Optional[str]]: + """Applies the front and back regexes to the string value, determines defaults and annotations + + Args: + end_regex_op: compiled regex for the back half of the match + clip_regex_op: compiled regex for the front half of the match + value: current string value to resolve + allow_default: if allowed to contain default value syntax + allow_annotation: if allowed to contain annotation syntax + + Returns: + tuple containing the resolved string reference, the default value, and the annotation string + + Raises: + _SpockResolverError if annotation isn't within the supported set, annotation is not supported, multiple `,` + values are used, or defaults are given yet not supported + + """ # Based on the start and end regex ops find the value the user set env_str = end_regex_op.split(clip_regex_op.split(value)[-1])[0] if ( @@ -105,7 +174,19 @@ def _apply_regex( class EnvResolver(BaseResolver): - # ENV Resolver -- full regex is ^\${spock\.env:.*}$ + """Class for resolving environmental variables + + Attributes: + _annotation_set: current set of supported resolver annotations + CLIP_ENV_PATTERN: regex for the front half + CLIP_REGEX_OP: compiled regex for front half + END_ENV_PATTERN: regex for back half + END_REGEX_OP: comiled regex for back half + FULL_ENV_PATTERN: full regex pattern + FULL_REGEX_OP: compiled regex for full regex + + """ + # ENV Resolver -- full regex is ^\${spock\.env\.?([a-z]*?):.*}$ CLIP_ENV_PATTERN = r"^\${spock\.env\.?([a-z]*?):" CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) END_ENV_PATTERN = r"}$" @@ -114,6 +195,7 @@ class EnvResolver(BaseResolver): FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) def __init__(self): + """Init for EnvResolver """ super(EnvResolver, self).__init__() def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: @@ -140,7 +222,20 @@ def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: return typed_env, annotation @staticmethod - def _get_from_env(default_value: Optional[str], env_value: str): + def _get_from_env(default_value: Optional[str], env_value: str) -> Optional[str]: + """Gets a value from an environmental variable + + Args: + default_value: default value to fall back on for the env resolver + env_value: current string of the env variable to get + + Returns: + string or None for the resolved env variable + + Raises: + _SpockResolverError if the env variable is not available or if no default is specified + + """ # Attempt to get the env variable if default_value == "None": maybe_env = os.getenv(env_value) @@ -155,6 +250,20 @@ def _get_from_env(default_value: Optional[str], env_value: str): class CryptoResolver(BaseResolver): + """Class for resolving cryptographic variables + + Attributes: + _annotation_set: current set of supported resolver annotations + CLIP_ENV_PATTERN: regex for the front half + CLIP_REGEX_OP: compiled regex for front half + END_ENV_PATTERN: regex for back half + END_REGEX_OP: comiled regex for back half + FULL_ENV_PATTERN: full regex pattern + FULL_REGEX_OP: compiled regex for full regex + _salt: current cryptographic salt + _key: current cryptographic key + + """ # ENV Resolver -- full regex is ^\${spock\.crypto:.*}$ CLIP_ENV_PATTERN = r"^\${spock\.crypto:" CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) @@ -164,6 +273,12 @@ class CryptoResolver(BaseResolver): FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) def __init__(self, salt: str, key: ByteString): + """Init for CryptoResolver + + Args: + salt: cryptographic salt to use + key: cryptographic key to use + """ super(CryptoResolver, self).__init__() self._salt = salt self._key = key diff --git a/spock/backend/saver.py b/spock/backend/saver.py index 5a8ede1a..32cb1375 100644 --- a/spock/backend/saver.py +++ b/spock/backend/saver.py @@ -13,7 +13,7 @@ from spock.backend.handler import BaseHandler from spock.backend.utils import _callable_2_str, _get_iter, _recurse_callables from spock.backend.wrappers import Spockspace -from spock.utils import add_info, get_packages +from spock.utils import add_info, get_packages, _T class BaseSaver(BaseHandler): # pylint: disable=too-few-public-methods @@ -28,11 +28,16 @@ class BaseSaver(BaseHandler): # pylint: disable=too-few-public-methods """ - def __init__(self, s3_config=None): + def __init__(self, s3_config: Optional[_T] = None): + """Init function for base class + + Args: + s3_config: optional s3Config object for S3 support + """ super(BaseSaver, self).__init__(s3_config=s3_config) def dict_payload(self, payload: Spockspace) -> Dict: - """Clean up the config payload so it can be returned as a dict representation + """Clean up the config payload so that it can be returned as a dict representation Args: payload: dirty payload @@ -81,7 +86,7 @@ def save( fname = "" if file_name is None else f"{file_name}." name = f"{fname}{uuid_str}.spock.cfg{file_extension}" # Fix up values -- parameters - out_dict = self._clean_up_values(payload) + out_dict = self.dict_payload(payload) # Handle any env annotations that are present # Just stuff them into the dictionary crypto_flag = False @@ -119,11 +124,12 @@ def save( raise e @abstractmethod - def _clean_up_values(self, payload: Spockspace) -> Dict: + def _clean_up_values(self, payload: Spockspace, remove_crypto: bool = True) -> Dict: """Clean up the config payload so it can be written to file Args: payload: dirty payload + remove_crypto: try and remove crypto values if present Returns: clean_dict: cleaned output payload @@ -198,13 +204,18 @@ class AttrSaver(BaseSaver): """ - def __init__(self, s3_config=None): + def __init__(self, s3_config: Optional[_T] = None): + """Init for AttrSaver class + + Args: + s3_config: s3Config object for S3 support + """ super().__init__(s3_config=s3_config) def __call__(self, *args, **kwargs): return AttrSaver(*args, **kwargs) - def _clean_up_values(self, payload: Spockspace) -> Dict: + def _clean_up_values(self, payload: Spockspace, remove_crypto: bool = True) -> Dict: # Dictionary to recursively write to out_dict = {} # All of the classes are defined at the top level @@ -216,6 +227,11 @@ def _clean_up_values(self, payload: Spockspace) -> Dict: clean_dict = self._clean_output(out_dict) # Clip any empty dictionaries clean_dict = {k: v for k, v in clean_dict.items() if len(v) > 0} + if remove_crypto: + if "__salt__" in clean_dict: + _ = clean_dict.pop("__salt__") + if "__key__" in clean_dict: + _ = clean_dict.pop("__key__") return clean_dict def _clean_tuner_values(self, payload: Spockspace) -> Dict: diff --git a/spock/backend/wrappers.py b/spock/backend/wrappers.py index 50a7d76a..93a7079b 100644 --- a/spock/backend/wrappers.py +++ b/spock/backend/wrappers.py @@ -22,13 +22,16 @@ def __init__(self, **kwargs): @property def __repr_dict__(self): + """Handles making a clean dict to hind the salt and key on print""" return {k: v for k, v in self.__dict__.items() if k not in {"__key__", "__salt__"}} def __repr__(self): + """Overloaded repr to pretty print the spock object""" # Remove aliases in YAML print yaml.Dumper.ignore_aliases = lambda *args: True return yaml.dump(self.__repr_dict__, default_flow_style=False) def __iter__(self): + """Iter for the underlying dictionary""" for k, v in self.__dict__.items(): yield k, v diff --git a/spock/builder.py b/spock/builder.py index 2427479c..2a89723a 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -74,9 +74,8 @@ def __init__( lazy: bool = False, no_cmd_line: bool = False, s3_config: Optional[_T] = None, - key: Optional[Union[str, ByteString]], - salt: Optional[str], - crypto: Optional[Union[str, Tuple[str, ByteString]]] = None, + key: Optional[Union[str, ByteString]] = None, + salt: Optional[str] = None, **kwargs, ): """Init call for ConfigArgBuilder @@ -806,7 +805,7 @@ def _maybe_crypto( salt: Optional[str], s3_config: Optional[_T] = None, salt_len: int = 16, - ): + ) -> Tuple[str, ByteString]: """Handles setting up the underlying cryptography needs Args: @@ -831,7 +830,7 @@ def _get_salt( env_resolver: EnvResolver, salt_len: int, s3_config: Optional[_T] = None, - ): + ) -> str: """ Args: @@ -858,7 +857,7 @@ def _get_key( key: Optional[Union[str, ByteString]], env_resolver: EnvResolver, s3_config: Optional[_T] = None, - ): + ) -> ByteString: """ Args: @@ -885,7 +884,7 @@ def _get_key( return key @staticmethod - def _handle_yaml_read(value: str, access: str, s3_config: Optional[_T] = None, encode: bool = False): + def _handle_yaml_read(value: str, access: str, s3_config: Optional[_T] = None, encode: bool = False) -> Union[str, ByteString]: """Reads in a salt/key yaml Args: diff --git a/spock/handlers.py b/spock/handlers.py index 2a1d370d..5022a39f 100644 --- a/spock/handlers.py +++ b/spock/handlers.py @@ -29,7 +29,7 @@ class Handler(ABC): """ - def load(self, path: Path, s3_config=None) -> Dict: + def load(self, path: Path, s3_config: Optional[_T] = None) -> Dict: """Load function for file type This handles s3 path conversion for all handler types pre load call @@ -202,7 +202,7 @@ def _save( raise NotImplementedError @staticmethod - def _handle_possible_s3_load_path(path: Path, s3_config=None) -> Union[str, Path]: + def _handle_possible_s3_load_path(path: Path, s3_config: Optional[_T] = None) -> Union[str, Path]: """Handles the possibility of having to handle loading from a S3 path Checks to see if it detects a S3 uri and if so triggers imports of s3 functionality and handles the file @@ -229,7 +229,7 @@ def _handle_possible_s3_load_path(path: Path, s3_config=None) -> Union[str, Path @staticmethod def _handle_possible_s3_save_path( - path: Path, name: str, create_path: bool, s3_config=None + path: Path, name: str, create_path: bool, s3_config: Optional[_T] = None ) -> Tuple[str, bool]: """Handles the possibility of having to save to a S3 path diff --git a/tests/base/test_state.py b/tests/base/test_state.py index fd6306ab..4957f9a9 100644 --- a/tests/base/test_state.py +++ b/tests/base/test_state.py @@ -46,4 +46,8 @@ def test_serialization_deserialization(self, monkeypatch, tmp_path): *all_configs, desc="Test Builder", ).generate() + delattr(config_values, '__key__') + delattr(config_values, '__salt__') + delattr(de_serial_config, '__key__') + delattr(de_serial_config, '__salt__') assert config_values == de_serial_config From 66ef5a2f3f8b4613158b7aa63201d273673b8a70 Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Mon, 16 May 2022 11:13:40 -0400 Subject: [PATCH 5/7] linted --- spock/addons/tune/payload.py | 6 +++--- spock/backend/resolvers.py | 7 +++++-- spock/backend/saver.py | 2 +- spock/backend/wrappers.py | 4 +++- spock/builder.py | 8 ++++++-- spock/handlers.py | 4 +++- 6 files changed, 21 insertions(+), 10 deletions(-) diff --git a/spock/addons/tune/payload.py b/spock/addons/tune/payload.py index 14fcb486..7f15bf65 100644 --- a/spock/addons/tune/payload.py +++ b/spock/addons/tune/payload.py @@ -5,11 +5,11 @@ """Handles the tuner payload backend""" -from spock.backend.payload import BasePayload -from spock.backend.utils import get_attr_fields, _T - from typing import Optional +from spock.backend.payload import BasePayload +from spock.backend.utils import _T, get_attr_fields + class TunerPayload(BasePayload): """Handles building the payload for tuners diff --git a/spock/backend/resolvers.py b/spock/backend/resolvers.py index 3d6d87c5..09e05d53 100644 --- a/spock/backend/resolvers.py +++ b/spock/backend/resolvers.py @@ -7,7 +7,7 @@ import os import re from abc import ABC, abstractmethod -from typing import Any, ByteString, Optional, Tuple, Pattern, Union +from typing import Any, ByteString, Optional, Pattern, Tuple, Union from spock.backend.utils import decrypt_value from spock.exceptions import _SpockResolverError @@ -23,6 +23,7 @@ class BaseResolver(ABC): _annotation_set: current set of supported resolver annotations """ + def __init__(self): """Init for BaseResolver class""" self._annotation_set = {"crypto", "inject"} @@ -186,6 +187,7 @@ class EnvResolver(BaseResolver): FULL_REGEX_OP: compiled regex for full regex """ + # ENV Resolver -- full regex is ^\${spock\.env\.?([a-z]*?):.*}$ CLIP_ENV_PATTERN = r"^\${spock\.env\.?([a-z]*?):" CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) @@ -195,7 +197,7 @@ class EnvResolver(BaseResolver): FULL_REGEX_OP = re.compile(FULL_ENV_PATTERN) def __init__(self): - """Init for EnvResolver """ + """Init for EnvResolver""" super(EnvResolver, self).__init__() def resolve(self, value: Any, value_type: _T) -> Tuple[Any, Optional[str]]: @@ -264,6 +266,7 @@ class CryptoResolver(BaseResolver): _key: current cryptographic key """ + # ENV Resolver -- full regex is ^\${spock\.crypto:.*}$ CLIP_ENV_PATTERN = r"^\${spock\.crypto:" CLIP_REGEX_OP = re.compile(CLIP_ENV_PATTERN) diff --git a/spock/backend/saver.py b/spock/backend/saver.py index 32cb1375..1b248656 100644 --- a/spock/backend/saver.py +++ b/spock/backend/saver.py @@ -13,7 +13,7 @@ from spock.backend.handler import BaseHandler from spock.backend.utils import _callable_2_str, _get_iter, _recurse_callables from spock.backend.wrappers import Spockspace -from spock.utils import add_info, get_packages, _T +from spock.utils import _T, add_info, get_packages class BaseSaver(BaseHandler): # pylint: disable=too-few-public-methods diff --git a/spock/backend/wrappers.py b/spock/backend/wrappers.py index 93a7079b..50dbc062 100644 --- a/spock/backend/wrappers.py +++ b/spock/backend/wrappers.py @@ -23,7 +23,9 @@ def __init__(self, **kwargs): @property def __repr_dict__(self): """Handles making a clean dict to hind the salt and key on print""" - return {k: v for k, v in self.__dict__.items() if k not in {"__key__", "__salt__"}} + return { + k: v for k, v in self.__dict__.items() if k not in {"__key__", "__salt__"} + } def __repr__(self): """Overloaded repr to pretty print the spock object""" diff --git a/spock/builder.py b/spock/builder.py index 2a89723a..ce1716b4 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -874,7 +874,9 @@ def _get_key( key = Fernet.generate_key() # Byte string is assumed to be a direct key elif os.path.splitext(key)[1] in {".yaml", ".YAML", ".yml", ".YML"}: - key = self._handle_yaml_read(key, access="key", s3_config=s3_config, encode=True) + key = self._handle_yaml_read( + key, access="key", s3_config=s3_config, encode=True + ) else: # Byte string is assumed to be a direct key # So only handle the str here @@ -884,7 +886,9 @@ def _get_key( return key @staticmethod - def _handle_yaml_read(value: str, access: str, s3_config: Optional[_T] = None, encode: bool = False) -> Union[str, ByteString]: + def _handle_yaml_read( + value: str, access: str, s3_config: Optional[_T] = None, encode: bool = False + ) -> Union[str, ByteString]: """Reads in a salt/key yaml Args: diff --git a/spock/handlers.py b/spock/handlers.py index 5022a39f..369e9fe3 100644 --- a/spock/handlers.py +++ b/spock/handlers.py @@ -202,7 +202,9 @@ def _save( raise NotImplementedError @staticmethod - def _handle_possible_s3_load_path(path: Path, s3_config: Optional[_T] = None) -> Union[str, Path]: + def _handle_possible_s3_load_path( + path: Path, s3_config: Optional[_T] = None + ) -> Union[str, Path]: """Handles the possibility of having to handle loading from a S3 path Checks to see if it detects a S3 uri and if so triggers imports of s3 functionality and handles the file From 075a2a81a54ebc8ae56d6430d11ec7708a22410f Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Mon, 16 May 2022 12:42:27 -0400 Subject: [PATCH 6/7] fixed tuner missing salt and key arg --- spock/builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spock/builder.py b/spock/builder.py index ce1716b4..fbd1755a 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -280,7 +280,7 @@ def _handle_tuner_objects( from spock.addons.tune.builder import TunerBuilder from spock.addons.tune.payload import TunerPayload - tuner_builder = TunerBuilder(*tune_args, **kwargs, lazy=self._lazy) + tuner_builder = TunerBuilder(*tune_args, **kwargs, lazy=self._lazy, salt=self.salt, key=self.key) tuner_payload = TunerPayload(s3_config=s3_config) return tuner_builder, tuner_payload except ImportError: From dd620ede3cc0d86fef1b3a4a1285d7eb2826db52 Mon Sep 17 00:00:00 2001 From: Nicholas Cilfone Date: Mon, 16 May 2022 12:45:14 -0400 Subject: [PATCH 7/7] linted --- spock/builder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/spock/builder.py b/spock/builder.py index fbd1755a..2048dde3 100644 --- a/spock/builder.py +++ b/spock/builder.py @@ -280,7 +280,9 @@ def _handle_tuner_objects( from spock.addons.tune.builder import TunerBuilder from spock.addons.tune.payload import TunerPayload - tuner_builder = TunerBuilder(*tune_args, **kwargs, lazy=self._lazy, salt=self.salt, key=self.key) + tuner_builder = TunerBuilder( + *tune_args, **kwargs, lazy=self._lazy, salt=self.salt, key=self.key + ) tuner_payload = TunerPayload(s3_config=s3_config) return tuner_builder, tuner_payload except ImportError: