Skip to content

Commit 31f46a2

Browse files
pyth0n1cpyth0n1c
committed
Remove use of .value on enums in code
now that they have been more strictly defined as IntEnum or StrEnum. This has not yet been tested.
1 parent 0999270 commit 31f46a2

16 files changed

+155
-116
lines changed

‎contentctl/actions/detection_testing/infrastructures/DetectionTestingInfrastructure.py‎

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ def test_detection(self, detection: Detection) -> None:
442442
self.format_pbar_string(
443443
TestReportingType.GROUP,
444444
test_group.name,
445-
FinalTestingStates.SKIP.value,
445+
FinalTestingStates.SKIP,
446446
start_time=time.time(),
447447
set_pbar=False,
448448
)
@@ -483,7 +483,7 @@ def test_detection(self, detection: Detection) -> None:
483483
self.format_pbar_string(
484484
TestReportingType.GROUP,
485485
test_group.name,
486-
TestingStates.DONE_GROUP.value,
486+
TestingStates.DONE_GROUP,
487487
start_time=setup_results.start_time,
488488
set_pbar=False,
489489
)
@@ -504,7 +504,7 @@ def setup_test_group(self, test_group: TestGroup) -> SetupTestGroupResults:
504504
self.format_pbar_string(
505505
TestReportingType.GROUP,
506506
test_group.name,
507-
TestingStates.BEGINNING_GROUP.value,
507+
TestingStates.BEGINNING_GROUP,
508508
start_time=setup_start_time
509509
)
510510
# https://github.com/WoLpH/python-progressbar/issues/164
@@ -544,7 +544,7 @@ def cleanup_test_group(
544544
self.format_pbar_string(
545545
TestReportingType.GROUP,
546546
test_group.name,
547-
TestingStates.DELETING.value,
547+
TestingStates.DELETING,
548548
start_time=test_group_start_time,
549549
)
550550

@@ -632,7 +632,7 @@ def execute_unit_test(
632632
self.format_pbar_string(
633633
TestReportingType.UNIT,
634634
f"{detection.name}:{test.name}",
635-
FinalTestingStates.SKIP.value,
635+
FinalTestingStates.SKIP,
636636
start_time=test_start_time,
637637
set_pbar=False,
638638
)
@@ -664,7 +664,7 @@ def execute_unit_test(
664664
self.format_pbar_string(
665665
TestReportingType.UNIT,
666666
f"{detection.name}:{test.name}",
667-
FinalTestingStates.ERROR.value,
667+
FinalTestingStates.ERROR,
668668
start_time=test_start_time,
669669
set_pbar=False,
670670
)
@@ -724,7 +724,7 @@ def execute_unit_test(
724724
res = "ERROR"
725725
link = detection.search
726726
else:
727-
res = test.result.status.value.upper() # type: ignore
727+
res = test.result.status.upper() # type: ignore
728728
link = test.result.get_summary_dict()["sid_link"]
729729

730730
self.format_pbar_string(
@@ -755,7 +755,7 @@ def execute_unit_test(
755755
self.format_pbar_string(
756756
TestReportingType.UNIT,
757757
f"{detection.name}:{test.name}",
758-
FinalTestingStates.PASS.value,
758+
FinalTestingStates.PASS,
759759
start_time=test_start_time,
760760
set_pbar=False,
761761
)
@@ -766,7 +766,7 @@ def execute_unit_test(
766766
self.format_pbar_string(
767767
TestReportingType.UNIT,
768768
f"{detection.name}:{test.name}",
769-
FinalTestingStates.SKIP.value,
769+
FinalTestingStates.SKIP,
770770
start_time=test_start_time,
771771
set_pbar=False,
772772
)
@@ -777,7 +777,7 @@ def execute_unit_test(
777777
self.format_pbar_string(
778778
TestReportingType.UNIT,
779779
f"{detection.name}:{test.name}",
780-
FinalTestingStates.FAIL.value,
780+
FinalTestingStates.FAIL,
781781
start_time=test_start_time,
782782
set_pbar=False,
783783
)
@@ -788,7 +788,7 @@ def execute_unit_test(
788788
self.format_pbar_string(
789789
TestReportingType.UNIT,
790790
f"{detection.name}:{test.name}",
791-
FinalTestingStates.ERROR.value,
791+
FinalTestingStates.ERROR,
792792
start_time=test_start_time,
793793
set_pbar=False,
794794
)
@@ -821,7 +821,7 @@ def execute_integration_test(
821821
test_start_time = time.time()
822822

823823
# First, check to see if the test should be skipped (Hunting or Correlation)
824-
if detection.type in [AnalyticsType.Hunting.value, AnalyticsType.Correlation.value]:
824+
if detection.type in [AnalyticsType.Hunting, AnalyticsType.Correlation]:
825825
test.skip(
826826
f"TEST SKIPPED: detection is type {detection.type} and cannot be integration "
827827
"tested at this time"
@@ -843,11 +843,11 @@ def execute_integration_test(
843843
# Determine the reporting state (we should only encounter SKIP/FAIL/ERROR)
844844
state: str
845845
if test.result.status == TestResultStatus.SKIP:
846-
state = FinalTestingStates.SKIP.value
846+
state = FinalTestingStates.SKIP
847847
elif test.result.status == TestResultStatus.FAIL:
848-
state = FinalTestingStates.FAIL.value
848+
state = FinalTestingStates.FAIL
849849
elif test.result.status == TestResultStatus.ERROR:
850-
state = FinalTestingStates.ERROR.value
850+
state = FinalTestingStates.ERROR
851851
else:
852852
raise ValueError(
853853
f"Status for (integration) '{detection.name}:{test.name}' was preemptively set"
@@ -891,7 +891,7 @@ def execute_integration_test(
891891
self.format_pbar_string(
892892
TestReportingType.INTEGRATION,
893893
f"{detection.name}:{test.name}",
894-
FinalTestingStates.FAIL.value,
894+
FinalTestingStates.FAIL,
895895
start_time=test_start_time,
896896
set_pbar=False,
897897
)
@@ -935,7 +935,7 @@ def execute_integration_test(
935935
if test.result is None:
936936
res = "ERROR"
937937
else:
938-
res = test.result.status.value.upper() # type: ignore
938+
res = test.result.status.upper() # type: ignore
939939

940940
# Get the link to the saved search in this specific instance
941941
link = f"https://{self.infrastructure.instance_address}:{self.infrastructure.web_ui_port}"
@@ -968,7 +968,7 @@ def execute_integration_test(
968968
self.format_pbar_string(
969969
TestReportingType.INTEGRATION,
970970
f"{detection.name}:{test.name}",
971-
FinalTestingStates.PASS.value,
971+
FinalTestingStates.PASS,
972972
start_time=test_start_time,
973973
set_pbar=False,
974974
)
@@ -979,7 +979,7 @@ def execute_integration_test(
979979
self.format_pbar_string(
980980
TestReportingType.INTEGRATION,
981981
f"{detection.name}:{test.name}",
982-
FinalTestingStates.SKIP.value,
982+
FinalTestingStates.SKIP,
983983
start_time=test_start_time,
984984
set_pbar=False,
985985
)
@@ -990,7 +990,7 @@ def execute_integration_test(
990990
self.format_pbar_string(
991991
TestReportingType.INTEGRATION,
992992
f"{detection.name}:{test.name}",
993-
FinalTestingStates.FAIL.value,
993+
FinalTestingStates.FAIL,
994994
start_time=test_start_time,
995995
set_pbar=False,
996996
)
@@ -1001,7 +1001,7 @@ def execute_integration_test(
10011001
self.format_pbar_string(
10021002
TestReportingType.INTEGRATION,
10031003
f"{detection.name}:{test.name}",
1004-
FinalTestingStates.ERROR.value,
1004+
FinalTestingStates.ERROR,
10051005
start_time=test_start_time,
10061006
set_pbar=False,
10071007
)
@@ -1077,7 +1077,7 @@ def retry_search_until_timeout(
10771077
self.format_pbar_string(
10781078
TestReportingType.UNIT,
10791079
f"{detection.name}:{test.name}",
1080-
TestingStates.PROCESSING.value,
1080+
TestingStates.PROCESSING,
10811081
start_time=start_time
10821082
)
10831083

@@ -1086,7 +1086,7 @@ def retry_search_until_timeout(
10861086
self.format_pbar_string(
10871087
TestReportingType.UNIT,
10881088
f"{detection.name}:{test.name}",
1089-
TestingStates.SEARCHING.value,
1089+
TestingStates.SEARCHING,
10901090
start_time=start_time,
10911091
)
10921092

@@ -1289,7 +1289,7 @@ def replay_attack_data_file(
12891289
self.format_pbar_string(
12901290
TestReportingType.GROUP,
12911291
test_group.name,
1292-
TestingStates.DOWNLOADING.value,
1292+
TestingStates.DOWNLOADING,
12931293
start_time=test_group_start_time
12941294
)
12951295

@@ -1307,7 +1307,7 @@ def replay_attack_data_file(
13071307
self.format_pbar_string(
13081308
TestReportingType.GROUP,
13091309
test_group.name,
1310-
TestingStates.REPLAYING.value,
1310+
TestingStates.REPLAYING,
13111311
start_time=test_group_start_time
13121312
)
13131313

‎contentctl/actions/detection_testing/progress_bar.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
import time
2-
from enum import Enum
2+
from enum import StrEnum
33
from tqdm import tqdm
44
import datetime
55

66

7-
class TestReportingType(str, Enum):
7+
class TestReportingType(StrEnum):
88
"""
99
5-char identifiers for the type of testing being reported on
1010
"""
@@ -21,7 +21,7 @@ class TestReportingType(str, Enum):
2121
INTEGRATION = "INTEG"
2222

2323

24-
class TestingStates(str, Enum):
24+
class TestingStates(StrEnum):
2525
"""
2626
Defined testing states
2727
"""
@@ -40,10 +40,10 @@ class TestingStates(str, Enum):
4040

4141

4242
# the longest length of any state
43-
LONGEST_STATE = max(len(w.value) for w in TestingStates)
43+
LONGEST_STATE = max(len(w) for w in TestingStates)
4444

4545

46-
class FinalTestingStates(str, Enum):
46+
class FinalTestingStates(StrEnum):
4747
"""
4848
The possible final states for a test (for pbar reporting)
4949
"""
@@ -82,7 +82,7 @@ def format_pbar_string(
8282
:returns: a formatted string for use w/ pbar
8383
"""
8484
# Extract and ljust our various fields
85-
field_one = test_reporting_type.value
85+
field_one = test_reporting_type
8686
field_two = test_name.ljust(MAX_TEST_NAME_LENGTH)
8787
field_three = state.ljust(LONGEST_STATE)
8888
field_four = datetime.timedelta(seconds=round(time.time() - start_time))

‎contentctl/actions/detection_testing/views/DetectionTestingView.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -110,11 +110,11 @@ def getSummaryObject(
110110
total_skipped += 1
111111

112112
# Aggregate production status metrics
113-
if detection.status == DetectionStatus.production.value: # type: ignore
113+
if detection.status == DetectionStatus.production:
114114
total_production += 1
115-
elif detection.status == DetectionStatus.experimental.value: # type: ignore
115+
elif detection.status == DetectionStatus.experimental:
116116
total_experimental += 1
117-
elif detection.status == DetectionStatus.deprecated.value: # type: ignore
117+
elif detection.status == DetectionStatus.deprecated:
118118
total_deprecated += 1
119119

120120
# Check if the detection is manual_test
@@ -178,7 +178,7 @@ def getSummaryObject(
178178
# Construct and return the larger results dict
179179
result_dict = {
180180
"summary": {
181-
"mode": self.config.getModeName(),
181+
"mode": self.config.mode.mode_name,
182182
"enable_integration_testing": self.config.enable_integration_testing,
183183
"success": overall_success,
184184
"total_detections": total_detections,

‎contentctl/actions/test.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from dataclasses import dataclass
22
from typing import List
33

4-
from contentctl.objects.config import test_common
4+
from contentctl.objects.config import test_common, Selected, Changes
55
from contentctl.objects.enums import DetectionTestingMode, DetectionStatus, AnalyticsType
66
from contentctl.objects.detection import Detection
77

@@ -78,19 +78,18 @@ def execute(self, input_dto: TestInputDto) -> bool:
7878
input_dto=manager_input_dto, output_dto=output_dto
7979
)
8080

81-
mode = input_dto.config.getModeName()
8281
if len(input_dto.detections) == 0:
8382
print(
84-
f"With Detection Testing Mode '{mode}', there were [0] detections found to test."
83+
f"With Detection Testing Mode '{input_dto.config.mode.mode_name}', there were [0] detections found to test."
8584
"\nAs such, we will quit immediately."
8685
)
8786
# Directly call stop so that the summary.yml will be generated. Of course it will not
8887
# have any test results, but we still want it to contain a summary showing that now
8988
# detections were tested.
9089
file.stop()
9190
else:
92-
print(f"MODE: [{mode}] - Test [{len(input_dto.detections)}] detections")
93-
if mode in [DetectionTestingMode.changes.value, DetectionTestingMode.selected.value]:
91+
print(f"MODE: [{input_dto.config.mode.mode_name}] - Test [{len(input_dto.detections)}] detections")
92+
if isinstance(input_dto.config.mode, Selected) or isinstance(input_dto.config.mode, Changes):
9493
files_string = '\n- '.join(
9594
[str(pathlib.Path(detection.file_path).relative_to(input_dto.config.path)) for detection in input_dto.detections]
9695
)

0 commit comments

Comments
 (0)