feat: adds check to TrestleRule to match compliance-trestle CSV fields (#173)

Signed-off-by: Jennifer Power <barnabei.jennifer@gmail.com>

Author: jpower432

tests/conftest.py

@@ -28,6 +28,7 @@
 )
 from trestlebot import const
 from trestlebot.transformers.trestle_rule import (
+    Check,
     ComponentInfo,
     Control,
     Parameter,
@@ -168,6 +169,7 @@ def test_rule() -> TrestleRule:
             alternative_values={},
             default_value="test",
         ),
+        check=Check(name="test_check", description="test check"),
         profile=Profile(
             description="test", href="test", include_controls=[Control(id="ac-1")]
         ),

tests/data/yaml/test_complete_rule.yaml

@@ -6,6 +6,9 @@ x-trestle-rule-info:
     description: prm_1 description
     alternative-values: {'default': '5%', '5pc': '5%', '10pc': '10%', '15pc': '15%', '20pc': '20%'}
     default-value: '5%'
+  check:
+    name: my_check
+    description: My check description
   profile:
     description: Simple NIST Profile
     href: profiles/simplified_nist_profile/profile.json

tests/trestlebot/tasks/test_rule_transform_task.py

@@ -62,7 +62,7 @@ def test_rule_transform_task(tmp_trestle_dir: str) -> None:
 
     assert component is not None
     assert component.props is not None
-    assert len(component.props) == 5
+    assert len(component.props) == 7
     assert component.props[0].name == RULE_ID
     assert component.props[0].value == "example_rule_1"
     assert component.props[1].name == RULE_DESCRIPTION

tests/trestlebot/transformers/test_csv_transformer.py

@@ -38,6 +38,8 @@ def test_csv_builder(test_rule: TrestleRule, tmp_trestle_dir: str) -> None:
     assert row["Parameter_Value_Default"] == test_rule.parameter.default_value  # type: ignore
     assert row["Profile_Description"] == test_rule.profile.description
     assert row["Profile_Source"] == test_rule.profile.href
+    assert row["Check_Id"] == test_rule.check.name  # type: ignore
+    assert row["Check_Description"] == test_rule.check.description  # type: ignore
 
     trestle_root = pathlib.Path(tmp_trestle_dir)
     tmp_csv_path = trestle_root.joinpath("test.csv")

tests/trestlebot/transformers/test_yaml_transformer.py

@@ -46,6 +46,9 @@ def test_rule_transformer() -> None:
     assert rule.parameter.default_value == "5%"
     assert rule.profile.description == "Simple NIST Profile"
     assert rule.profile.href == "profiles/simplified_nist_profile/profile.json"
+    assert rule.check is not None
+    assert rule.check.name == "my_check"
+    assert rule.check.description == "My check description"
 
 
 def test_rules_transform_with_incomplete_rule() -> None:

trestlebot/const.py

@@ -25,6 +25,7 @@
 NAME = "name"
 DESCRIPTION = "description"
 PARAMETER = "parameter"
+CHECK = "check"
 PROFILE = "profile"
 HREF = "href"
 ALTERNATIVE_VALUES = "alternative-values"

trestlebot/transformers/csv_transformer.py

@@ -13,6 +13,8 @@
 import trestle.tasks.csv_to_oscal_cd as csv_to_oscal_cd
 from trestle.common.const import TRESTLE_GENERIC_NS
 from trestle.tasks.csv_to_oscal_cd import (
+    CHECK_DESCRIPTION,
+    CHECK_ID,
     COMPONENT_DESCRIPTION,
     COMPONENT_TITLE,
     COMPONENT_TYPE,
@@ -35,6 +37,7 @@
     ToRulesTransformer,
 )
 from trestlebot.transformers.trestle_rule import (
+    Check,
     ComponentInfo,
     Control,
     Parameter,
@@ -65,13 +68,15 @@ def transform(self, row: Dict[str, str]) -> TrestleRule:
         profile = self._extract_profile(row)
         component_info = self._extract_component_info(row)
         parameter = self._extract_parameter(row)
+        check = self._extract_check(row)
 
         return TrestleRule(
             name=rule_info[const.NAME],
             description=rule_info[const.DESCRIPTION],
             component=component_info,
             parameter=parameter,
             profile=profile,
+            check=check,
         )
 
     def _extract_rule_info(self, row: Dict[str, str]) -> Dict[str, str]:
@@ -106,6 +111,16 @@ def _extract_parameter(self, row: Dict[str, str]) -> Optional[Parameter]:
             )
         return None
 
+    def _extract_check(self, row: Dict[str, str]) -> Optional[Check]:
+        """Extract check information from a CSV row."""
+        check_name = row.get(csv_to_oscal_cd.CHECK_ID, None)
+        if check_name:
+            return Check(
+                name=check_name,
+                description=row.get(csv_to_oscal_cd.CHECK_DESCRIPTION, ""),
+            )
+        return None
+
     def _extract_component_info(self, row: Dict[str, str]) -> ComponentInfo:
         """Extract component information from a CSV row."""
         return ComponentInfo(
@@ -141,6 +156,12 @@ def transform(self, rule: TrestleRule) -> Dict[str, str]:
         }
         if rule.parameter is not None:
             merged_dict.update(self._add_parameter(rule.parameter))
+        if rule.check is not None:
+            check: Dict[str, str] = {
+                CHECK_ID: rule.check.name,
+                CHECK_DESCRIPTION: rule.check.description,
+            }
+            merged_dict.update(check)
         return merged_dict
 
     def _add_profile(self, profile: Profile) -> Dict[str, str]:

trestlebot/transformers/trestle_rule.py

@@ -46,14 +46,25 @@ class ComponentInfo(BaseModel):
     description: str
 
 
+class Check(BaseModel):
+    """Check dataclass."""
+
+    name: str
+    description: str
+
+    class Config:
+        allow_population_by_field_name = True
+
+
 class TrestleRule(BaseModel):
     """TrestleRule dataclass."""
 
     name: str
     description: str
     component: ComponentInfo
-    parameter: Optional[Parameter]
     profile: Profile
+    check: Optional[Check]
+    parameter: Optional[Parameter]
 
 
 def get_default_rule() -> TrestleRule:

trestlebot/transformers/yaml_transformer.py

@@ -18,6 +18,7 @@
     ToRulesTransformer,
 )
 from trestlebot.transformers.trestle_rule import (
+    Check,
     ComponentInfo,
     Parameter,
     Profile,
@@ -66,12 +67,17 @@ def transform(self, blob: str) -> TrestleRule:
                     rule_info_data[const.PARAMETER]
                 )
 
+            check_instance: Optional[Check] = None
+            if const.CHECK in rule_info_data:
+                check_instance = Check.parse_obj(rule_info_data[const.CHECK])
+
             rule_info_instance: TrestleRule = TrestleRule(
                 name=rule_info_data[const.NAME],
                 description=rule_info_data[const.DESCRIPTION],
                 component=component_info_instance,
                 parameter=parameter_instance,
                 profile=profile_info_instance,
+                check=check_instance,
             )
 
         except KeyError as e:
@@ -134,4 +140,8 @@ def _to_rule_info(rule: TrestleRule) -> Dict[str, Any]:
             rule_info[const.RULE_INFO_TAG][const.PARAMETER] = rule.parameter.dict(
                 by_alias=True, exclude_unset=True
             )
+        if rule.check is not None:
+            rule_info[const.RULE_INFO_TAG][const.CHECK] = rule.check.dict(
+                by_alias=True, exclude_unset=True
+            )
         return rule_info