diff --git a/pyoaev/apis/signature.py b/pyoaev/apis/signature.py index 2450c2c..0fe1971 100644 --- a/pyoaev/apis/signature.py +++ b/pyoaev/apis/signature.py @@ -107,7 +107,7 @@ def send_signatures( return sig_data = json.loads(payload["execution_output_structured"]) - targets = sig_data["signatures"]["targets"] + targets = sig_data["expectation_signatures"]["targets"] envelopes = self._split_into_envelopes( payload, sig_data, @@ -199,8 +199,8 @@ def _build_envelope( targets_subset: list[dict[str, Any]], ) -> dict[str, Any]: subset_sig = dict(sig_data) - subset_sig["signatures"] = dict(sig_data["signatures"]) - subset_sig["signatures"]["targets"] = targets_subset + subset_sig["expectation_signatures"] = dict(sig_data["expectation_signatures"]) + subset_sig["expectation_signatures"]["targets"] = targets_subset envelope = dict(base_payload) envelope["execution_output_structured"] = json.dumps(subset_sig) diff --git a/pyoaev/signatures/models.py b/pyoaev/signatures/models.py index b8a2c25..a5640f2 100644 --- a/pyoaev/signatures/models.py +++ b/pyoaev/signatures/models.py @@ -17,7 +17,11 @@ model_validator, ) -from pyoaev.signatures.types import ExpectationType, InjectExecutionActions +from pyoaev.signatures.types import ( + ExecutionStatus, + ExpectationType, + InjectExecutionActions, +) class ToolErrorInfo(BaseModel): @@ -122,7 +126,7 @@ class SignatureOutputStructure(BaseModel): model_config = ConfigDict(populate_by_name=True, extra="forbid") - signatures: SignaturePayload + expectation_signatures: SignaturePayload def normalize_signature_payload(self) -> None: """ @@ -130,7 +134,7 @@ def normalize_signature_payload(self) -> None: """ normalized_targets: list[TargetSignatures] = [] - for target in self.signatures.targets: + for target in self.expectation_signatures.targets: if not target.signature_values: normalized_targets.append(target) continue @@ -156,7 +160,7 @@ def normalize_signature_payload(self) -> None: normalized_targets.append(normalized_target) - self.signatures.targets = normalized_targets + self.expectation_signatures.targets = normalized_targets class ExecutionDetails(BaseModel): @@ -167,14 +171,15 @@ class ExecutionDetails(BaseModel): start_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) end_time: datetime | None = None - execution_status: str = "unknown" + execution_status: ExecutionStatus | None = None execution_action: InjectExecutionActions | None = None @computed_field @property def execution_message(self) -> str: action = self.execution_action.value if self.execution_action else "unknown" - return f"Current action: {action} - Current status: {self.execution_status}" + status = self.execution_action.value if self.execution_action else "unknown" + return f"Current action: {action} - Current status: {status}" @computed_field @property @@ -191,19 +196,19 @@ def post_execution_update(self, tool_output: ToolOutput, now: datetime) -> None: self.end_time = now if tool_output.error_info and tool_output.error_info.exit_code != 0: - self.execution_status = "failed" + self.execution_status = ExecutionStatus.ERROR if tool_output.error_info.crash_timestamp: self.end_time = datetime.strptime( tool_output.error_info.crash_timestamp, "%Y-%m-%dT%H:%M:%SZ" ) elif tool_output.timeout_info: - self.execution_status = "timeout" + self.execution_status = ExecutionStatus.TIMEOUT elif tool_output.status == "partial": - self.execution_status = "partial" + self.execution_status = ExecutionStatus.ERROR else: - self.execution_status = "success" + self.execution_status = ExecutionStatus.EXECUTED - self.execution_action = InjectExecutionActions("complete") + self.execution_action = InjectExecutionActions("command_execution") class SignatureCallbackPayload(BaseModel): diff --git a/pyoaev/signatures/signature_manager.py b/pyoaev/signatures/signature_manager.py index 81db6ef..2845234 100644 --- a/pyoaev/signatures/signature_manager.py +++ b/pyoaev/signatures/signature_manager.py @@ -221,7 +221,9 @@ def send_signatures( Raises: SignatureTransmissionError: Validation failed, 4xx hit, or retries exhausted. """ - sig_output = SignatureOutputStructure(signatures=SignaturePayload(**signatures)) + sig_output = SignatureOutputStructure( + expectation_signatures=SignaturePayload(**signatures) + ) self.client.signature.send_signatures( inject_id, diff --git a/pyoaev/signatures/types.py b/pyoaev/signatures/types.py index 0c4f78b..436b0b4 100644 --- a/pyoaev/signatures/types.py +++ b/pyoaev/signatures/types.py @@ -7,6 +7,27 @@ class ExpectationType(str, Enum): VULNERABILITY = "VULNERABILITY" +class ExecutionStatus(str, Enum): + EXECUTED = "EXECUTED" + EXECUTED_WITH_CLEANUP_FAILURE = "EXECUTED_WITH_CLEANUP_FAILURE" + WARNING = "WARNING" + ACCESS_DENIED = "ACCESS_DENIED" + ERROR = "ERROR" + COMMAND_NOT_FOUND = "COMMAND_NOT_FOUND" + COMMAND_CANNOT_BE_EXECUTED = "COMMAND_CANNOT_BE_EXECUTED" + PREREQUISITE_FAILED = "PREREQUISITE_FAILED" + INVALID_USAGE = "INVALID_USAGE" + TIMEOUT = "TIMEOUT" + INTERRUPTED = "INTERRUPTED" + ASSET_AGENTLESS = "ASSET_AGENTLESS" + AGENT_INACTIVE = "AGENT_INACTIVE" + AGENT_OVERLOADED = "AGENT_OVERLOADED" + INFO = "INFO" + PARTIAL = "PARTIAL" + MAYBE_PREVENTED = "MAYBE_PREVENTED" + MAYBE_PARTIAL_PREVENTED = "MAYBE_PARTIAL_PREVENTED" + + class InjectExecutionActions(str, Enum): PREREQUISITE_CHECK = "prerequisite_check" PREREQUISITE_EXECUTION = "prerequisite_execution" diff --git a/test/signatures/constraints/signature_manager_post_execution_constraints.feature b/test/signatures/constraints/signature_manager_post_execution_constraints.feature index 16d7b47..17914e7 100644 --- a/test/signatures/constraints/signature_manager_post_execution_constraints.feature +++ b/test/signatures/constraints/signature_manager_post_execution_constraints.feature @@ -18,7 +18,7 @@ Feature: SignatureManager post-execution constraints Scenario: Tool crash sets execution_status to failed and uses crash timestamp as end_time Given a tool_output containing error_info with exit_code=1 and crash_timestamp="2024-06-26T06:05:00Z" When I call post_execution_updates with the execution_details, execution_signatures and tool_output - Then execution_status equals "failed" + Then execution_status equals "ERROR" And end_time equals "2024-06-26T06:05:00Z" And the execution signature model contains every previous parameter unchanged And the execution details model contain every previous parameter pair unchanged @@ -26,7 +26,7 @@ Feature: SignatureManager post-execution constraints Scenario: Timeout sets execution_status to timeout and includes available partial results Given a tool_output containing timeout_info with partial_results=["result-A", "result-B"] When I call post_execution_updates with the execution_details, execution_signatures and tool_output - Then execution_status equals "timeout" + Then execution_status equals "TIMEOUT" And the returned dict contains the partial results ["result-A", "result-B"] from timeout_info And the execution signature model contains every previous parameter unchanged And the execution details model contain every previous parameter pair unchanged @@ -34,6 +34,6 @@ Feature: SignatureManager post-execution constraints Scenario: Timeout with no partial results still sets execution_status to timeout Given a tool_output containing timeout_info with no partial results available When I call post_execution_updates with the execution_details, execution_signatures and tool_output - Then execution_status equals "timeout" + Then execution_status equals "TIMEOUT" And the execution signature model contains every previous parameter unchanged And the execution details model contain every previous parameter pair unchanged diff --git a/test/signatures/features/signature_manager_post_execution.feature b/test/signatures/features/signature_manager_post_execution.feature index ffa78a4..4673429 100644 --- a/test/signatures/features/signature_manager_post_execution.feature +++ b/test/signatures/features/signature_manager_post_execution.feature @@ -24,5 +24,5 @@ Feature: SignatureManager post-execution execution elements update And the execution details model contain every previous parameter pair unchanged And the end_time parameter in the execution details model is a datetime object And this end_time is chronologically greater than or equal to start_time "2024-06-26T06:00:00Z" - And the execution_status parameter in the execution details model is equal to "success" - And the execution_action parameter in the execution details model is equal to "complete" + And the execution_status parameter in the execution details model is equal to "EXECUTED" + And the execution_action parameter in the execution details model is equal to "command_execution" diff --git a/test/signatures/test_signature_manager_post_execution.py b/test/signatures/test_signature_manager_post_execution.py index 356959b..22f2d0e 100644 --- a/test/signatures/test_signature_manager_post_execution.py +++ b/test/signatures/test_signature_manager_post_execution.py @@ -6,6 +6,7 @@ from pyoaev.signatures.models import ExecutionDetails, ExecutionSignature from pyoaev.signatures.signature_manager import SignatureManager +from pyoaev.signatures.types import ExecutionStatus @scenario( @@ -75,7 +76,7 @@ def execution_signatures(): ) def execution_details(): return ExecutionDetails( - execution_status="", + execution_status=ExecutionStatus.INFO, start_time=datetime.strptime("2024-06-26T06:00:00Z", "%Y-%m-%dT%H:%M:%SZ"), ) @@ -85,7 +86,7 @@ def execution_details(): target_fixture="tool_output", ) def successful_tool_output(): - return {"status": "success"} + return {"status": "EXECUTED"} @given( @@ -203,7 +204,9 @@ def result_contains_datetime_end_time(context): ) @then(parsers.parse('execution_status equals "{status}"')) def execution_status_equals(context, status): - assert context["execution_details_result"].execution_status == status + assert context["execution_details_result"].execution_status == ExecutionStatus( + status + ) @then( diff --git a/test/signatures/test_signature_manager_transmission.py b/test/signatures/test_signature_manager_transmission.py index e414983..b06759b 100644 --- a/test/signatures/test_signature_manager_transmission.py +++ b/test/signatures/test_signature_manager_transmission.py @@ -11,6 +11,7 @@ from pyoaev.exceptions import OpenAEVUpdateError, SignatureTransmissionError from pyoaev.signatures.models import ExecutionDetails from pyoaev.signatures.signature_manager import SignatureManager +from pyoaev.signatures.types import ExecutionStatus @scenario( @@ -92,7 +93,7 @@ def context(): def _extract_targets(body: dict) -> list[dict]: """Parse targets from the SignatureCallbackPayload wire format.""" sig_data = json.loads(body["execution_output_structured"]) - return sig_data["signatures"]["targets"] + return sig_data["expectation_signatures"]["targets"] def _build_signature_payload( @@ -172,20 +173,20 @@ def _http_post(*args, **kwargs): context["status_plan"] = [200] context["error_body"] = "" context["inject_id"] = "inject-abc-001" - context["signatures"] = _build_signature_payload() + context["expectation_signatures"] = _build_signature_payload() context["signature_manager"] = SignatureManager(mock_client, logger=logger) @given(parsers.parse('a compiled post-execution payload for inject_id "{inject_id}"')) def compiled_post_execution_payload(context, inject_id): context["inject_id"] = inject_id - context["signatures"] = _build_signature_payload() + context["expectation_signatures"] = _build_signature_payload() @given(parsers.parse("an updated post-execution execution details object")) def updated_post_execution_execution_details(context): execution_details = ExecutionDetails( - execution_status="success", + execution_status=ExecutionStatus.EXECUTED, execution_action="complete", ) execution_details.end_time = execution_details.start_time + timedelta(0.1) @@ -203,7 +204,7 @@ def compiled_payload_single_target( signature_type, signature_value, ): - context["signatures"] = { + context["expectation_signatures"] = { "targets": [ { "signature_target": dict(_CANONICAL_SIGNATURE_TARGET), @@ -230,9 +231,9 @@ def compiled_large_payload(context): context["signature_manager"] = SignatureManager( context["mock_client"], logger=context["logger"], - max_payload_size=700, + max_payload_size=800, ) - context["signatures"] = { + context["expectation_signatures"] = { "targets": [ { "signature_target": { @@ -335,7 +336,7 @@ def compiled_payload_grouped_by_expectation( expectation_a, expectation_b, ): - context["signatures"] = { + context["expectation_signatures"] = { "targets": [ { "signature_target": dict(_CANONICAL_SIGNATURE_TARGET), @@ -376,7 +377,7 @@ def call_send_signatures(context, inject_id): context["signature_manager"].send_signatures( inject_id, context["execution_details"], - context["signatures"], + context["expectation_signatures"], ) except Exception as exc: context["send_exception"] = exc @@ -498,7 +499,7 @@ def assert_no_chunk_metadata(context): @then("the union of targets across all POST requests equals the original target set") def assert_targets_union_matches_original(context): - original_targets = context["signatures"]["targets"] + original_targets = context["expectation_signatures"]["targets"] sent_targets = [ target for call_item in context["captured_calls"]