Files
atomic-red-team/atomic_red_team/tests/test_docs.py
T
2026-05-02 18:30:22 -04:00

474 lines
16 KiB
Python

import json
import yaml
from mitreattack import navlayers
from atomic_red_team.docs import AtomicRedTeam, AtomicRedTeamDocs
def test_generate_technique_docs_writes_markdown(tmp_path):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
atomic_yaml_path = technique_dir / "T1001.yaml"
atomic_yaml_path.write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "Run command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"auto_generated_guid": "11111111-1111-4111-8111-111111111111",
"input_arguments": {
"path": {
"description": "File path",
"type": "path",
"default": "C:\\Temp\\file.txt",
}
},
"executor": {
"name": "command_prompt",
"command": "type #{path}",
"cleanup_command": "del #{path}",
},
}
],
}
)
)
attack_file = tmp_path / "enterprise-attack.json"
attack_file.write_text(
json.dumps(
{
"type": "bundle",
"id": "bundle--11111111-1111-4111-8111-111111111111",
"spec_version": "2.0",
"objects": [
{
"type": "attack-pattern",
"id": "attack-pattern--11111111-1111-4111-8111-111111111111",
"created": "2024-01-01T00:00:00.000Z",
"modified": "2024-01-01T00:00:00.000Z",
"name": "Data Obfuscation",
"description": "Technique description.",
"external_references": [
{"source_name": "mitre-attack", "external_id": "T1001"}
],
"x_mitre_platforms": ["Windows"],
"kill_chain_phases": [
{
"kill_chain_name": "mitre-attack",
"phase_name": "command-and-control",
}
],
}
],
}
)
)
output_path = technique_dir / "T1001.md"
atomic_red_team = AtomicRedTeam(atomics_dir=atomics_dir)
docs = AtomicRedTeamDocs(attack_file=attack_file, atomic_red_team=atomic_red_team)
docs.generate_technique_docs(
atomic_red_team.atomic_tests[0],
output_path,
)
output = output_path.read_text()
assert "# T1001 - Data Obfuscation" in output
assert "[Atomic Test #1: Run command](#atomic-test-1-run-command)" in output
assert "| path | File path | path | C:\Temp\file.txt|" in output
assert "```cmd\ntype #{path}\n```" in output
assert "```cmd\ndel #{path}\n```" in output
def test_generate_technique_docs_separates_markdown_blocks(tmp_path):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
atomic_yaml_path = technique_dir / "T1001.yaml"
atomic_yaml_path.write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "First command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"auto_generated_guid": "11111111-1111-4111-8111-111111111111",
"input_arguments": {
"path": {
"description": "File path",
"type": "path",
"default": "C:\\Temp",
}
},
"executor": {
"name": "command_prompt",
"command": "type #{path}",
"cleanup_command": "del #{path}",
},
"dependencies": [
{
"description": "Target must exist",
"prereq_command": "if exist #{path} (exit /b 0)",
"get_prereq_command": "mkdir #{path}",
}
],
},
{
"name": "Second command",
"description": "Runs another command.",
"supported_platforms": ["windows"],
"auto_generated_guid": "22222222-2222-4222-8222-222222222222",
"executor": {"name": "command_prompt", "command": "dir"},
},
],
}
)
)
attack_file = tmp_path / "enterprise-attack.json"
attack_file.write_text(
json.dumps(
{
"type": "bundle",
"id": "bundle--11111111-1111-4111-8111-111111111111",
"spec_version": "2.0",
"objects": [
{
"type": "attack-pattern",
"id": "attack-pattern--11111111-1111-4111-8111-111111111111",
"created": "2024-01-01T00:00:00.000Z",
"modified": "2024-01-01T00:00:00.000Z",
"name": "Data Obfuscation",
"description": "Technique description.",
"external_references": [
{"source_name": "mitre-attack", "external_id": "T1001"}
],
"x_mitre_platforms": ["Windows"],
"kill_chain_phases": [
{
"kill_chain_name": "mitre-attack",
"phase_name": "command-and-control",
}
],
}
],
}
)
)
output_path = technique_dir / "T1001.md"
atomic_red_team = AtomicRedTeam(atomics_dir=atomics_dir)
docs = AtomicRedTeamDocs(attack_file=attack_file, atomic_red_team=atomic_red_team)
docs.generate_technique_docs(atomic_red_team.atomic_tests[0], output_path)
output = output_path.read_text()
assert (
output
== """# T1001 - Data Obfuscation
## Description from ATT&CK
> Technique description.
[Source](https://attack.mitre.org/techniques/T1001)
## Atomic Tests
- [Atomic Test #1: First command](#atomic-test-1-first-command)
- [Atomic Test #2: Second command](#atomic-test-2-second-command)
### Atomic Test #1: First command
Runs a command.
**Supported Platforms:** Windows
**auto_generated_guid:** `11111111-1111-4111-8111-111111111111`
#### Inputs
| Name | Description | Type | Default Value |
|------|-------------|------|---------------|
| path | File path | path | C:\Temp|
#### Attack Commands: Run with `command_prompt`!
```cmd
type #{path}
```
#### Cleanup Commands
```cmd
del #{path}
```
#### Dependencies: Run with `command_prompt`!
##### Description: Target must exist
###### Check Prereq Commands
```cmd
if exist #{path} (exit /b 0)
```
###### Get Prereq Commands
```cmd
mkdir #{path}
```
### Atomic Test #2: Second command
Runs another command.
**Supported Platforms:** Windows
**auto_generated_guid:** `22222222-2222-4222-8222-222222222222`
#### Attack Commands: Run with `command_prompt`!
```cmd
dir
```
"""
)
def test_github_link_prompts_for_contribution_when_platform_missing(tmp_path):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
(technique_dir / "T1001.yaml").write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "Run command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"executor": {"name": "command_prompt", "command": "dir"},
}
],
}
)
)
atomic_red_team = AtomicRedTeam(atomics_dir=atomics_dir)
technique = {
"name": "Data Obfuscation",
"external_references": [
{"source_name": "mitre-attack", "external_id": "T1001"}
],
}
link = atomic_red_team.github_link_to_technique(
technique,
include_identifier=True,
only_platform="linux",
)
assert link == (
"T1001 Data Obfuscation "
"[CONTRIBUTE A TEST](https://github.com/redcanaryco/atomic-red-team/wiki/Contributing)"
)
def test_generate_navigator_layers_writes_mitreattack_layer(tmp_path):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
(technique_dir / "T1001.yaml").write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "Windows command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"executor": {"name": "command_prompt", "command": "dir"},
},
{
"name": "Linux command",
"description": "Runs a command.",
"supported_platforms": ["linux"],
"executor": {"name": "sh", "command": "ls"},
},
],
}
)
)
output_dir = tmp_path / "layers"
output_dir.mkdir()
docs = AtomicRedTeamDocs(
base_dir=tmp_path,
attack_file=tmp_path / "enterprise-attack.json",
atomic_red_team=AtomicRedTeam(atomics_dir=atomics_dir),
)
docs.generate_navigator_layers(output_dir)
layer = navlayers.Layer()
layer.from_file(str(output_dir / "art-navigator-layer-windows.json"))
technique = layer.layer.techniques[0]
assert layer.layer.name == "Atomic Red Team (Windows)"
assert layer.layer.versions.attack == "19"
assert layer.layer.filters.platforms == ["Windows"]
assert technique.techniqueID == "T1001"
assert technique.score == 1
assert technique.comment == "\n- Windows command\n"
def test_generate_all_the_docs_prints_generated_file_paths(
tmp_path, capsys, monkeypatch
):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
atomic_yaml_path = technique_dir / "T1001.yaml"
atomic_yaml_path.write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "Run command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"executor": {"name": "command_prompt", "command": "dir"},
}
],
}
)
)
attack_file = tmp_path / "enterprise-attack.json"
attack_file.write_text(
json.dumps(
{
"type": "bundle",
"id": "bundle--11111111-1111-4111-8111-111111111111",
"spec_version": "2.0",
"objects": [
{
"type": "attack-pattern",
"id": "attack-pattern--11111111-1111-4111-8111-111111111111",
"created": "2024-01-01T00:00:00.000Z",
"modified": "2024-01-01T00:00:00.000Z",
"name": "Data Obfuscation",
"description": "Technique description.",
"external_references": [
{"source_name": "mitre-attack", "external_id": "T1001"}
],
"x_mitre_platforms": ["Windows"],
"kill_chain_phases": [
{
"kill_chain_name": "mitre-attack",
"phase_name": "command-and-control",
}
],
}
],
}
)
)
indexes_dir = atomics_dir / "Indexes"
for path in [
indexes_dir / "Matrices",
indexes_dir / "Indexes-Markdown",
indexes_dir / "Indexes-CSV",
indexes_dir / "Attack-Navigator-Layers",
]:
path.mkdir(parents=True)
docs = AtomicRedTeamDocs(
base_dir=tmp_path,
attack_file=attack_file,
atomic_red_team=AtomicRedTeam(atomics_dir=atomics_dir),
)
monkeypatch.setattr(docs, "generate_indexes", lambda: None)
docs.generate_all_the_docs()
output = capsys.readouterr().out
expected_markdown_path = technique_dir / "T1001.md"
assert f"Generating docs for {atomic_yaml_path}" in output
assert f"=> {expected_markdown_path} => OK" in output
assert f"Generated docs for 1 techniques, 0 failures" in output
def test_generate_yaml_index_serializes_mitreattack_stix_dates(tmp_path):
atomics_dir = tmp_path / "atomics"
technique_dir = atomics_dir / "T1001"
technique_dir.mkdir(parents=True)
(technique_dir / "T1001.yaml").write_text(
yaml.safe_dump(
{
"attack_technique": "T1001",
"display_name": "Data Obfuscation",
"atomic_tests": [
{
"name": "Run command",
"description": "Runs a command.",
"supported_platforms": ["windows"],
"executor": {"name": "command_prompt", "command": "dir"},
}
],
}
)
)
attack_file = tmp_path / "enterprise-attack.json"
attack_file.write_text(
json.dumps(
{
"type": "bundle",
"id": "bundle--11111111-1111-4111-8111-111111111111",
"spec_version": "2.0",
"objects": [
{
"type": "attack-pattern",
"id": "attack-pattern--11111111-1111-4111-8111-111111111111",
"created": "2024-01-01T00:00:00.000Z",
"modified": "2024-01-01T00:00:00.000Z",
"name": "Data Obfuscation",
"description": "Technique description.",
"external_references": [
{"source_name": "mitre-attack", "external_id": "T1001"}
],
"x_mitre_platforms": ["Windows"],
"kill_chain_phases": [
{
"kill_chain_name": "mitre-attack",
"phase_name": "command-and-control",
}
],
}
],
}
)
)
output_path = tmp_path / "index.yaml"
docs = AtomicRedTeamDocs(
base_dir=tmp_path,
attack_file=attack_file,
atomic_red_team=AtomicRedTeam(atomics_dir=atomics_dir),
)
docs.generate_yaml_index(output_path)
output = yaml.safe_load(output_path.read_text())
assert output["command-and-control"]["T1001"]["technique"]["created"] == (
"2024-01-01T00:00:00.000Z"
)