def _run(self, args):
"""Add an OSCAL component/subcomponent to the specified component.
This method takes input a filename and a list of comma-seperated element path. Element paths are field aliases.
The method first finds the parent model from the file and loads the file into the model.
Then the method executes 'add' for each of the element paths specified.
"""
args = args.__dict__
if args[const.ARG_FILE] is None:
raise err.TrestleError(
f'Argument "-{const.ARG_FILE_SHORT}" is required')
if args[const.ARG_ELEMENT] is None:
raise err.TrestleError(
f'Argument "-{const.ARG_ELEMENT}" is required')
file_path = pathlib.Path(args[const.ARG_FILE])
# Get parent model and then load json into parent model
parent_model, parent_alias = fs.get_contextual_model_type(
file_path.absolute())
parent_object = parent_model.oscal_read(file_path.absolute())
parent_element = Element(
parent_object,
utils.classname_to_alias(parent_model.__name__, 'json'))
# Do _add for each element_path specified in args
element_paths: list[str] = args[const.ARG_ELEMENT].split(',')
for elm_path_str in element_paths:
element_path = ElementPath(elm_path_str)
self.add(file_path, element_path, parent_model, parent_element)
def get_singular_alias(alias_path: str, contextual_mode: bool = False) -> str:
"""
Get the alias in the singular form from a jsonpath.
If contextual_mode is True and contextual_path is None, it assumes alias_path is relative to the directory the user
is running trestle from.
"""
if len(alias_path.strip()) == 0:
raise err.TrestleError('Invalid jsonpath.')
singular_alias: str = ''
full_alias_path = alias_path
if contextual_mode:
_, full_model_alias = get_contextual_model_type()
first_alias_a = full_model_alias.split('.')[-1]
first_alias_b = alias_path.split('.')[0]
if first_alias_a == first_alias_b:
full_model_alias = '.'.join(full_model_alias.split('.')[:-1])
full_alias_path = '.'.join([full_model_alias, alias_path]).strip('.')
path_parts = full_alias_path.split(const.ALIAS_PATH_SEPARATOR)
if len(path_parts) < 2:
raise err.TrestleError('Invalid jsonpath.')
model_types = []
root_model_alias = path_parts[0]
found = False
for module_name in const.MODELTYPE_TO_MODELMODULE.values():
model_type, model_alias = utils.get_root_model(module_name)
if root_model_alias == model_alias:
found = True
model_types.append(model_type)
break
if not found:
raise err.TrestleError(
f'{root_model_alias} is an invalid root model alias.')
model_type = model_types[0]
for i in range(1, len(path_parts)):
if utils.is_collection_field_type(model_type):
model_type = utils.get_inner_type(model_type)
i = i + 1
else:
model_type = model_type.alias_to_field_map()[
path_parts[i]].outer_type_
model_types.append(model_type)
if not utils.is_collection_field_type(model_type):
raise err.TrestleError('Not a valid generic collection model.')
last_alias = path_parts[-1]
parent_model_type = model_types[-2]
singular_alias = utils.classname_to_alias(
utils.get_inner_type(parent_model_type.alias_to_field_map()
[last_alias].outer_type_).__name__, 'json')
return singular_alias
def create_stripped_model_type(
cls,
stripped_fields: List[str] = None,
stripped_fields_aliases: List[str] = None
) -> Type['OscalBaseModel']:
"""Use introspection to create a model that removes the fields.
Either 'stripped_fields' or 'stripped_fields_aliases' need to be passed, not both.
Returns a model class definition that can be used to instanciate a model.
"""
if stripped_fields is not None and stripped_fields_aliases is not None:
raise err.TrestleError(
'Either "stripped_fields" or "stripped_fields_aliases" need to be passed, not both.'
)
# create alias to field_name mapping
excluded_fields = []
if stripped_fields is not None:
excluded_fields = stripped_fields
elif stripped_fields_aliases is not None:
alias_to_field = cls.alias_to_field_map()
try:
excluded_fields = [
alias_to_field[key].name for key in stripped_fields_aliases
]
except KeyError as e:
raise err.TrestleError(
f'Field {str(e)} does not exists in the model')
current_fields = cls.__fields__
new_fields_for_model = {}
# Build field list
for current_mfield in current_fields.values():
if current_mfield.name in excluded_fields:
continue
# Validate name in the field
# Cehcke behaviour with an alias
if current_mfield.required:
new_fields_for_model[current_mfield.name] = (
current_mfield.outer_type_,
Field(...,
title=current_mfield.name,
alias=current_mfield.alias))
else:
new_fields_for_model[current_mfield.name] = (
Optional[current_mfield.outer_type_],
Field(None,
title=current_mfield.name,
alias=current_mfield.alias))
new_model = create_model(cls.__name__,
__base__=OscalBaseModel,
**new_fields_for_model) # type: ignore
# TODO: This typing cast should NOT be necessary. Potentially fixable with a fix to pydantic. Issue #175
new_model = cast(Type[OscalBaseModel], new_model)
return new_model
def get_root_model(module_name: str) -> Tuple[Type[Any], str]:
"""Get the root model class and alias based on the module."""
try:
module = importlib.import_module(module_name)
except ModuleNotFoundError as e:
raise err.TrestleError(str(e))
if hasattr(module, 'Model'):
model_metadata = next(iter(module.Model.__fields__.values()))
return (model_metadata.type_, model_metadata.alias)
else:
raise err.TrestleError('Invalid module')
def add(cls, file_path, element_path, parent_model, parent_element):
"""For a file_path and element_path, add a child model to the parent_element of a given parent_model.
First we find the child model at the specified element path and instantiate it with default values.
Then we check if there's already existing element at that path, in which case we append the child model
to the existing list of dict.
Then we set up an action plan to update the model (specified by file_path) in memory, create a file
at the same location and write the file.
"""
element_path_list = element_path.get_full_path_parts()
if '*' in element_path_list:
raise err.TrestleError(
'trestle add does not support Wildcard element path.')
# Get child model
try:
child_model = utils.get_target_model(element_path_list,
parent_model)
# Create child element with sample values
child_object = utils.get_sample_model(child_model)
if parent_element.get_at(element_path) is not None:
# The element already exists
if type(parent_element.get_at(element_path)) is list:
child_object = parent_element.get_at(
element_path) + child_object
elif type(parent_element.get_at(element_path)) is dict:
child_object = {
**parent_element.get_at(element_path),
**child_object
}
else:
raise err.TrestleError(
'Already exists and is not a list or dictionary.')
except Exception as e:
raise err.TrestleError(f'Bad element path. {str(e)}')
update_action = UpdateAction(sub_element=child_object,
dest_element=parent_element,
sub_element_path=element_path)
create_action = CreatePathAction(file_path.absolute(), True)
write_action = WriteFileAction(
file_path.absolute(), parent_element,
FileContentType.to_content_type(file_path.suffix))
add_plan = Plan()
add_plan.add_action(update_action)
add_plan.add_action(create_action)
add_plan.add_action(write_action)
add_plan.simulate()
add_plan.execute()
def get_contextual_file_type(path: pathlib.Path) -> FileContentType:
"""Return the file content type for files in the given directory, if it's a trestle project."""
if not is_valid_project_model_path(path):
raise err.TrestleError('Trestle project not found.')
for file_or_directory in path.iterdir():
if file_or_directory.is_file():
return FileContentType.to_content_type(file_or_directory.suffix)
for file_or_directory in path.iterdir():
if file_or_directory.is_dir():
return get_contextual_file_type(file_or_directory)
raise err.TrestleError('No files found in the project.')
def copy_to(self,
new_oscal_type: Type['OscalBaseModel']) -> 'OscalBaseModel':
"""
Opportunistic copy operation between similar types of data classes.
Due to the way in which oscal is constructed we get a set of similar / the same definition across various
oscal models. Due to the lack of guarantees that they are the same we cannot easily 'collapse' the mode.
Args:
new_oscal_type: The desired type of oscal model
Returns:
Opportunistic copy of the data into the new model type.
"""
logger.debug('Copy to started')
# FIXME: This needs to be tested. Unsure of behavior.
if self.__class__.__name__ == new_oscal_type.__name__:
logger.debug('Dict based copy too ')
return new_oscal_type.parse_obj(
self.dict(exclude_none=True, by_alias=True))
if ('__root__' in self.__fields__ and len(self.__fields__) == 1
and '__root__' in new_oscal_type.__fields__
and len(new_oscal_type.__fields__) == 1):
logger.debug('Root element based copy too')
return new_oscal_type.parse_obj(self.__root__)
# bad place here.
raise err.TrestleError('Provided inconsistent classes.')
def oscal_write(self, path: pathlib.Path, minimize_json=False) -> None:
"""
Write oscal objects.
OSCAL schema mandates that top level elements are wrapped in a singular
json/yaml field. This function handles both json and yaml output as well
as
"""
class_name = self.__class__.__name__
# It would be nice to pass through the description but I can't seem to and
# it does not affect the output
dynamic_parser = {}
dynamic_parser[classname_to_alias(class_name, 'field')] = (
self.__class__,
Field(self, title=classname_to_alias(class_name, 'field'), alias=classname_to_alias(class_name, 'json'))
)
wrapper_model = create_model(class_name, __base__=OscalBaseModel, **dynamic_parser) # type: ignore
# Default behaviour is strange here.
wrapped_model = wrapper_model(**{classname_to_alias(class_name, 'json'): self})
yaml_suffix = ['.yaml', '.yml']
json_suffix = ['.json']
encoding = 'utf8'
write_file = pathlib.Path(path).open('w', encoding=encoding)
if path.suffix in yaml_suffix:
yaml.dump(yaml.safe_load(wrapped_model.json(exclude_none=True, by_alias=True)), write_file)
pass
elif path.suffix in json_suffix:
write_file.write(wrapped_model.json(exclude_none=True, by_alias=True, indent=2))
else:
raise err.TrestleError('Unknown file type')
def oscal_read(cls, path: pathlib.Path) -> 'OscalBaseModel':
"""
Read OSCAL objects.
Handles the fact OSCAL wrap's top level elements and also deals with both yaml and json.
"""
# Define valid extensions
yaml_suffix = ['.yaml', '.yml']
json_suffix = ['.json']
# Create the wrapper model.
class_name = cls.__name__
dynamic_passer = {}
dynamic_passer[utils.class_to_oscal(
class_name,
'field')] = (cls,
Field(...,
title=utils.class_to_oscal(class_name, 'json'),
alias=utils.class_to_oscal(class_name, 'json')))
wrapper_model = create_model('Wrapped' + class_name,
__base__=OscalBaseModel,
**dynamic_passer)
if path.suffix in yaml_suffix:
return wrapper_model.parse_obj(yaml.safe_load(
path.open())).__dict__[utils.class_to_oscal(
class_name, 'field')]
elif path.suffix in json_suffix:
return wrapper_model.parse_file(path).__dict__[
utils.class_to_oscal(class_name, 'field')]
else:
raise err.TrestleError('Unknown file type')
def test_execute_failure(tmp_trestle_dir: pathlib.Path) -> None:
"""Ensure create plan failure will return clean return codes from run."""
args = argparse.Namespace(extension='json', name='my_catalog')
with mock.patch('trestle.core.models.plans.Plan.simulate') as simulate_mock:
simulate_mock.side_effect = err.TrestleError('stuff')
rc = create.CreateCmd.create_object('catalog', Catalog, args)
assert rc == 1
def add(cls, element_path, parent_model, parent_element):
"""For a element_path, add a child model to the parent_element of a given parent_model.
First we find the child model at the specified element path and instantiate it with default values.
Then we check if there's already existing element at that path, in which case we append the child model
to the existing list of dict.
Then we set up an action plan to update the model (specified by file_path) in memory, create a file
at the same location and write the file.
We update the parent_element to prepare for next adds in the chain
"""
element_path_list = element_path.get_full_path_parts()
if '*' in element_path_list:
raise err.TrestleError(
'trestle add does not support Wildcard element path.')
# Get child model
try:
child_model = utils.get_target_model(element_path_list,
parent_model)
# Create child element with sample values
child_object = gens.generate_sample_model(child_model)
if parent_element.get_at(element_path) is not None:
# The element already exists
if type(parent_element.get_at(element_path)) is list:
child_object = parent_element.get_at(
element_path) + child_object
elif type(parent_element.get_at(element_path)) is dict:
child_object = {
**parent_element.get_at(element_path),
**child_object
}
else:
raise err.TrestleError(
'Already exists and is not a list or dictionary.')
except Exception as e:
raise err.TrestleError(f'Bad element path. {str(e)}')
update_action = UpdateAction(sub_element=child_object,
dest_element=parent_element,
sub_element_path=element_path)
parent_element = parent_element.set_at(element_path, child_object)
return update_action, parent_element
def alias_to_classname(alias: str, mode: str) -> str:
"""
Return class name based dashed or snake alias.
This is applicable creating dynamic wrapper model for a list or dict field.
"""
if mode == 'json':
return snake_to_upper_camel(alias.replace('-', '_'))
elif mode == 'field':
return snake_to_upper_camel(alias)
else:
raise err.TrestleError('Bad option')
def remove(cls, element_path: ElementPath, parent_model: Type[OscalBaseModel],
parent_element: Element) -> Tuple[RemoveAction, Element]:
"""For the element_path, remove a model from the parent_element of a given parent_model.
First we check if there is an existing element at that path
If not, we complain.
Then we set up an action plan to update the model (specified by file_path) in memory,
return the action and return the parent_element.
LIMITATIONS:
1. This does not remove elements of a list or dict. Instead, the entire list or dict is removed.
2. This cannot remove arbitrarily named elements that are not specified in the schema.
For example, "responsible-parties" contains named elements, e.g., "organisation". The tool will not
remove the "organisation" as it is not in the schema, but one can remove its elements, e.g., "party-uuids".
"""
element_path_list = element_path.get_full_path_parts()
if '*' in element_path_list:
raise err.TrestleError('trestle remove does not support Wildcard element path.')
deleting_element = parent_element.get_at(element_path)
if deleting_element is not None:
# The element already exists
if type(deleting_element) is list:
logger.warning(
'Warning: trestle remove does not support removing elements of a list: '
'this removes the entire list'
)
elif type(deleting_element) is dict:
logger.warning(
'Warning: trestle remove does not support removing dict elements: '
'this removes the entire dict element'
)
else:
raise err.TrestleError(f'Bad element path: {str(element_path)}')
remove_action = RemoveAction(parent_element, element_path)
return remove_action, parent_element
def class_to_oscal(class_name: str, mode: str) -> str:
"""
Return oscal json or field element name based on class name.
This is applicable when asking for a singular element.
"""
parts = pascal_case_split(class_name)
if mode == 'json':
return '-'.join(map(str.lower, parts))
elif mode == 'field':
return '_'.join(map(str.lower, parts))
else:
raise err.TrestleError('Bad option')
def get_contextual_model_type(
path: pathlib.Path = None) -> Tuple[Type[OscalBaseModel], str]:
"""Get the full contextual model class and full jsonpath for the alias based on the contextual path."""
if path is None:
path = pathlib.Path.cwd()
if not is_valid_project_model_path(path):
raise err.TrestleError(f'Trestle project not found at {path}')
root_path = get_trestle_project_root(path)
project_model_path = get_project_model_path(path)
if root_path is None or project_model_path is None:
raise err.TrestleError('Trestle project not found')
relative_path = path.relative_to(str(root_path))
project_type = relative_path.parts[0] # catalogs, profiles, etc
module_name = const.MODELTYPE_TO_MODELMODULE[project_type]
model_relative_path = pathlib.Path(*relative_path.parts[2:])
model_type, model_alias = utils.get_root_model(module_name)
full_alias = model_alias
for i in range(len(model_relative_path.parts)):
tmp_path = root_path.joinpath(*relative_path.parts[:2],
*model_relative_path.parts[:i + 1])
alias = extract_alias(tmp_path)
if i > 0 or model_alias != alias:
model_alias = alias
full_alias = f'{full_alias}.{model_alias}'
if utils.is_collection_field_type(model_type):
model_type = utils.get_inner_type(model_type)
else:
model_type = model_type.alias_to_field_map()[alias].outer_type_
return model_type, full_alias
def test_assemble_execution_failure(testdata_dir: pathlib.Path, tmp_trestle_dir: pathlib.Path) -> None:
"""Test execution of assemble plan fails."""
test_data_source = testdata_dir / 'split_merge/step4_split_groups_array/catalogs'
catalogs_dir = pathlib.Path('catalogs/')
# Copy files from test/data/split_merge/step4
shutil.rmtree(catalogs_dir)
shutil.rmtree(pathlib.Path('dist'))
shutil.copytree(test_data_source, catalogs_dir)
with mock.patch('trestle.core.models.plans.Plan.simulate') as simulate_mock:
simulate_mock.side_effect = err.TrestleError('simulation error')
rc = AssembleCmd().assemble_model(
'catalog', Catalog, argparse.Namespace(name='mycatalog', extension='json', verbose=1)
)
assert rc == 1
def classname_to_alias(classname: str, mode: str) -> str:
"""
Return oscal key name or field element name based on class name.
This is applicable when asking for a singular element.
"""
suffix = classname.split('.')[-1]
if mode == 'json':
return camel_to_dash(suffix)
elif mode == 'field':
return camel_to_snake(suffix)
else:
raise err.TrestleError('Bad option')
def robust_datetime_serialization(input_dt: datetime.datetime) -> str:
"""Return a nicely formatted string for in a format compatible with OSCAL specifications.
Args:
input_dt: Input datetime to convert to a string.
Returns:
String in isoformat to the millisecond enforcing that timezone offset is provided.
Raises:
TrestleError: Error is raised if datetime object does not contain sufficient timezone information.
"""
# fail if the input datetime is not aware - ie it has no associated timezone
if input_dt.tzinfo is None:
raise err.TrestleError('Missing timezone in datetime')
if input_dt.tzinfo.utcoffset(input_dt) is None:
raise err.TrestleError('Missing utcoffset in datetime')
# use this leave in original timezone rather than utc
# return input_dt.astimezone().isoformat(timespec='milliseconds') noqa: E800
# force it to be utc
return input_dt.astimezone(
datetime.timezone.utc).isoformat(timespec='milliseconds')
def test_import_failure_simulate_plan(tmp_trestle_dir: pathlib.Path) -> None:
"""Test model failures throw errors and exit badly."""
rand_str = ''.join(random.choice(string.ascii_letters) for x in range(16))
catalog_file = f'{tmp_trestle_dir.parent}/{rand_str}.json'
catalog_data = generators.generate_sample_model(
trestle.oscal.catalog.Catalog)
catalog_data.oscal_write(pathlib.Path(catalog_file))
with patch(
'trestle.core.models.plans.Plan.simulate') as simulate_plan_mock:
simulate_plan_mock.side_effect = err.TrestleError('stuff')
args = argparse.Namespace(file=catalog_file,
output='imported',
verbose=True)
i = importcmd.ImportCmd()
rc = i._run(args)
assert rc == 1
def get_target_model(element_path_parts: List[str], current_model: BaseModel) -> BaseModel:
"""Get the target model from the parts of a Element Path.
Takes as input a list, containing parts of an ElementPath as str and expressed in aliases,
and the parent model to follow the ElementPath in.
Returns the type of the model at the specified ElementPath of the input model.
"""
try:
for index in range(1, len(element_path_parts)):
if is_collection_field_type(current_model):
# Return the model class inside the collection
current_model = get_inner_type(current_model)
else:
current_model = current_model.alias_to_field_map()[element_path_parts[index]].outer_type_
return current_model
except Exception as e:
raise err.TrestleError(f'Possibly bad element path. {str(e)}')
def generate_sample_value_by_type(
type_: type,
field_name: str,
) -> Union[datetime, bool, int, str, float, Enum]:
"""Given a type, return sample value.
Includes the Optional use of passing down a parent_model
"""
# FIXME: Should be in separate generator module as it inherits EVERYTHING
if type_ is datetime:
return datetime.now().astimezone()
elif type_ is bool:
return False
elif type_ is int:
return 0
elif type_ is str:
return 'REPLACE_ME'
elif type_ is float:
return 0.00
elif issubclass(type_, ConstrainedStr) or 'ConstrainedStr' in type_.__name__:
# This code here is messy. we need to meet a set of constraints. If we do
# TODO: explore whether there is a
if 'uuid' == field_name:
return str(uuid.uuid4())
elif field_name == 'date_authorized':
return date.today().isoformat()
return '00000000-0000-4000-8000-000000000000'
elif 'ConstrainedIntValue' in type_.__name__:
# create an int value as close to the floor as possible does not test upper bound
multiple = type_.multiple_of or 1 # default to every integer
floor = type_.ge or type_.gt + 1 or 0 # default to 0
if math.remainder(floor, multiple) == 0:
return floor
else:
return (floor + 1) * multiple
elif issubclass(type_, Enum):
# keys and values diverge due to hypens in oscal names
return type_(list(type_.__members__.values())[0])
elif type_ is pydantic.networks.EmailStr:
return pydantic.networks.EmailStr('*****@*****.**')
elif type_ is pydantic.networks.AnyUrl:
# TODO: Cleanup: this should be usable from a url.. but it's not inuitive.
return pydantic.networks.AnyUrl('https://sample.com/replaceme.html', scheme='http', host='sample.com')
else:
raise err.TrestleError(f'Fatal: Bad type in model {type_}')
def test_import_failure_parse_file(tmp_trestle_dir: pathlib.Path) -> None:
"""Test model failures throw errors and exit badly."""
sample_data = {'id': '0000'}
rand_str = ''.join(random.choice(string.ascii_letters) for x in range(16))
sample_file = pathlib.Path(
f'{tmp_trestle_dir.parent}/{rand_str}.json').open('w+',
encoding='utf8')
sample_file.write(json.dumps(sample_data))
sample_file.close()
with patch('trestle.core.parser.parse_file') as parse_file_mock:
parse_file_mock.side_effect = err.TrestleError('stuff')
args = argparse.Namespace(
file=f'{tmp_trestle_dir.parent}/{rand_str}.json',
output='catalog',
verbose=True)
i = importcmd.ImportCmd()
rc = i._run(args)
assert rc == 1
def test_trestle_fail_on_task_exception(tmp_trestle_dir: pathlib.Path) -> None:
"""Force an exception in a running trestle task and ensure it is captured."""
# Load and overwrite config file
section_name = 'task.pass-fail'
config_object = configparser.ConfigParser()
config_path = pathlib.Path(const.TRESTLE_CONFIG_DIR) / const.TRESTLE_CONFIG_FILE
config_object.read(config_path)
# add section
config_object.add_section(section_name)
config_object[section_name]['execute_status'] = 'True'
config_object[section_name]['simulate_status'] = 'True'
config_object.write(config_path.open('w'))
# Now good.
with mock.patch('trestle.tasks.base_task.PassFail.execute') as simulate_mock:
simulate_mock.side_effect = err.TrestleError('stuff')
args = argparse.Namespace(name='task', list=False, verbose=1, task='pass-fail', config=None, info=False)
rc = taskcmd.TaskCmd()._run(args)
assert rc > 0
def copy_to(self, new_oscal_type: Type['OscalBaseModel']) -> 'OscalBaseModel':
"""
Copy operation that explicilty does type conversion.
Input parameter is a class of type OscalBaseModel NOT a a class isntance.
"""
logger.debug('Copy to started')
if self.__class__.__name__ == new_oscal_type.__name__:
logger.debug('Dict based copy too ')
return new_oscal_type.parse_obj(self.dict(exclude_none=True, by_alias=True))
if ('__root__' in self.__fields__ and len(self.__fields__) == 1 and '__root__' in new_oscal_type.__fields__
and len(new_oscal_type.__fields__) == 1):
logger.debug('Root element based copy too')
return new_oscal_type.parse_obj(self.__root__)
# bad place here.
raise err.TrestleError('Provided inconsistent classes.')
def test_run_failure_plan_execute(tmp_path, sample_catalog_minimal):
"""Test failure plan execute() in _run on RemoveCmd."""
# Create a temporary file as a valid arg for trestle remove:
content_type = FileContentType.JSON
catalog_def_dir, catalog_def_file = test_utils.prepare_trestle_project_dir(
tmp_path, content_type, sample_catalog_minimal,
test_utils.CATALOGS_DIR)
testargs = [
'trestle', 'remove', '-f',
str(catalog_def_file), '-e', 'catalog.metadata.responsible-parties'
]
with mock.patch('trestle.core.models.plans.Plan.simulate'):
with mock.patch(
'trestle.core.models.plans.Plan.execute') as execute_mock:
execute_mock.side_effect = err.TrestleError('stuff')
with patch.object(sys, 'argv', testargs):
exitcode = Trestle().run()
assert exitcode == 1
def generate_sample_model(model: Union[Type[TG], List[TG], Dict[str, TG]]) -> TG:
"""Given a model class, generate an object of that class with sample values."""
# FIXME: Typing is wrong.
# TODO: The typing here is very generic - which may cause some pain. It may be more appropriate to create a wrapper
# Function for the to level execution. This would imply restructuring some other parts of the code.
model_type = model
# This block normalizes model type down to
if utils.is_collection_field_type(model): # type: ignore
model_type = utils.get_origin(model) # type: ignore
model = utils.get_inner_type(model) # type: ignore
model = cast(TG, model) # type: ignore
model_dict = {}
# this block is needed to avoid situations where an inbuilt is inside a list / dict.
if issubclass(model, OscalBaseModel):
for field in model.__fields__:
outer_type = model.__fields__[field].outer_type_
# Check for unions. This is awkward due to allow support for python 3.7
# It also does not inspect for which union we want. Should be removable with oscal 1.0.0
if utils.get_origin(outer_type) == Union:
outer_type = outer_type.__args__[0]
if model.__fields__[field].required:
""" FIXME: This type_ could be a List or a Dict """
if utils.is_collection_field_type(outer_type) or issubclass(outer_type, OscalBaseModel):
model_dict[field] = generate_sample_model(outer_type)
else:
model_dict[field] = generate_sample_value_by_type(outer_type, field)
# Note: this assumes list constrains in oscal are always 1 as a minimum size. if two this may still fail.
else:
# There is set of circumstances where a m
if model_type is list:
return [generate_sample_value_by_type(model, '')]
elif model_type is dict:
return {'REPLACE_ME': generate_sample_value_by_type(model, '')}
err.TrestleError('Unhandled collection type.')
if model_type is list:
return [model(**model_dict)]
elif model_type is dict:
return {'REPLACE_ME': model(**model_dict)}
return model(**model_dict)
def oscal_read(cls, path: pathlib.Path) -> 'OscalBaseModel':
"""
Read OSCAL objects.
Handles the fact OSCAL wrap's top level elements and also deals with both yaml and json.
"""
# Create the wrapper model.
alias = classname_to_alias(cls.__name__, 'json')
content_type = FileContentType.to_content_type(path.suffix)
if content_type == FileContentType.YAML:
return cls.parse_obj(yaml.safe_load(path.open())[alias])
elif content_type == FileContentType.JSON:
obj = load_file(
path,
json_loads=cls.__config__.json_loads,
)
return cls.parse_obj(obj[alias])
else:
raise err.TrestleError('Unknown file type')
def test_import_load_file_failure(tmp_trestle_dir: pathlib.Path) -> None:
"""Test model load failures."""
# Create a file with bad json
sample_data = '"star": {'
rand_str = ''.join(random.choice(string.ascii_letters) for x in range(16))
bad_file = pathlib.Path(f'{tmp_trestle_dir.parent}/{rand_str}.json').open(
'w+', encoding='utf8')
bad_file.write(sample_data)
bad_file.close()
with patch('trestle.utils.fs.load_file') as load_file_mock:
load_file_mock.side_effect = err.TrestleError('stuff')
args = argparse.Namespace(file=bad_file.name,
output='imported',
verbose=True)
i = importcmd.ImportCmd()
rc = i._run(args)
assert rc == 1
# Force PermissionError:
with patch('trestle.utils.fs.load_file') as load_file_mock:
load_file_mock.side_effect = PermissionError
args = argparse.Namespace(file=bad_file.name,
output='imported',
verbose=True)
i = importcmd.ImportCmd()
rc = i._run(args)
assert rc == 1
# Force JSONDecodeError:
with patch('trestle.utils.fs.load_file') as load_file_mock:
load_file_mock.side_effect = JSONDecodeError(msg='Extra data:',
doc=bad_file.name,
pos=0)
args = argparse.Namespace(file=bad_file.name,
output='imported',
verbose=True)
i = importcmd.ImportCmd()
rc = i._run(args)
assert rc == 1
# This is in case the same tmp_trestle_dir.parent is used, as across succeeding scopes of one pytest
os.chmod(bad_file.name, 0o600)
os.remove(bad_file.name)
def generate_sample_value_by_type(
type_: type,
field_name: str,
parent_model: Optional[Type[OscalBaseModel]] = None
) -> Union[datetime, bool, int, str, float, Enum]:
"""Given a type, return sample value.
Includes the Optional use of passing down a parent_model
"""
# FIXME: Should be in separate generator module as it inherits EVERYTHING
if type_ is datetime:
return datetime.now().astimezone()
elif type_ is bool:
return False
elif type_ is int:
return 0
elif type_ is str:
return 'REPLACE_ME'
elif type_ is float:
return 0.00
elif issubclass(type_, ConstrainedStr) or 'ConstrainedStr' in str(type):
# This code here is messy. we need to meet a set of constraints. If we do
# not do so it fails to generate.
if 'uuid' == field_name:
return str(uuid.uuid4())
elif parent_model == trestle.oscal.ssp.DateAuthorized:
return date.today().isoformat()
return '00000000-0000-4000-8000-000000000000'
elif issubclass(type_, Enum):
# keys and values diverge due to hypens in oscal names
return type_(list(type_.__members__.values())[0])
elif type_ is pydantic.networks.EmailStr:
return pydantic.networks.EmailStr('*****@*****.**')
elif type_ is pydantic.networks.AnyUrl:
# TODO: Cleanup: this should be usable from a url.. but it's not inuitive.
return pydantic.networks.AnyUrl('https://sample.com/replaceme.html',
scheme='http',
host='sample.com')
else:
raise err.TrestleError('Fatal: Bad type in model')
def get_inner_type(
collection_field_type: Union[Type[List[TG]], Type[Dict[str,
TG]]]) -> Type[TG]:
"""Get the inner model in a generic collection model such as a List or a Dict.
For a dict the return type is of the value and not the key.
Args:
collection_field_type: Provided type annotation from a pydantic object
Returns:
The desired type.
"""
try:
# Pydantic special cases ust be dealt with here:
if getattr(collection_field_type, '__name__',
None) == 'ConstrainedListValue':
return collection_field_type.item_type # type: ignore
return typing_extensions.get_args(collection_field_type)[-1]
except Exception as e:
logger.debug(e)
raise err.TrestleError('Model type is not a Dict or List') from e
