diff --git a/contentctl/actions/new_content.py b/contentctl/actions/new_content.py index 0a54cf11..74b4dbce 100644 --- a/contentctl/actions/new_content.py +++ b/contentctl/actions/new_content.py @@ -4,7 +4,6 @@ import questionary from typing import Any from contentctl.input.new_content_questions import NewContentQuestions -from contentctl.output.new_content_yml_output import NewContentYmlOutput from contentctl.objects.config import new, NewContentType import uuid from datetime import datetime diff --git a/contentctl/output/detection_writer.py b/contentctl/output/detection_writer.py deleted file mode 100644 index 2f439ca9..00000000 --- a/contentctl/output/detection_writer.py +++ /dev/null @@ -1,28 +0,0 @@ - -import yaml - - -class DetectionWriter: - - @staticmethod - def writeYmlFile(file_path : str, obj : dict) -> None: - - new_obj = dict() - new_obj["name"] = obj["name"] - new_obj["id"] = obj["id"] - new_obj["version"] = obj["version"] - new_obj["date"] = obj["date"] - new_obj["author"] = obj["author"] - new_obj["type"] = obj["type"] - new_obj["status"] = obj["status"] - new_obj["description"] = obj["description"] - new_obj["data_source"] = obj["data_source"] - new_obj["search"] = obj["search"] - new_obj["how_to_implement"] = obj["how_to_implement"] - new_obj["known_false_positives"] = obj["known_false_positives"] - new_obj["references"] = obj["references"] - new_obj["tags"] = obj["tags"] - new_obj["tests"] = obj["tests"] - - with open(file_path, 'w') as outfile: - yaml.safe_dump(new_obj, outfile, default_flow_style=False, sort_keys=False) \ No newline at end of file diff --git a/contentctl/output/new_content_yml_output.py b/contentctl/output/new_content_yml_output.py deleted file mode 100644 index 38730b37..00000000 --- a/contentctl/output/new_content_yml_output.py +++ /dev/null @@ -1,56 +0,0 @@ -import os -import pathlib -from contentctl.objects.enums import SecurityContentType -from contentctl.output.yml_writer import YmlWriter -import pathlib -from contentctl.objects.config import NewContentType -class NewContentYmlOutput(): - output_path: pathlib.Path - - def __init__(self, output_path:pathlib.Path): - self.output_path = output_path - - - def writeObjectNewContent(self, object: dict, subdirectory_name: str, type: NewContentType) -> None: - if type == NewContentType.detection: - - file_path = os.path.join(self.output_path, 'detections', subdirectory_name, self.convertNameToFileName(object['name'], object['tags']['product'])) - output_folder = pathlib.Path(self.output_path)/'detections'/subdirectory_name - #make sure the output folder exists for this detection - output_folder.mkdir(exist_ok=True) - - YmlWriter.writeYmlFile(file_path, object) - print("Successfully created detection " + file_path) - - elif type == NewContentType.story: - file_path = os.path.join(self.output_path, 'stories', self.convertNameToFileName(object['name'], object['tags']['product'])) - YmlWriter.writeYmlFile(file_path, object) - print("Successfully created story " + file_path) - - else: - raise(Exception(f"Object Must be Story or Detection, but is not: {object}")) - - - - def convertNameToFileName(self, name: str, product: list): - file_name = name \ - .replace(' ', '_') \ - .replace('-','_') \ - .replace('.','_') \ - .replace('/','_') \ - .lower() - - file_name = file_name + '.yml' - return file_name - - - def convertNameToTestFileName(self, name: str, product: list): - file_name = name \ - .replace(' ', '_') \ - .replace('-','_') \ - .replace('.','_') \ - .replace('/','_') \ - .lower() - - file_name = file_name + '.test.yml' - return file_name \ No newline at end of file diff --git a/contentctl/output/yml_output.py b/contentctl/output/yml_output.py deleted file mode 100644 index 93eae5dc..00000000 --- a/contentctl/output/yml_output.py +++ /dev/null @@ -1,66 +0,0 @@ -import os - -from contentctl.output.detection_writer import DetectionWriter -from contentctl.objects.detection import Detection - - -class YmlOutput(): - - - def writeDetections(self, objects: list, output_path : str) -> None: - for obj in objects: - file_path = obj.file_path - obj.id = str(obj.id) - - DetectionWriter.writeYmlFile(os.path.join(output_path, file_path), obj.dict( - exclude_none=True, - include = - { - "name": True, - "id": True, - "version": True, - "date": True, - "author": True, - "type": True, - "status": True, - "description": True, - "data_source": True, - "search": True, - "how_to_implement": True, - "known_false_positives": True, - "references": True, - "tags": - { - "analytic_story": True, - "asset_type": True, - "atomic_guid": True, - "confidence": True, - "impact": True, - "drilldown_search": True, - "mappings": True, - "message": True, - "mitre_attack_id": True, - "kill_chain_phases:": True, - "observable": True, - "product": True, - "required_fields": True, - "risk_score": True, - "security_domain": True - }, - "tests": - { - '__all__': - { - "name": True, - "attack_data": { - '__all__': - { - "data": True, - "source": True, - "sourcetype": True - } - } - } - } - } - )) \ No newline at end of file diff --git a/contentctl/output/yml_writer.py b/contentctl/output/yml_writer.py index 7d71762b..2e408c83 100644 --- a/contentctl/output/yml_writer.py +++ b/contentctl/output/yml_writer.py @@ -1,6 +1,21 @@ import yaml from typing import Any +from enum import StrEnum, IntEnum + +# Set the following so that we can write StrEnum and IntEnum +# to files. Otherwise, we will get the following errors when trying +# to write to files: +# yaml.representer.RepresenterError: ('cannot represent an object',..... +yaml.SafeDumper.add_multi_representer( + StrEnum, + yaml.representer.SafeRepresenter.represent_str +) + +yaml.SafeDumper.add_multi_representer( + IntEnum, + yaml.representer.SafeRepresenter.represent_int +) class YmlWriter: diff --git a/enums_test.py b/enums_test.py deleted file mode 100644 index fea6a545..00000000 --- a/enums_test.py +++ /dev/null @@ -1,46 +0,0 @@ -from pydantic import BaseModel, ConfigDict -from enum import StrEnum, IntEnum - -class SomeString(StrEnum): - one = "one" - two = "TWO" - -class SomeInt(IntEnum): - one = 1 - two = 2 - -class WithUseEnum(BaseModel): - ConfigDict(use_enum_values=True) - strval: SomeString - intval: SomeInt - - -class WithOutUseEnum(BaseModel): - strval: SomeString - intval: SomeInt - -withObj = WithUseEnum.model_validate({"strval": "one", "intval": "2"}) -withoutObj = WithOutUseEnum.model_validate({"strval": "one", "intval": "2"}) - - -print("With tests") -print(withObj.strval) -print(withObj.strval.upper()) -print(withObj.strval.value) -print(withObj.intval) -print(withObj.intval.value) -print(withObj.strval == SomeString.one) -print(withObj.strval == "ONE") -print(withObj.intval == SomeInt.two) -print(withObj.intval == 2) - - -print("Without tests") -print(withoutObj.strval) -print(withoutObj.strval.value) -print(withoutObj.intval) -print(withoutObj.intval.value) -print(withoutObj.strval == SomeString.one) -print(withoutObj.strval == "ONE") -print(withoutObj.intval == SomeInt.two) -print(withoutObj.intval == 2) \ No newline at end of file