use oatf::normalize::normalize;
use oatf::parse::parse;
use oatf::serialize::serialize;
use proptest::prelude::*;
fn arb_mode() -> impl Strategy<Value = String> {
prop_oneof![
Just("mcp_server".to_string()),
Just("mcp_client".to_string()),
Just("a2a_server".to_string()),
Just("a2a_client".to_string()),
]
}
fn arb_severity() -> impl Strategy<Value = String> {
prop_oneof![
Just("informational"),
Just("low"),
Just("medium"),
Just("high"),
Just("critical"),
]
.prop_map(|s| s.to_string())
}
fn arb_surface() -> impl Strategy<Value = String> {
prop_oneof![
Just("tools/list"),
Just("tools/call"),
Just("resources/read"),
Just("prompts/list"),
Just("message/send"),
]
.prop_map(|s| s.to_string())
}
fn build_single_phase_doc(mode: &str, tool_name: &str) -> String {
format!(
r#"oatf: "0.1"
attack:
execution:
mode: {mode}
state:
tools:
- name: {tool_name}
description: "A test tool"
inputSchema:
type: object"#,
)
}
fn build_multi_phase_doc(mode: &str, tool_name: &str) -> String {
format!(
r#"oatf: "0.1"
attack:
execution:
mode: {mode}
phases:
- name: exploit
state:
tools:
- name: {tool_name}
description: "A test tool"
inputSchema:
type: object
trigger:
event: tools/call
- name: terminal"#,
)
}
fn build_multi_actor_doc(mode: &str, tool_name: &str) -> String {
format!(
r#"oatf: "0.1"
attack:
execution:
actors:
- name: default
mode: {mode}
phases:
- name: exploit
state:
tools:
- name: {tool_name}
description: "A test tool"
inputSchema:
type: object
trigger:
event: tools/call
- name: terminal"#,
)
}
fn build_doc_with_indicators(mode: &str, surface: &str, severity: &str) -> String {
format!(
r#"oatf: "0.1"
attack:
severity: {severity}
execution:
mode: {mode}
phases:
- name: exploit
state:
tools:
- name: test-tool
description: "desc"
inputSchema:
type: object
trigger:
event: tools/call
- name: terminal
indicators:
- surface: {surface}
target: "tools[*].description"
pattern:
contains: malicious"#,
)
}
fn docs_equal(a: &oatf::types::Document, b: &oatf::types::Document) -> bool {
let a_json = serde_json::to_value(a).unwrap();
let b_json = serde_json::to_value(b).unwrap();
a_json == b_json
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(64))]
#[test]
fn single_phase_idempotent(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_single_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let n1 = normalize(doc);
let n2 = normalize(n1.clone());
prop_assert!(docs_equal(&n1, &n2),
"normalize not idempotent for single-phase form with mode={}", mode);
}
#[test]
fn multi_phase_idempotent(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_multi_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let n1 = normalize(doc);
let n2 = normalize(n1.clone());
prop_assert!(docs_equal(&n1, &n2),
"normalize not idempotent for multi-phase form with mode={}", mode);
}
#[test]
fn multi_actor_idempotent(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_multi_actor_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let n1 = normalize(doc);
let n2 = normalize(n1.clone());
prop_assert!(docs_equal(&n1, &n2),
"normalize not idempotent for multi-actor form with mode={}", mode);
}
#[test]
fn single_phase_becomes_actors(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_single_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
prop_assert!(normalized.attack.execution.actors.is_some(),
"single-phase should normalize to actors form");
prop_assert!(normalized.attack.execution.state.is_none(),
"single-phase state should be cleared after normalization");
}
#[test]
fn multi_phase_becomes_actors(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_multi_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
prop_assert!(normalized.attack.execution.actors.is_some(),
"multi-phase should normalize to actors form");
prop_assert!(normalized.attack.execution.phases.is_none(),
"multi-phase phases should be cleared after normalization");
}
#[test]
fn defaults_applied(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_multi_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
prop_assert_eq!(normalized.attack.name.as_deref(), Some("Untitled"));
prop_assert_eq!(normalized.attack.version, Some(1));
prop_assert_eq!(normalized.attack.status, Some(oatf::enums::Status::Draft));
}
#[test]
fn severity_expanded(
mode in arb_mode(),
severity in arb_severity(),
surface in arb_surface(),
) {
let yaml = build_doc_with_indicators(&mode, &surface, &severity);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
if let Some(ref sev) = normalized.attack.severity {
match sev {
oatf::types::Severity::Object { confidence, .. } => {
prop_assert_eq!(*confidence, Some(50),
"severity.confidence should default to 50");
}
oatf::types::Severity::Scalar(_) => {
prop_assert!(false, "severity should be expanded to object form");
}
}
}
}
#[test]
fn indicator_ids_generated(
mode in arb_mode(),
surface in arb_surface(),
severity in arb_severity(),
) {
let yaml = build_doc_with_indicators(&mode, &surface, &severity);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
if let Some(indicators) = &normalized.attack.indicators {
for ind in indicators {
prop_assert!(ind.id.is_some(), "indicator should have auto-generated id");
}
}
}
#[test]
fn normalized_is_reparseable(
mode in arb_mode(),
tool_name in "[a-z]{2,8}",
) {
let yaml = build_multi_phase_doc(&mode, &tool_name);
let doc = parse(&yaml).expect("parse should succeed");
let normalized = normalize(doc);
let serialized = serialize(&normalized).expect("serialize should succeed");
let reparsed = parse(&serialized);
prop_assert!(reparsed.is_ok(),
"re-parsing normalized document failed: {:?}", reparsed.err());
}
}