AI Security — Threat Modeling for Agentic AI
Best time to fix security issue is during system design
Series: AI Security Mini Series (2 articles)
- 2AI Security — Threat Modeling for Agentic AI This article
Enterprise Architecture | AI, Cloud & DevOps | Speaker | Mentor | Tech Evangelist
We identified all the places where an attacker could potentially interact with our system—prompts, tools, memory, external APIs, and knowledge sources.
But identifying an attack surface answers only one question:
📍 Where can an attacker interact with my system?
The next question is even more important:
🤔 What could actually go wrong?
That's exactly what Threat Modeling helps us answer.
Unlike penetration testing or code reviews, Threat Modeling is a design exercise. It helps us identify potential risks before writing code, making it one of the most cost-effective security practices.
Let's understand this using a simple example.
🏡 Think of Building a House
Imagine you're building your dream home.
Before construction starts, an architect prepares the blueprint.
It shows:
- 🚪 Doors
- 🪟 Windows
- 🚗 Garage
- 🛋️ Living Room
- 🍳 Kitchen
- 🌳 Backyard
Now imagine inviting a security expert to review that blueprint.
Do they inspect the quality of the bricks?
❌ No.
Do they check whether the walls have been painted?
❌ No.
Instead, they ask questions like:
- Is the front door secure?
- Can someone enter through the garage?
- Are valuables visible through the windows?
- Is there another way into the house?
Notice something interesting.
👉 The house hasn't even been built yet.
Yet, the expert can already identify security risks simply by understanding the design.
That's exactly how Threat Modeling works.
Instead of reviewing a house, we're reviewing the architecture of an application.
❓ The Four Questions of Threat Modeling
Adam Shostack, author of Threat Modeling: Designing for Security, simplifies the entire process into four simple questions.
These questions form the foundation of almost every threat modeling exercise.
1️⃣ What are we building?
Before thinking about security, we first need to understand the system.
For our house, we'd ask:
- 🚪 Where are the doors?
- 🪟 Where are the windows?
- 🚗 Is there a garage?
- 👨👩👧 Who lives here?
- 👥 Who visits?
For an AI application, the questions become:
- 👤 Who are the users?
- 🤖 Which AI model are we using?
- 🛠️ Which tools can the agent access?
- 🧠 Where is memory stored?
- 🌐 Which external APIs are called?
💡 Key Idea
You can't secure something you don't understand.
Understanding the architecture always comes first.
2️⃣ What can go wrong?
Now it's time to think like an attacker.
For the house:
- 🚪 Can someone pretend to be a family member?
- 🪟 Can someone climb through a window?
- 💎 Can valuables be stolen?
- 🔥 Can someone damage the property?
We're not solving problems yet.
We're simply identifying possible threats.
For an AI application, the questions become:
- Can someone manipulate prompts?
- Can sensitive information leak?
- Can external tools be misused?
- Can an attacker overwhelm the application?
- Can the AI be tricked into making unsafe decisions?
💡 Threat modeling is simply asking:
"If I were the attacker, where would I start?"
3️⃣ What are we going to do about it?
Once we know the risks, we choose appropriate security controls.
🏡 House examples:
Threat | Mitigation |
|---|---|
Someone enters through the front door | 🔐 Strong lock, Smart Doorbell, CCTV |
Someone climbs through a window | 🪟 Window locks, Alarm sensors |
Someone steals valuables | 🏠 Safe, Security System |
🤖 AI examples:
- 🔑 Authentication
- 👤 Authorization
- 🛡️ Guardrails
- 🔒 Encryption
- ✋ Human Approval
- 📊 Monitoring
- ⏱️ Rate Limiting
Security controls should always be driven by identified threats, not guesswork.
4️⃣ Did we do a good enough job?
No system is ever 100% secure.
Threat modeling helps us make better design decisions, not perfect ones.
For example, you probably wouldn't install fingerprint scanners on every bedroom door.
Instead, you'd focus on protecting the main entrance, because that's where the highest risk exists.
The same applies to software.
Focus on reducing the most important risks first.
🚧 Trust Boundaries
One of the most important concepts in threat modeling is the Trust Boundary.
Let's go back to our house.
Inside the house, you trust your family.
Outside the house, you trust strangers much less.
What separates those two worlds?
🚪 The Front Door
The front door is a Trust Boundary.
Crossing it changes the level of trust.
Other examples include:
- 🚪 Front Gate
- 🚗 Garage Door
- 🌳 Backyard Fence
- 🔐 Personal Safe
Every one of these separates areas with different levels of trust.
🤖 Trust Boundaries in AI Applications
Applications also have trust boundaries.
For example:
text
👤 User│▼🤖 AI Application│▼🧠 LLM│▼🌐 Internet Search
Every transition changes the trust level.
Questions we should ask include:
- ❓ Should we trust user input?
- ❓ Should we trust the LLM's response?
- ❓ Should we trust Internet search results?
- ❓ Should we trust external APIs?
💡 Key Idea
Every time data crosses a Trust Boundary, ask yourself:
"What could go wrong here?"
This simple question uncovers many security risks before the application is even built.
📊 Data Flow Diagrams (DFDs)
Now that we've identified Trust Boundaries, the next step is to understand how data moves across them.
This is where Data Flow Diagrams (DFDs) come in.
Think about receiving a package at your house.
text
📦 Delivery Person│▼🚪 Front Gate│▼🚪 Front Door│▼🏡 Living Room
The diagram isn't showing the structure of your house.
It's showing how something moves.
A DFD does exactly the same thing for software.
It answers questions like:
- 📍 Where does data originate?
- 📍 Where does it travel?
- 📍 Where is it stored?
- 📍 Which systems receive it?
For an AI application, a simple DFD might look like:
text
👤 User│▼🤖 AI Research Assistant│├────────► 🧠 LLM│├────────► 🌐 Web Search│└────────► 💾 Memory
💡 Key Idea
Every arrow represents a potential attack path.
The combination of DFDs and Trust Boundaries gives architects a clear picture of where threats may exist.
🧰 Threat Modeling Frameworks
The Four Questions tell us how to think.
The following frameworks help us answer those questions consistently.
Think of them as tools in a security architect's toolbox.
🛡️ STRIDE
After understanding the architecture, we need a systematic way to identify threats.
That's where STRIDE comes in.
Threat | 🏡 House Example |
|---|---|
S – Spoofing | Someone pretends to be a family member. |
T – Tampering | Someone changes or damages your property. |
R – Repudiation | A visitor denies entering because there's no evidence. |
I – Information Disclosure | Someone sees valuables through a window. |
D – Denial of Service | Someone blocks your driveway. |
E – Elevation of Privilege | A guest gains access to your personal safe. |
Now replace the house with your AI application.
Ask the same questions.
- Can someone impersonate a user?
- Can prompts or responses be modified?
- Can sensitive data leak?
- Can the application be overwhelmed?
- Can users gain permissions they shouldn't have?
💡 Key Idea
STRIDE gives us a structured checklist to answer "What can go wrong?"
🤖 OWASP LLM Top 10
Traditional security frameworks weren't designed for Large Language Models.
AI introduces new risks that require additional thinking.
Some examples include:
- Prompt Injection
- Sensitive Information Disclosure
- Excessive Agency
- Supply Chain Risks
- Training Data Poisoning
Rather than replacing STRIDE, the OWASP LLM Top 10 extends threat modeling specifically for AI applications.
Whenever you ask:
"What can go wrong?"
remember to think about both traditional software threats and AI-specific threats.
🎯 MITRE ATLAS
If STRIDE tells us what the threats are...
MITRE ATLAS helps us understand how attackers perform them.
It documents real-world attack techniques against AI systems, including:
- Model Evasion
- Prompt Manipulation
- Data Poisoning
- Model Theft
- Resource Exhaustion
💡 Think of it as an attacker playbook for AI.
It helps architects think beyond theory and consider realistic attack scenarios.
📋 NIST AI Risk Management Framework
After identifying threats and selecting mitigations, one final question remains.
"Did we do a good enough job?"
This is where governance becomes important.
The NIST AI Risk Management Framework (AI RMF) helps organizations manage AI risks throughout the system lifecycle.
Instead of focusing on individual threats, it asks broader questions:
- Are risks continuously monitored?
- Are responsibilities clearly defined?
- Are governance processes in place?
- Are AI systems trustworthy?
- Are privacy and security considered throughout the lifecycle?
💡 Key Idea
Think of STRIDE, OWASP, and MITRE ATLAS as helping us identify risks.
NIST AI RMF helps organizations manage those risks over time.
🔄 Bringing It All Together
Let's revisit Adam Shostack's Four Questions.
Question | Supporting Tools |
|---|---|
🏗️ What are we building? | Architecture Docs, C4 Diagrams, DFDs |
⚠️ What can go wrong? | STRIDE, OWASP LLM Top 10, MITRE ATLAS |
🛡️ What are we going to do about it? | Security Controls, Design Decisions, Mitigations |
✅ Did we do a good enough job? | NIST AI RMF, Reviews, Governance |
Notice something interesting.
None of these frameworks compete with each other.
They work together to help us think systematically about security.
🤖 Can We Automate Threat Modeling?
By now, you may have noticed something.
Threat modeling follows a repeatable process.
- 📖 Understand the architecture.
- 🏗️ Build a system model.
- 📊 Generate a Data Flow Diagram.
- 🚧 Identify Trust Boundaries.
- 🛡️ Apply STRIDE.
- 📋 Produce a Threat Register.
These are structured activities.
Which naturally leads to an exciting question.
Can an AI agent perform these steps and generate a draft threat model?
The answer is Yes.
But with an important caveat.
AI should assist architects, not replace them.
An AI agent can quickly analyze architecture artifacts, identify potential threats, and generate a high-quality first draft.
Architects and security engineers then review, validate, and refine the findings using business context and experience.
🚀 What's Next?
In this blog, we learned how security architects think.
In the next blog, we'll put these concepts into practice by building an AI Threat Modeling Agent using the AWS Strands Agent SDK.
Our agent will:
- 📖 Read architecture documents
- 🗂️ Understand C4 diagrams and ADRs
- 📊 Generate or validate a DFD
- 🚧 Identify Trust Boundaries
- 🛡️ Apply STRIDE
- 📋 Generate a Threat Register
- 📄 Produce a Threat Modeling Report
Instead of replacing architects, it will become an AI Security Architect Assistant, helping teams perform faster and more consistent architecture reviews.
🎯 Key Takeaways
✅ Threat modeling is a design activity, not a code review.
✅ Start by understanding the system before thinking about security.
✅ The Four Questions provide a simple and practical way to approach threat modeling.
✅ Trust Boundaries identify where trust changes across the architecture.
✅ DFDs visualize how data moves through the system.
✅ STRIDE helps identify traditional security threats.
✅ OWASP LLM Top 10 and MITRE ATLAS extend threat modeling for AI systems.
✅ NIST AI RMF helps organizations govern and manage AI risks.
✅ AI can automate much of the threat modeling process while keeping humans responsible for the final decisions.
📚 References
- 📖 Adam Shostack – Threat Modeling: Designing for Security
- 🛡️ Microsoft STRIDE
- 🤖 OWASP LLM Top 10
- 🎯 MITRE ATLAS
- 📋 NIST AI Risk Management Framework
💬 Final Thoughts
Threat modeling isn't about predicting every possible attack.
It's about asking the right questions early, understanding your architecture, and making informed design decisions before the first line of code is written.
As AI systems become increasingly autonomous and interconnected, architecture-driven security will become more important than ever.
In the next blog, we'll see how these concepts can be automated by building an AI Threat Modeling Agent that assists architects throughout the design review process.
Happy learning, and see you in the next post! 👋
Series: AI Security Mini Series (2 articles)
- 2AI Security — Threat Modeling for Agentic AI This article
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article