From b841a4be9dd36e3bf0be5eb1b99710cea8ee5f10 Mon Sep 17 00:00:00 2001 From: Matt Adams Date: Wed, 15 Jul 2026 10:24:57 +0100 Subject: [PATCH 1/3] feat: let the Assistant refine the Detection & Response narrative (#44) The AttackGen Assistant could only refine the generated scenario, not the purple-team Detection & Response narrative introduced in v0.13. Hand the narrative across in the cross-page handoff and add a target selector so the chat can edit either artifact. - scenario_page: _persist_and_render stashes last_defense_narrative (None when there's no narrative, so a stale one can't survive a plain-scenario regen). - Assistant page: 'Editing: Scenario / Detection & Response' radio (shown only when a narrative exists) routes the artifact, system prompt, and LangSmith trace tags; per-target chat histories avoid cross-contamination. Editing the narrative passes the scenario as read-only reference context. Co-Authored-By: Claude Opus 4.8 --- core/scenario_page.py | 8 +- ...4_\360\237\222\254_AttackGen_Assistant.py" | 209 +++++++++++------- tests/test_scenario_page.py | 6 + 3 files changed, 146 insertions(+), 77 deletions(-) diff --git a/core/scenario_page.py b/core/scenario_page.py index 5f6003c..9c7d47b 100644 --- a/core/scenario_page.py +++ b/core/scenario_page.py @@ -346,9 +346,15 @@ def _persist_and_render( st.session_state[layer_key] = layer_payload st.session_state[filename_key] = download_name st.session_state[defense_key] = defense_state - # Cross-page handoff for the AttackGen Assistant chat page. + # Cross-page handoff for the AttackGen Assistant chat page. The defense + # narrative rides along so the Assistant can refine it too; set it + # unconditionally (None when there's no narrative) so a stale one from an + # earlier generation can't linger after a plain-scenario regen. st.session_state["last_scenario"] = True st.session_state["last_scenario_text"] = cleaned + st.session_state["last_defense_narrative"] = ( + defense_state.get("narrative_md") if defense_state else None + ) _render_result( page_id=page_id, diff --git "a/pages/4_\360\237\222\254_AttackGen_Assistant.py" "b/pages/4_\360\237\222\254_AttackGen_Assistant.py" index 0f3b22b..6bd3763 100644 --- "a/pages/4_\360\237\222\254_AttackGen_Assistant.py" +++ "b/pages/4_\360\237\222\254_AttackGen_Assistant.py" @@ -11,13 +11,23 @@ restore_from_query_params() -SYSTEM_PROMPT = ( +SCENARIO_SYSTEM_PROMPT = ( "You are an AI assistant that helps users update and ask questions about their incident " "response scenario. Only respond to questions or requests relating to the scenario, or " "incident response testing in general. Format your responses using proper Markdown syntax " "with headers, bullet points, and formatting for readability." ) +DEFENSE_SYSTEM_PROMPT = ( + "You are an AI assistant that helps users refine the purple-team Detection & Response " + "narrative that accompanies their incident response scenario. The narrative walks the " + "scenario from the defender's side β€” detection opportunities, log sources, and response " + "actions, stage by stage. Only respond to questions or requests relating to the detection " + "and response of this scenario, or purple-team testing in general. Keep your suggestions " + "grounded in the scenario provided for reference. Format your responses using proper " + "Markdown syntax with headers, bullet points, and formatting for readability." +) + st.set_page_config(page_title="AttackGen Assistant", page_icon=":speech_balloon:") inject_emoji_fonts() @@ -25,87 +35,134 @@ st.markdown("# AttackGen AssistantπŸ’¬", unsafe_allow_html=True) -if 'last_scenario_text' in st.session_state and st.session_state.get('last_scenario'): - input_scenario = st.session_state['last_scenario_text'] - with st.expander("Generated Scenario"): - with st.container(height=400, border=True): - st.markdown(input_scenario) +scenario_text = ( + st.session_state["last_scenario_text"] + if st.session_state.get("last_scenario") and "last_scenario_text" in st.session_state + else None +) +defense_narrative = st.session_state.get("last_defense_narrative") + +if not scenario_text: + st.info("No scenario found. Please generate a scenario first.") + st.stop() + + +# Pick the artifact to edit. The Detection & Response option only appears when a +# purple-team narrative was generated alongside the scenario (page 1/2 toggle). +if defense_narrative: + choice = st.radio( + "Editing:", + ["Scenario", "Detection & Response"], + horizontal=True, + key="assistant_target", + ) +else: + choice = "Scenario" +target = "defense" if choice == "Detection & Response" else "scenario" + +if target == "defense": + artifact = defense_narrative + system_prompt = DEFENSE_SYSTEM_PROMPT + expander_label = "Detection & Response Narrative" + greeting = "Hi, I can help you refine the Detection & Response narrative for your scenario." + trace_name = "AttackGen Assistant β€” Detection & Response" + trace_tags = ("assistant", "purple_team_narrative") +else: + artifact = scenario_text + system_prompt = SCENARIO_SYSTEM_PROMPT + expander_label = "Generated Scenario" + greeting = "Hi, I can help you update and ask questions about your incident response scenario." + trace_name = "AttackGen Assistant" + trace_tags = ("assistant",) + + +with st.expander(expander_label): + with st.container(height=400, border=True): + st.markdown(artifact) + +chat_container = st.empty() + +# Keep a separate history per target so switching between the scenario and the +# Detection & Response narrative doesn't feed one artifact's chat into the other. +messages_key = f"assistant_messages_{target}" +if messages_key not in st.session_state: + st.session_state[messages_key] = [{"role": "assistant", "content": greeting}] + +with chat_container: + for message in st.session_state[messages_key]: + with st.chat_message(message["role"]): + st.markdown(message["content"]) + + +def generate_response(user_input, chat_history): + if target == "defense": + context = ( + f"Here is the scenario, for reference:\n\n{scenario_text}\n\n" + f"Here is the current Detection & Response narrative the user wants to refine:" + f"\n\n{artifact}\n\n" + f"Chat history:\n{chat_history}\n\nUser: {user_input}" + ) + else: + context = ( + f"Here is the scenario that the user previously generated:\n\n{artifact}\n\n" + f"Chat history:\n{chat_history}\n\nUser: {user_input}" + ) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": context}, + ] + config = LLMConfig.from_session_state( + trace_name=trace_name, + trace_tags=trace_tags, + ) + raw_chunks: list[str] = [] + + def _tee(chunks): + for chunk in chunks: + raw_chunks.append(chunk) + yield chunk + + try: + yield from stream_filter_thinking(_tee(call_llm_stream(config, messages))) + except Exception as e: + yield f"\n\nAn error occurred while calling the model: {e}" + st.session_state["_last_assistant_cleaned"] = ( + f"An error occurred while calling the model: {e}" + ) + return + + raw = "".join(raw_chunks) + thinking, cleaned = clean_model_response(raw) + if thinking: + with st.expander("View Model's Reasoning"): + st.markdown(thinking) + st.session_state["_last_assistant_cleaned"] = cleaned - chat_container = st.empty() - if 'messages' not in st.session_state: - st.session_state.messages = [ - {"role": "assistant", "content": "Hi, I can help you update and ask questions about your incident response scenario."} - ] +if prompt := st.chat_input("Type your message here..."): + st.session_state[messages_key].append({"role": "user", "content": prompt}) + with st.chat_message("user"): + st.markdown(prompt) - with chat_container: - for message in st.session_state.messages: - with st.chat_message(message["role"]): - st.markdown(message["content"]) - - def generate_response(user_input, chat_history): - messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, - { - "role": "user", - "content": ( - f"Here is the scenario that the user previously generated:\n\n{input_scenario}\n\n" - f"Chat history:\n{chat_history}\n\nUser: {user_input}" - ), - }, - ] - config = LLMConfig.from_session_state( - trace_name="AttackGen Assistant", - trace_tags=("assistant",), + with st.chat_message("assistant"): + history = "\n".join( + f"{m['role']}: {m['content']}" for m in st.session_state[messages_key][:-1] ) - raw_chunks: list[str] = [] - - def _tee(chunks): - for chunk in chunks: - raw_chunks.append(chunk) - yield chunk - - try: - yield from stream_filter_thinking(_tee(call_llm_stream(config, messages))) - except Exception as e: - yield f"\n\nAn error occurred while calling the model: {e}" - st.session_state["_last_assistant_cleaned"] = ( - f"An error occurred while calling the model: {e}" - ) - return - - raw = "".join(raw_chunks) - thinking, cleaned = clean_model_response(raw) - if thinking: - with st.expander("View Model's Reasoning"): - st.markdown(thinking) - st.session_state["_last_assistant_cleaned"] = cleaned - - if prompt := st.chat_input("Type your message here..."): - st.session_state.messages.append({"role": "user", "content": prompt}) - with st.chat_message("user"): - st.markdown(prompt) + st.write_stream(generate_response(prompt, history)) - with st.chat_message("assistant"): - history = "\n".join(f"{m['role']}: {m['content']}" for m in st.session_state.messages[:-1]) - st.write_stream(generate_response(prompt, history)) + st.session_state[messages_key].append( + {"role": "assistant", "content": st.session_state.pop("_last_assistant_cleaned", "")} + ) - st.session_state.messages.append( - {"role": "assistant", "content": st.session_state.pop("_last_assistant_cleaned", "")} - ) - def clear_conversation(): - st.session_state.messages = [ - {"role": "assistant", "content": "Hi, I can help you update and ask questions about your incident response scenario."} - ] - chat_container.empty() - with chat_container: - with st.chat_message("assistant"): - st.markdown(st.session_state.messages[0]["content"]) +def clear_conversation(): + st.session_state[messages_key] = [{"role": "assistant", "content": greeting}] + chat_container.empty() + with chat_container: + with st.chat_message("assistant"): + st.markdown(st.session_state[messages_key][0]["content"]) - with st.container(): - if st.button("Clear Conversation", key='clear_button'): - clear_conversation() -else: - st.info("No scenario found. Please generate a scenario first.") +with st.container(): + if st.button("Clear Conversation", key='clear_button'): + clear_conversation() diff --git a/tests/test_scenario_page.py b/tests/test_scenario_page.py index afba353..816621c 100644 --- a/tests/test_scenario_page.py +++ b/tests/test_scenario_page.py @@ -203,6 +203,8 @@ def test_happy_path_calls_llm_cleans_response_and_persists( # Cross-page handoff for the Assistant page. assert fake_session_state["last_scenario"] is True assert fake_session_state["last_scenario_text"] == cleaned + # No defense companion here, so nothing for the Assistant to refine there. + assert fake_session_state["last_defense_narrative"] is None # The artifact flag is set. assert fake_session_state["threat_group_scenario_generated"] is True @@ -520,6 +522,8 @@ def test_defense_persisted_and_offered_for_download( state = fake_session_state["threat_group_scenario_defense"] assert state["narrative_md"] is None assert "Command and Scripting Interpreter (T1059)" in state["deterministic_md"] + # Deterministic-only: no narrative for the Assistant to refine. + assert fake_session_state["last_defense_narrative"] is None md_name = fake_session_state["threat_group_scenario_filename"] detection_downloads = [ @@ -562,6 +566,8 @@ def test_defense_narrative_makes_second_llm_call_and_persists( assert state["narrative_md"] and "Detection walkthrough" in state["narrative_md"] # The combined download carries both the narrative and the reference section. assert "Detection & Response Reference" in state["download_md"] + # The narrative is handed to the Assistant so it can be refined there too. + assert fake_session_state["last_defense_narrative"] == state["narrative_md"] def test_no_defense_download_when_build_defense_returns_none( From fa95d822a318e1b0693591576c8754689e5a7be9 Mon Sep 17 00:00:00 2001 From: Matt Adams Date: Thu, 16 Jul 2026 09:18:27 +0100 Subject: [PATCH 2/3] feat: add combined Scenario + Detection & Response editing mode to the Assistant The Assistant could refine the scenario or the Detection & Response narrative, but not both together, so a cross-cutting change (industry, threat actor, technique, timeline) had to be made twice and could drift between the two. Add a third 'Scenario + Detection & Response' target that shows both artifacts, feeds both into the model context, and instructs the model to apply changes consistently across them. Each target keeps its own chat history. Verified end-to-end in the browser: a single 'education -> hospital' instruction in combined mode rewrote both the scenario and the narrative coherently. Co-Authored-By: Claude Opus 4.8 --- ...4_\360\237\222\254_AttackGen_Assistant.py" | 63 ++++++++++++++----- 1 file changed, 49 insertions(+), 14 deletions(-) diff --git "a/pages/4_\360\237\222\254_AttackGen_Assistant.py" "b/pages/4_\360\237\222\254_AttackGen_Assistant.py" index 6bd3763..dd4c53d 100644 --- "a/pages/4_\360\237\222\254_AttackGen_Assistant.py" +++ "b/pages/4_\360\237\222\254_AttackGen_Assistant.py" @@ -28,6 +28,17 @@ "Markdown syntax with headers, bullet points, and formatting for readability." ) +BOTH_SYSTEM_PROMPT = ( + "You are an AI assistant that helps users refine an incident response scenario and its " + "accompanying purple-team Detection & Response narrative together. When a requested change " + "affects both β€” a different threat actor, industry, technique, or timeline β€” apply it " + "consistently across the two so the attacker's scenario and the defender's walkthrough stay " + "aligned, and make clear which output each part of your response applies to. Only respond to " + "questions or requests relating to the scenario, its detection and response, or incident " + "response testing in general. Format your responses using proper Markdown syntax with " + "headers, bullet points, and formatting for readability." +) + st.set_page_config(page_title="AttackGen Assistant", page_icon=":speech_balloon:") inject_emoji_fonts() @@ -47,38 +58,54 @@ st.stop() -# Pick the artifact to edit. The Detection & Response option only appears when a -# purple-team narrative was generated alongside the scenario (page 1/2 toggle). +# Pick what to edit. The Detection & Response and combined options only appear +# when a purple-team narrative was generated alongside the scenario (page 1/2 +# toggle). The combined option refines both together so a change made to one can +# be carried consistently into the other. if defense_narrative: choice = st.radio( "Editing:", - ["Scenario", "Detection & Response"], + ["Scenario", "Detection & Response", "Scenario + Detection & Response"], horizontal=True, key="assistant_target", ) else: choice = "Scenario" -target = "defense" if choice == "Detection & Response" else "scenario" +target = { + "Detection & Response": "defense", + "Scenario + Detection & Response": "both", +}.get(choice, "scenario") if target == "defense": - artifact = defense_narrative + panels = [("Detection & Response Narrative", defense_narrative)] system_prompt = DEFENSE_SYSTEM_PROMPT - expander_label = "Detection & Response Narrative" greeting = "Hi, I can help you refine the Detection & Response narrative for your scenario." trace_name = "AttackGen Assistant β€” Detection & Response" trace_tags = ("assistant", "purple_team_narrative") +elif target == "both": + panels = [ + ("Generated Scenario", scenario_text), + ("Detection & Response Narrative", defense_narrative), + ] + system_prompt = BOTH_SYSTEM_PROMPT + greeting = ( + "Hi, I can help you refine the scenario and its Detection & Response narrative " + "together, keeping changes consistent across both." + ) + trace_name = "AttackGen Assistant β€” Scenario + Detection & Response" + trace_tags = ("assistant", "purple_team_narrative") else: - artifact = scenario_text + panels = [("Generated Scenario", scenario_text)] system_prompt = SCENARIO_SYSTEM_PROMPT - expander_label = "Generated Scenario" greeting = "Hi, I can help you update and ask questions about your incident response scenario." trace_name = "AttackGen Assistant" trace_tags = ("assistant",) -with st.expander(expander_label): - with st.container(height=400, border=True): - st.markdown(artifact) +for label, content in panels: + with st.expander(label): + with st.container(height=400, border=True): + st.markdown(content) chat_container = st.empty() @@ -95,16 +122,24 @@ def generate_response(user_input, chat_history): - if target == "defense": + if target == "both": + context = ( + f"Here is the incident response scenario:\n\n{scenario_text}\n\n" + f"Here is the accompanying Detection & Response narrative:\n\n{defense_narrative}\n\n" + f"The user wants to refine both together. When a requested change affects both, " + f"apply it consistently across them and show the update to each.\n\n" + f"Chat history:\n{chat_history}\n\nUser: {user_input}" + ) + elif target == "defense": context = ( f"Here is the scenario, for reference:\n\n{scenario_text}\n\n" f"Here is the current Detection & Response narrative the user wants to refine:" - f"\n\n{artifact}\n\n" + f"\n\n{defense_narrative}\n\n" f"Chat history:\n{chat_history}\n\nUser: {user_input}" ) else: context = ( - f"Here is the scenario that the user previously generated:\n\n{artifact}\n\n" + f"Here is the scenario that the user previously generated:\n\n{scenario_text}\n\n" f"Chat history:\n{chat_history}\n\nUser: {user_input}" ) messages = [ From f3f60c295f4e144119ba1aaedc2d536acfbbbe46 Mon Sep 17 00:00:00 2001 From: Matt Adams Date: Thu, 16 Jul 2026 09:40:09 +0100 Subject: [PATCH 3/3] docs: add v0.14 (unreleased) release notes for the Assistant purple-team editing Covers the Assistant editing the Detection & Response narrative (#44) and the new combined 'Scenario + Detection & Response' editing mode. Co-Authored-By: Claude Opus 4.8 --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index 1140c14..7ea2a34 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,11 @@ If you find AttackGen useful, please consider starring the repository on GitHub. ## Releases +### v0.14 (unreleased) +| What's new? | Why is it useful? | +| ----------- | ----------------- | +| Assistant Refines the Detection & Response Output | - Purple-Team Editing: The AttackGen Assistant can now refine the **Detection & Response** narrative, not just the scenario. When a scenario is generated with the 🟣 Purple-team narrative toggle (Threat Group / Custom pages), an **Editing:** selector appears on the Assistant page so the chat can target either the scenario or the defender's walkthrough β€” previously the Assistant only had the scenario text.

- Combined Editing Mode: A third **Scenario + Detection & Response** target refines both together, so a cross-cutting change β€” a different industry, threat actor, technique or timeline β€” is applied consistently across the attacker's scenario and the defender's narrative rather than being made twice and drifting apart.

- Deterministic Facts Untouched: Only the LLM narrative is a refinement target; the local STIX detection join is never LLM-edited. Each editing mode keeps its own chat history, and the combined pass is tagged `purple_team_narrative` in LangSmith. | + ### v0.13.1 | What's new? | Why is it useful? | | ----------- | ----------------- |