├── AuthService-example.yaml ├── OrderService-example.yaml ├── README.md ├── TaaC-AI.py ├── requirements.txt ├── src ├── AuthService_2024-03-27_ThreatModelingReport.html ├── design1.png ├── download1.gif ├── execution1.png ├── execution2.png ├── manual1.gif ├── review1.gif ├── template.md └── yaml_generator.png ├── taac_yaml_generator.py └── template.html /AuthService-example.yaml: -------------------------------------------------------------------------------- 1 | Version: '1.0' 2 | Date: 14.11.2023 3 | 4 | # Service Description 5 | Description: 6 | Name: AuthService 7 | Type: Service 8 | Criticality: Tier1 9 | 10 | # Service Functionality 11 | Functionality: Handles user authentication, including login and token generation. 12 | 13 | # Data Processing Details 14 | DataProcessed: 15 | Type: Confidential 16 | DataCategory: Auth 17 | EncryptionAtRest: Yes 18 | 19 | # Components Used 20 | Components: 21 | Internal: 22 | Exist: Yes 23 | Source: Private 24 | Note: Scoped Package Access 25 | External: 26 | Exist: Yes 27 | PackageManager: NPM 28 | 29 | # Pipeline Configuration 30 | Pipeline: 31 | Type: GithubActions 32 | CODEOWNERS: Yes 33 | BranchProtection: Yes 34 | SignCommits: Yes 35 | PinActions: Yes 36 | 37 | # Network Information 38 | Network: 39 | Access: Private 40 | 41 | # Data Flow 42 | dataFlow: # Removed the dash here 43 | - name: UserAuthenticationFlow 44 | description: Authenticates users and issues tokens. 45 | source: UserLoginInterface 46 | EncryptionTransit: Yes 47 | Authentication: 48 | Exist: Yes 49 | Type: JWT 50 | Authorization: read-write 51 | Protocol: HTTPS 52 | Communication: 53 | Type: RESTful API 54 | interactions: 55 | - from: UserLoginInterface 56 | to: AuthService 57 | method: RESTful API 58 | protocol: HTTPS 59 | - from: AuthService 60 | to: UserDatabase 61 | method: Query 62 | protocol: JDBC 63 | servicesInvolved: [UserLoginInterface, AuthService, UserDatabase] 64 | -------------------------------------------------------------------------------- /OrderService-example.yaml: -------------------------------------------------------------------------------- 1 | Version: '1.0' 2 | Date: 14.11.2023 3 | 4 | # Order Processing Service Description 5 | Description: 6 | Name: OrderProcessingService 7 | Type: Service 8 | Criticality: Tier2 9 | 10 | # Service Functionality 11 | Functionality: Processes customer orders and manages the order database. 12 | 13 | # Data Processing Details 14 | DataProcessed: 15 | Type: Internal 16 | DataCategory: PCI 17 | EncryptionAtRest: Yes 18 | 19 | # Components Used by the Service 20 | Components: 21 | Internal: 22 | Exist: Yes 23 | Source: Private 24 | Note: Namespacing 25 | External: 26 | Exist: Yes 27 | PackageManager: Maven 28 | 29 | # Pipeline Configuration 30 | Pipeline: 31 | Type: Jenkins 32 | CODEOWNERS: No 33 | BranchProtection: Yes 34 | SignCommits: No 35 | PinActions: No 36 | 37 | # Network Information 38 | Network: 39 | Access: Public 40 | 41 | # Order Processing Service Data Flow 42 | dataFlow: 43 | - name: OrderProcessingFlow 44 | description: Processes customer orders and updates inventory. 45 | source: OrderSubmissionPortal 46 | EncryptionTransit: Yes 47 | Authentication: 48 | Exist: Yes 49 | Type: API Keys 50 | Authorization: admin 51 | Protocol: HTTPS 52 | Communication: 53 | Type: REST APIs 54 | interactions: 55 | - from: OrderSubmissionPortal 56 | to: OrderProcessingService 57 | method: RESTful API 58 | protocol: HTTPS 59 | - from: OrderProcessingService 60 | to: InventoryDatabase 61 | method: Update 62 | protocol: JDBC 63 | servicesInvolved: [OrderSubmissionPortal, OrderProcessingService, InventoryDatabase] 64 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AI-driven Threat Modeling-as-a-Code (TaaC-AI) 2 | 3 | It leverages AI to analyze service descriptions, identify security threats, visualize data flow and suggest remediations based on the STRIDE framework and OWASP guidelines. The script generates in-depth HTML report and includes feature for manual risk assessment. 4 | 5 | ## Design 6 | 7 | ![design](/src/design1.png) 8 | 9 | ## How It Works 10 | 11 | - **YAML File Processing:** Loads and validates a YAML file containing the service details 12 | - **AI Threat Analysis:** If an OpenAI API key is provided, the script uses AI to generate a comprehensive threat modeling analysis 13 | - **AI Threat Cross validation:** Various LLMs performs a validation results 14 | - **Data Flow Generation:** Automatically generates a visual representation of the data flow within the service 15 | - **Manual Risk Management:** Users can manually add, modify, or cross out risks in the generated report 16 | - **Report Generation:** Produces a detailed HTML report, including both AI-generated and manually added risks 17 | 18 | ## Model supported 19 | 20 | - GPT-3.5 21 | - GPT-4 22 | - Claude 3 Haiku 23 | - Mistral 7b (through ollama) 24 | - ‎Gemini (Planned) 25 | 26 | ## Set API Key 27 | 28 | Set the openai api key as an environment variable OPENAI_KEY and ANTHROPIC_KEY in your operating system 29 | 30 | [OpenAI API](https://openai.com/blog/openai-api) 31 | [Anthropic API](https://docs.anthropic.com/claude/reference/getting-started-with-the-api) 32 | 33 | **Linux/Mac** 34 | 35 | ```bash 36 | export OPENAI_KEY=sk-ApiKeyExample 37 | ``` 38 | 39 | ```bash 40 | export ANTHROPIC_KEY=sk-ant-api03-ApiKeyExample 41 | ``` 42 | 43 | **Windows** 44 | 45 | ```bash 46 | set OPENAI_KEY=sk-ApiKeyExample 47 | ``` 48 | 49 | ```bash 50 | set ANTHROPIC_KEY=sk-ant-api03-ApiKeyExample 51 | ``` 52 | 53 | ## Install dependencies 54 | 55 | ```bash 56 | pip3 install -r requirements.txt 57 | ``` 58 | ## Install Mistral and Ollama 59 | 60 | 1. Download and install Ollama [ollama.ai](https://ollama.ai) 61 | 2. Install Mistral 62 | 63 | ```bash 64 | ollama pull mistral 65 | ``` 66 | 3. Start ollama service (make sure the Ollama desktop app is closed) 67 | 68 | ```bash 69 | ollama serve 70 | ``` 71 | Now you can use mistral as the main model: 72 | 73 | ```bash 74 | python3 TaaC-AI.py --model mistral 75 | ``` 76 | Or for cross-validation 77 | 78 | ```bash 79 | python3 TaaC-AI.py --model claude --cross-validation mistral 80 | ``` 81 | 82 | ## How to Use ❓ 83 | 1. Create a valid service description using [these guidelines](src/template.md) or use ```taac_yaml_generator.py``` that will guide you through the process of generating one 84 | 85 | ![yaml_generator!](src/yaml_generator.png) 86 | 87 | 2. Execute the script (GPT-3.5 is used by default) 88 | 89 | ```bash 90 | python3 TaaC-AI.py 91 | ``` 92 | 93 | Use ```gpt-4```, ```claude``` or ```mistral``` as a model to identify threats by specifying the ```--model``` option 94 | 95 | ```bash 96 | python3 TaaC-AI.py --model gpt-4 97 | ``` 98 | 99 | To perform Threats result validation by another LLM use ```--cross-validation``` option. 100 | 101 | ```bash 102 | python3 TaaC-AI.py --model claude --cross-validation claude 103 | ``` 104 | 105 | 3. Open generate .html report 106 | 4. Review/Edit AI-driven Threat Modeling Analysis table, and for false positives or resolved issues, mark the 'Status' checkbox 107 | 5. Add manually identified threats to the table (optional) 108 | 6. Download the report via the Download Report button 109 | 110 | ## Usage Example 🏁 111 | 112 | 1. Valid service description example 113 | ```yaml 114 | Version: '1.0' 115 | Date: 14.11.2023 116 | 117 | # Authentication Service Description 118 | Description: 119 | Name: AuthService 120 | Type: Service 121 | Criticality: Tier1 122 | 123 | # Service Functionality 124 | Functionality: Handles user authentication, including login and token generation. 125 | 126 | # Data Processing Details 127 | DataProcessed: 128 | Type: Confidential 129 | DataCategory: Auth 130 | EncryptionAtRest: Yes 131 | 132 | # Components Used by the Service 133 | Components: 134 | Internal: 135 | Exist: Yes 136 | Source: Private 137 | Note: Scoped Package Access 138 | External: 139 | Exist: Yes 140 | PackageManager: NPM 141 | 142 | # Pipeline Configuration 143 | Pipeline: 144 | Type: GithubActions 145 | CODEOWNERS: Yes 146 | BranchProtection: Yes 147 | SignCommits: Yes 148 | PinActions: Yes 149 | 150 | # Network Information 151 | Network: 152 | Access: Private 153 | 154 | # Authentication Service Data Flow 155 | dataFlow: # Removed the dash here 156 | - name: UserAuthenticationFlow 157 | description: Authenticates users and issues tokens. 158 | source: UserLoginInterface 159 | EncryptionTransit: Yes 160 | Authentication: 161 | Exist: Yes 162 | Type: JWT 163 | Authorization: read-write 164 | Protocol: HTTPS 165 | Communication: 166 | Type: RESTful API 167 | interactions: 168 | - from: UserLoginInterface 169 | to: AuthService 170 | method: RESTful API 171 | protocol: HTTPS 172 | - from: AuthService 173 | to: UserDatabase 174 | method: Query 175 | protocol: JDBC 176 | servicesInvolved: [UserLoginInterface, AuthService, UserDatabase] 177 | ``` 178 | 2. Script execution 179 | ![execution!](src/execution1.png) 180 | 3. Download and Review the generated [HTML report](src/AuthService_2024-03-27_ThreatModelingReport.html) 181 | 182 | ![review!](src/review1.gif) 183 | 184 | 5. Add Threats manually 185 | 186 | ![manual!](src/manual1.gif) 187 | 188 | 6. Download the latest report 189 | 190 | ![download!](src/download1.gif) 191 | 192 | ## Roadmap 🗓️ 193 | 194 | - ~~Template Design~~ 195 | - ~~Basic Functionality~~ 196 | - ~~GPT-3 Integration~~ 197 | - ~~Report generation~~ 198 | - ~~Manually adding identified threats~~ 199 | - ~~GPT-4 Integration~~ 200 | - ~~Claude Integration~~ 201 | - ~~LLM Cross Validation~~ 202 | - ~~Mistral Integration via Ollama~~ 203 | - ‎Gemini Integration 204 | - Accuracy Comparison 205 | 206 | ## Contact 📧 207 | 208 | All suggestions write to yevhsec1@gmail.com 209 | -------------------------------------------------------------------------------- /TaaC-AI.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import yaml 4 | import json 5 | from openai import OpenAI 6 | from anthropic import Client 7 | from ollama import Client as OllamaClient 8 | import argparse 9 | from datetime import datetime, date 10 | from jinja2 import Environment, FileSystemLoader 11 | 12 | class Config: 13 | OPENAI_KEY = "OPENAI_KEY" 14 | ANTHROPIC_KEY = "ANTHROPIC_KEY" 15 | HTML_OUTPUT_FILE = 'report.html' 16 | TEMPLATE_FILE = 'template.html' 17 | MODEL = 'gpt-3.5-turbo' 18 | CROSS_VALIDATION = False 19 | DEBUG = False 20 | 21 | @staticmethod 22 | def set_output_file(service_name): 23 | current_date = date.today().strftime("%Y-%m-%d") 24 | Config.HTML_OUTPUT_FILE = f"{service_name.replace(' ', '_')}_{current_date}_ThreatModelingReport.html" 25 | 26 | class YAMLDataHandler: 27 | @staticmethod 28 | def load_and_validate_yaml_file(file_path): 29 | try: 30 | with open(file_path, 'r') as file: 31 | data = yaml.safe_load(file) 32 | except FileNotFoundError: 33 | return False, None, f"Error: File '{file_path}' not found." 34 | except yaml.YAMLError as e: 35 | return False, None, f"YAML syntax error in '{file_path}': {e}" 36 | valid, message = YAMLDataHandler.validate_yaml_data(data) 37 | if not valid: 38 | return False, None, message 39 | return True, data, "YAML file is valid." 40 | 41 | @staticmethod 42 | def validate_yaml_data(data): 43 | required_keys = { 44 | 'Version': str, 45 | 'Date': '%d.%m.%Y', 46 | 'Description': {'Name', 'Type', 'Criticality'}, 47 | 'DataProcessed': {'Type', 'DataCategory', 'EncryptionAtRest'}, 48 | 'Components': {'Internal', 'External'}, 49 | 'Pipeline': {'Type', 'CODEOWNERS', 'BranchProtection', 'SignCommits', 'PinActions'}, 50 | 'Network': {'Access'}, 51 | 'dataFlow': list 52 | } 53 | for key, expected in required_keys.items(): 54 | if key not in data: 55 | return False, f"Missing key: '{key}' in YAML data." 56 | if isinstance(expected, set) and not expected.issubset(data[key].keys()): 57 | return False, f"Missing keys in '{key}': {expected - data[key].keys()}" 58 | if isinstance(expected, type) and not isinstance(data[key], expected): 59 | return False, f"'{key}' should be of type {expected.__name__}." 60 | if key == 'Date': 61 | try: 62 | datetime.strptime(data[key], required_keys[key]) 63 | except ValueError: 64 | return False, f"'Date' is not in the correct format (DD.MM.YYYY)." 65 | return True, "YAML data is valid." 66 | 67 | class ThreatModeling: 68 | def __init__(self, service_description, model): 69 | self.service_description = service_description 70 | self.model = model 71 | self.openai_client = OpenAI(api_key=os.getenv(Config.OPENAI_KEY)) 72 | self.anthropic_client = Client(api_key=os.getenv(Config.ANTHROPIC_KEY)) 73 | self.ollama_client = OllamaClient() 74 | 75 | @staticmethod 76 | def convert_data_flow_to_json(data_flows): 77 | nodes = set() 78 | links = [] 79 | for flow in data_flows: 80 | for interaction in flow['interactions']: 81 | nodes.add(interaction['from']) 82 | nodes.add(interaction['to']) 83 | links.append({ 84 | "source": interaction['from'], 85 | "target": interaction['to'], 86 | "type": interaction['method'] 87 | }) 88 | return json.dumps({"nodes": list(map(lambda x: {"id": x}, nodes)), "links": links}) 89 | 90 | def generate_threat_modeling(self): 91 | if self.model in ['gpt-3.5-turbo', 'gpt-4']: 92 | return self.generate_threat_modeling_openai() 93 | elif self.model == 'claude': 94 | return self.generate_threat_modeling_anthropic() 95 | elif self.model == 'mistral': 96 | return self.generate_threat_modeling_ollama() 97 | else: 98 | raise ValueError(f"Unsupported model: {self.model}") 99 | 100 | def generate_threat_modeling_openai(self): 101 | if not self.openai_client.api_key: 102 | return "

OpenAI key was not provided or is incorrect. AI Threat Modeling was not performed.

" 103 | 104 | prompt = f"""Perform a thorough threat modeling analysis for the provided service, utilizing the STRIDE framework, OWASP Top 10 2021, and OWASP Top 10 CI/CD Security Risks guidelines. Return the analysis in JSON format with the following structure: 105 | {{ 106 | "threats": [ 107 | {{ 108 | "title": "Threat Title", 109 | "description": "Detailed threat description.", 110 | "categories": ["STRIDE Category", "OWASP Top 10 2021 Category", "OWASP Top 10 CI/CD Security Risks Category"], 111 | "remediation": "Recommended steps or strategies to mitigate or resolve the threat.", 112 | "validator": "🟢 {self.model}" 113 | }}, 114 | ... 115 | ] 116 | }} 117 | 118 | Service data: 119 | {self.service_description} 120 | """ 121 | 122 | try: 123 | response = self.openai_client.chat.completions.create( 124 | model=self.model, 125 | messages=[ 126 | {"role": "system", "content": "You are a security expert. Provide a threat analysis."}, 127 | {"role": "user", "content": prompt} 128 | ] 129 | ) 130 | response_text = response.choices[0].message.content.strip() 131 | log(f"OpenAI API Response: {response_text}") 132 | 133 | json_start = response_text.find("{") 134 | json_end = response_text.rfind("}") 135 | if json_start != -1 and json_end != -1: 136 | json_content = response_text[json_start:json_end+1] 137 | try: 138 | threat_analysis_json = json.loads(json_content) 139 | return json.dumps(threat_analysis_json) 140 | except json.JSONDecodeError: 141 | log("Failed to parse extracted JSON content.") 142 | else: 143 | log("Failed to extract JSON content from the response.") 144 | 145 | return "[]" 146 | except Exception as e: 147 | log(f"Error generating threat modeling with OpenAI: {str(e)}") 148 | return f"

Error generating threat modeling: {str(e)}

" 149 | 150 | def generate_threat_modeling_anthropic(self): 151 | if not self.anthropic_client.api_key: 152 | return "

Anthropic key was not provided or is incorrect. AI Threat Modeling was not performed.

" 153 | 154 | prompt = f"""Perform a thorough threat modeling analysis for the provided service, utilizing the STRIDE framework, OWASP Top 10 2021, and OWASP Top 10 CI/CD Security Risks guidelines. Return the analysis in JSON format with the following structure: 155 | {{ 156 | "threats": [ 157 | {{ 158 | "title": "Threat Title", 159 | "description": "Detailed threat description.", 160 | "categories": ["STRIDE Category", "OWASP Top 10 2021 Category", "OWASP Top 10 CI/CD Security Risks Category"], 161 | "remediation": "Recommended steps or strategies to mitigate or resolve the threat.", 162 | "validator": "🟢 {self.model}" 163 | }}, 164 | ... 165 | ] 166 | }} 167 | 168 | Service data: 169 | {self.service_description} 170 | """ 171 | 172 | log(f"Anthropic API Request: {prompt}") 173 | 174 | try: 175 | response = self.anthropic_client.messages.create( 176 | max_tokens=2048, 177 | model="claude-3-haiku-20240307", 178 | messages=[ 179 | { 180 | "role": "user", 181 | "content": prompt 182 | } 183 | ] 184 | ) 185 | log(f"Anthropic API Response: {response}") 186 | response_text = response.content[0].text.strip() 187 | log(f"Anthropic API Response Content: {response_text}") 188 | 189 | json_start = response_text.find("{") 190 | json_end = response_text.rfind("}") + 1 191 | if json_start != -1 and json_end != -1: 192 | json_content = response_text[json_start:json_end].strip() 193 | try: 194 | threat_analysis_json = json.loads(json_content) 195 | return json.dumps(threat_analysis_json) 196 | except json.JSONDecodeError: 197 | log("Failed to parse extracted JSON content.") 198 | else: 199 | log("Failed to extract JSON content from the response.") 200 | 201 | return "[]" 202 | 203 | except Exception as e: 204 | log(f"Error generating threat modeling with Anthropic: {str(e)}") 205 | return f"

Error generating threat modeling: {str(e)}

" 206 | 207 | def generate_threat_modeling_ollama(self): 208 | 209 | prompt = f"""Perform a thorough threat modeling analysis for the provided service, utilizing the STRIDE framework, OWASP Top 10 2021, and OWASP Top 10 CI/CD Security Risks guidelines. Return the analysis in JSON format with the following structure: 210 | {{ 211 | "threats": [ 212 | {{ 213 | "title": "Threat Title", 214 | "description": "Detailed threat description.", 215 | "categories": ["STRIDE Category", "OWASP Top 10 2021 Category", "OWASP Top 10 CI/CD Security Risks Category"], 216 | "remediation": "Recommended steps or strategies to mitigate or resolve the threat.", 217 | "validator": "🟢 {self.model}" 218 | }}, 219 | ... 220 | ] 221 | }} 222 | 223 | Service data: 224 | {self.service_description} 225 | """ 226 | 227 | log(f"Ollama API Request: {prompt}") 228 | 229 | try: 230 | response = self.ollama_client.generate( 231 | model="mistral", 232 | prompt=prompt, 233 | format="json", 234 | stream= False, 235 | system="You are a security expert." 236 | ) 237 | log(f"Ollama API Response: {response}") 238 | response_text = response['response'].strip() 239 | 240 | log(f"Ollama API Response Content: {response_text}") 241 | 242 | json_start = response_text.find("{") 243 | json_end = response_text.rfind("}") + 1 244 | if json_start != -1 and json_end != -1: 245 | json_content = response_text[json_start:json_end].strip() 246 | try: 247 | threat_analysis_json = json.loads(json_content) 248 | return json.dumps(threat_analysis_json) 249 | except json.JSONDecodeError: 250 | log("Failed to parse extracted JSON content.") 251 | else: 252 | log("Failed to extract JSON content from the response.") 253 | 254 | return "[]" 255 | 256 | except Exception as e: 257 | log(f"Error generating threat modeling with Ollama: {str(e)}") 258 | return f"

Error generating threat modeling: {str(e)}

" 259 | 260 | @staticmethod 261 | def validate_threats(threats, validation_model, client): 262 | validated_threats = [] 263 | for threat in threats: 264 | prompt = f""" 265 | Please validate the following threat: 266 | {{ 267 | "title": "{threat['title']}", 268 | "description": "{threat['description']}", 269 | "categories": {threat['categories']}, 270 | "remediation": "{threat['remediation']}" 271 | }} 272 | Is this a valid threat? Respond with 'Yes' or 'No'. 273 | """ 274 | if validation_model in ['gpt-3.5-turbo', 'gpt-4']: 275 | response = client.chat.completions.create( 276 | model=validation_model, 277 | messages=[ 278 | {"role": "system", "content": "You are a security expert. Validate the threat."}, 279 | {"role": "user", "content": prompt} 280 | ] 281 | ) 282 | is_valid = response.choices[0].message.content.strip().lower() == 'yes' 283 | elif validation_model == 'claude': 284 | response = client.messages.create( 285 | max_tokens=5, 286 | model="claude-3-haiku-20240307", 287 | messages=[ 288 | { 289 | "role": "user", 290 | "content": prompt 291 | } 292 | ] 293 | ) 294 | log(f"Validation prompt for Claude: {prompt}") 295 | log(f"Validation response from Claude: {response.content}") 296 | is_valid = 'yes' in response.content[0].text.strip().lower() 297 | elif validation_model == 'mistral': 298 | response = client.generate( 299 | model="mistral", 300 | prompt=prompt, 301 | format="json", 302 | stream=False, 303 | system="You are a security expert. Validate the threat." 304 | ) 305 | log(f"Validation prompt for Mistral: {prompt}") 306 | log(f"Validation response from Mistral: {response['response']}") 307 | is_valid = 'yes' in response['response'].lower() 308 | else: 309 | raise ValueError(f"Unsupported validation model: {validation_model}") 310 | 311 | if is_valid: 312 | threat['validator'] = f"{threat['validator']} 🟢 {validation_model}" 313 | else: 314 | threat['validator'] = f"{threat['validator']} 🔴 {validation_model}" 315 | 316 | validated_threats.append(threat) 317 | return validated_threats 318 | 319 | @staticmethod 320 | def remove_duplicate_threats(threats): 321 | unique_threats = [] 322 | seen_titles = set() 323 | for threat in threats: 324 | if threat['title'] not in seen_titles: 325 | unique_threats.append(threat) 326 | seen_titles.add(threat['title']) 327 | return unique_threats 328 | 329 | class HTMLReportRenderer: 330 | def __init__(self, service_info, data_flow_json, threat_analysis_html): 331 | self.service_info = service_info 332 | self.data_flow_json = data_flow_json 333 | self.threat_analysis_html = threat_analysis_html 334 | self.service_name = service_info.get('Description', {}).get('Name', 'Report') 335 | self.current_date = date.today().strftime("%Y-%m-%d") 336 | 337 | def render(self): 338 | env = Environment(loader=FileSystemLoader('.')) 339 | template = env.get_template(Config.TEMPLATE_FILE) 340 | return template.render( 341 | service=self.service_info, 342 | data_flow_json=self.data_flow_json, 343 | threat_analysis_html=self.threat_analysis_html, 344 | service_name=self.service_name, 345 | current_date=self.current_date 346 | ) 347 | 348 | class PrintManager: 349 | TITLE_STYLE = '\033[1;34m' 350 | NORMAL_STYLE = '\033[0m' 351 | NAME_STYLE = '\033[1;32m' 352 | HIGHLIGHT_STYLE = '\033[1;33m' 353 | FILE_STYLE = '\033[1;33m' 354 | ERROR_STYLE = '\033[1;31m' 355 | 356 | @staticmethod 357 | def print_usage(): 358 | print(f"{PrintManager.TITLE_STYLE}AI-driven Threat modeling-as-a-Code (TaaC) v1.1{PrintManager.NORMAL_STYLE}") 359 | print(f"Created by YevhSec1\n") 360 | print(f"{PrintManager.HIGHLIGHT_STYLE}Usage:{PrintManager.NORMAL_STYLE} python3 TaaC.py [options] \n") 361 | print(f"{PrintManager.HIGHLIGHT_STYLE}Options:{PrintManager.NORMAL_STYLE}") 362 | print(" -h, --help show this help message and exit") 363 | print(" --model Select the model version: gpt-3.5-turbo or gpt-4") 364 | print(" --cross-validation Perform cross-validation using two LLMs") 365 | print(" --debug Enable debug logging\n") 366 | print(f"{PrintManager.HIGHLIGHT_STYLE}Arguments:{PrintManager.NORMAL_STYLE}") 367 | print(" yaml_file Path to the YAML file containing the service information.\n") 368 | print(f"{PrintManager.HIGHLIGHT_STYLE}Example:{PrintManager.NORMAL_STYLE}") 369 | print(" python3 TaaC.py auth_service.yaml --model gpt-3.5-turbo --cross-validation --debug") 370 | 371 | @staticmethod 372 | def print_progress(file_name): 373 | print(f"{PrintManager.TITLE_STYLE}Processing: {file_name}{PrintManager.NORMAL_STYLE}") 374 | print("Generating report... Please wait.") 375 | 376 | @staticmethod 377 | def print_completion(): 378 | print(f"\n{PrintManager.NAME_STYLE}Done! The report has been successfully generated.{PrintManager.NORMAL_STYLE}\n") 379 | 380 | @staticmethod 381 | def print_error(message): 382 | print(f"{PrintManager.ERROR_STYLE}Error:{PrintManager.NORMAL_STYLE} {message}") 383 | 384 | def convert_json_to_html(json_data): 385 | try: 386 | threats = json.loads(json_data) 387 | if isinstance(threats, list): 388 | return 'No threats found.' 389 | else: 390 | threats = threats['threats'] 391 | except (json.JSONDecodeError, KeyError) as e: 392 | log(f"Error parsing JSON data: {str(e)}") 393 | threats = [] 394 | 395 | html = '' 396 | for threat in threats: 397 | html += '\n' 398 | html += f'{threat["title"]}\n' 399 | html += f'{threat["validator"]}\n' 400 | html += f'{threat["description"]}\n' 401 | html += f'{", ".join(threat["categories"])}\n' 402 | html += f'{threat["remediation"]}\n' 403 | html += '\n' 404 | html += '' 405 | html += '' 406 | html += '' 407 | html += '\n' 408 | html += '\n' 409 | return html 410 | 411 | def parse_arguments(): 412 | parser = argparse.ArgumentParser(description='Generate a threat modeling report from a YAML file.') 413 | parser.add_argument('yaml_file', help='Path to the YAML file containing the service information.') 414 | parser.add_argument('--model', choices=['gpt-3.5-turbo', 'gpt-4', 'claude', 'mistral'], default='gpt-3.5-turbo', help='Choice of LLM for generating the report.') 415 | parser.add_argument('--cross-validation', choices=['gpt-3.5-turbo', 'gpt-4', 'claude', 'mistral'], help='Perform cross-validation using two LLMs.') 416 | parser.add_argument('--debug', action='store_true', help='Enable debug logging.') 417 | return parser.parse_args() 418 | 419 | def log(message): 420 | if Config.DEBUG: 421 | print(f"[DEBUG] {message}") 422 | 423 | def main(): 424 | args = parse_arguments() 425 | 426 | Config.MODEL = args.model 427 | Config.CROSS_VALIDATION = args.cross_validation 428 | Config.DEBUG = args.debug 429 | 430 | valid, service_info, message = YAMLDataHandler.load_and_validate_yaml_file(args.yaml_file) 431 | if not valid: 432 | PrintManager.print_error(message) 433 | return 434 | 435 | service_name = service_info.get('Description', {}).get('Name', 'Report') 436 | Config.set_output_file(service_name) 437 | 438 | PrintManager.print_progress(args.yaml_file) 439 | 440 | data_flows = service_info.get('dataFlow', []) 441 | data_flow_json = ThreatModeling.convert_data_flow_to_json(data_flows) 442 | 443 | threat_modeling = ThreatModeling(json.dumps(service_info, indent=2), args.model) 444 | threat_analysis_json = threat_modeling.generate_threat_modeling() 445 | log(f"Threat Analysis JSON: {threat_analysis_json}") 446 | 447 | if Config.CROSS_VALIDATION: 448 | validation_model = Config.CROSS_VALIDATION 449 | log(f"Performing cross-validation using {validation_model}") 450 | threat_modeling_validation = ThreatModeling(json.dumps(service_info, indent=2), validation_model) 451 | threats = json.loads(threat_analysis_json)['threats'] 452 | log(f"Threats identified by {args.model}: {len(threats)}") 453 | 454 | validated_threats = ThreatModeling.validate_threats( 455 | threats, 456 | validation_model, 457 | threat_modeling_validation.openai_client if validation_model.startswith('gpt') else threat_modeling_validation.ollama_client if validation_model == 'mistral' else threat_modeling_validation.anthropic_client 458 | ) 459 | log(f"Validated threats: {len(validated_threats)}") 460 | 461 | threat_analysis_json = json.dumps({'threats': validated_threats}) 462 | log(f"Updated threat analysis JSON with validation results: {threat_analysis_json}") 463 | 464 | threat_analysis_html = convert_json_to_html(threat_analysis_json) 465 | 466 | renderer = HTMLReportRenderer(service_info, data_flow_json, threat_analysis_html) 467 | html_report = renderer.render() 468 | 469 | with open(Config.HTML_OUTPUT_FILE, 'w') as file: 470 | file.write(html_report) 471 | PrintManager.print_completion() 472 | 473 | if __name__ == "__main__": 474 | main() 475 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | openai 2 | pyyaml 3 | jinja2 4 | argparse 5 | anthropic 6 | ollama 7 | termcolor 8 | -------------------------------------------------------------------------------- /src/AuthService_2024-03-27_ThreatModelingReport.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | AuthService Report 7 | 8 | 155 | 156 | 157 |
158 |
159 |

AuthService - Threat Modeling Report

160 | 161 |
162 | 163 |
164 | 165 |
166 |
Service Description
167 |
168 |

Name: AuthService

169 |

Type: Service

170 |

Criticality: Tier1

171 |
172 |
173 | 174 | 175 |
176 |
Service Functionality
177 |
178 |

Handles user authentication, including login and token generation.

179 |
180 |
181 | 182 | 183 |
184 |
Data Processing Details
185 |
186 |

Type: Confidential

187 |

Data Category: Auth

188 |

Encryption at Rest: True

189 |
190 |
191 | 192 | 193 |
194 |
Network Information
195 |
196 |

Access: Private

197 |
198 |
199 | 200 | 201 |
202 |
Components
203 |
204 |

Internal Components

205 |

Exist: True

206 |

Source: Private

207 |

External Components

208 |

Exist: True

209 |

Package Manager: NPM

210 |
211 |
212 | 213 | 214 |
215 |
Pipeline Configuration
216 |
217 |

Type: GithubActions

218 |

CODEOWNERS: True

219 |

Branch Protection: True

220 |

Sign Commits: True

221 |

Pin Actions: True

222 |
223 |
224 |
225 | 226 | 227 |
228 |
229 |

Data Flow Information

230 |
231 |
232 |
233 |
234 |
235 | 236 | 237 |
238 |
Threats
239 |
240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 269 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | 290 | 291 | 292 | 293 | 294 | 295 | 296 | 297 | 298 | 299 | 300 |
TitleValidatorDescriptionCategoriesRemediationStatusActions
Injection Attacks🟢 claude 🟢 gpt-4Attackers may attempt to inject malicious code or SQL queries into the application, potentially allowing them to gain unauthorized access to sensitive data or perform other malicious actions.Spoofing, OWASP Top 10 2021: A03 - Injection, OWASP Top 10 CI/CD Security Risks: Injection FlawsImplement input validation and sanitization to prevent injection attacks. Use parameterized queries or prepared statements when interacting with the database. Validate and sanitize all user input before using it in queries, commands, or function calls.
Broken Authentication🟢 claude 🟢 gpt-4Weaknesses in the authentication mechanism, such as weak password policies, lack of multi-factor authentication, or insecure session management, could allow attackers to gain unauthorized access to user accounts.Tampering, OWASP Top 10 2021: A07 - Identification and Authentication Failures, OWASP Top 10 CI/CD Security Risks: Authentication and AuthorizationImplement strong password policies, enforce multi-factor authentication, and follow secure session management practices. Regularly review and update authentication and authorization mechanisms to address any vulnerabilities.
Sensitive Data Exposure🟢 claude 🟢 gpt-4If sensitive data, such as authentication credentials or user information, is not properly protected during transmission or storage, it could be accessed by unauthorized parties.Disclosure of Information, OWASP Top 10 2021: A02 - Cryptographic Failures, OWASP Top 10 CI/CD Security Risks: Secrets ManagementEnsure that all sensitive data is encrypted both in transit and at rest. Use secure communication protocols like HTTPS and follow best practices for storing and handling sensitive information.
Insecure Dependencies🟢 claude 🟢 gpt-4The use of external components or dependencies, such as libraries or frameworks, could introduce vulnerabilities if they are not properly managed and kept up-to-date.Elevation of Privilege, OWASP Top 10 2021: A06 - Vulnerable and Outdated Components, OWASP Top 10 CI/CD Security Risks: DependenciesRegularly review and update all external dependencies to the latest secure versions. Implement a dependency management process to ensure that all components are properly vetted and kept up-to-date.
Insecure CICD Pipelines🟢 claude 🟢 gpt-4Vulnerabilities or misconfigurations in the CICD pipeline could allow attackers to tamper with the build process, inject malicious code, or gain unauthorized access to sensitive information.Denial of Service, OWASP Top 10 2021: A05 - Security Misconfiguration, OWASP Top 10 CI/CD Security Risks: CICD PipelineImplement secure CICD practices, such as enforcing branch protection, using signed commits, pinning actions, and following the principle of least privilege. Regularly review and audit the CICD pipeline to identify and address any security issues.
301 |
302 | 303 |
304 |
305 |
306 |
307 | 308 | 309 | 450 | 451 | 511 | 512 | 520 | 521 | 522 | -------------------------------------------------------------------------------- /src/design1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/design1.png -------------------------------------------------------------------------------- /src/download1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/download1.gif -------------------------------------------------------------------------------- /src/execution1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/execution1.png -------------------------------------------------------------------------------- /src/execution2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/execution2.png -------------------------------------------------------------------------------- /src/manual1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/manual1.gif -------------------------------------------------------------------------------- /src/review1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/review1.gif -------------------------------------------------------------------------------- /src/template.md: -------------------------------------------------------------------------------- 1 | ## Creating a Valid Service Description in YAML Format 2 | 3 | ### Introduction 4 | To ensure successful threat modeling with TaaC, it's crucial to follow a specific structure for your service description in the YAML file. Below is a step-by-step guide on how to structure your YAML file. 5 | 6 | ### Tips: 7 | 8 | - Indentation is crucial in YAML. Use spaces (not tabs) for indentation 9 | - Ensure that keys and nested items are properly aligned 10 | - Use quotes if values contain special characters or spaces 11 | - Validate your YAML file with online tools to catch syntax errors 12 | 13 | ### YAML File Structure 14 | Your YAML file should consist of several key sections: 15 | 16 | - **Version:** Specify the version of the service or the document. 17 | - **Date:** The date when the document was created or last updated, in DD.MM.YYYY format. 18 | - **Description:** Basic details of the service, including its name, type, and criticality. 19 | - **Functionality:** A brief description of what the service does. 20 | - **Data Processed:** Details about the type and category of data processed by the service, and its encryption status at rest. 21 | - **Components:** Information about internal and external components used by the service. 22 | - **Pipeline:** Details of the CI/CD pipeline configuration. 23 | - **Network:** Information about the network access level. 24 | - **Data Flow:** A detailed representation of how data moves within the service, including interactions between different components. 25 | 26 | ### Instructions 27 | #### Version and Date 28 | 29 | ```yaml 30 | Version: '1.0' 31 | Date: 15.01.2024 32 | ``` 33 | 34 | #### Description: Provide the name, type, and criticality of the service 35 | 36 | ```yaml 37 | Description: 38 | Name: ExampleService 39 | Type: Web Application 40 | Criticality: High 41 | ``` 42 | #### Functionality: A brief summary of what the service does 43 | 44 | ```yaml 45 | Functionality: Manages user data and authentication. 46 | ``` 47 | 48 | #### Data Processed: Specify the type of data, its category, and encryption status 49 | 50 | ```yaml 51 | DataProcessed: 52 | Type: User Data 53 | DataCategory: PII 54 | EncryptionAtRest: Yes 55 | ``` 56 | 57 | #### Components: Detail any internal and external components 58 | 59 | ```yaml 60 | Components: 61 | Internal: 62 | Exist: Yes 63 | Source: Internal Repository 64 | External: 65 | Exist: Yes 66 | PackageManager: npm 67 | ``` 68 | 69 | #### Pipeline: Configuration of your CI/CD pipeline 70 | 71 | ```yaml 72 | Pipeline: 73 | Type: Jenkins 74 | CODEOWNERS: Yes 75 | BranchProtection: Yes 76 | SignCommits: Yes 77 | PinActions: No 78 | ``` 79 | 80 | #### Network: Define the network accessibility 81 | 82 | ```yaml 83 | Network: 84 | Access: Restricted 85 | ``` 86 | 87 | #### Data Flow: Illustrate the flow of data within the service, including sources, targets, and methods 88 | 89 | ```yaml 90 | DataFlow: 91 | - name: UserAuthenticationFlow 92 | description: Handles user login and authentication. 93 | interactions: 94 | - from: UserInterface 95 | to: AuthenticationServer 96 | method: HTTPS POST 97 | - from: AuthenticationServer 98 | to: Database 99 | method: SQL Query 100 | ``` 101 | 102 | ### Examples 103 | 104 | 1. **Template.yaml** 105 | ```yaml 106 | Version: '1.0' 107 | Date: 14.11.2023 108 | 109 | # Service Description 110 | Description: 111 | Name: Name1 112 | Type: Service 113 | Criticality: Tier1/Tier2/Tier3 114 | 115 | # Service Functionality 116 | Functionality: # Add a short description of what the service does 117 | 118 | # Data Processing Details 119 | DataProcessed: 120 | Type: Secret/Confidential/Internal/Public 121 | DataCategory: Auth/PCI/PII/etc 122 | EncryptionAtRest: Yes/No 123 | 124 | # Components Used by the Service 125 | Components: 126 | Internal: 127 | Exist: Yes/No 128 | Source: Private/Public 129 | Note: Namespacing/Scoped Package Access/etc 130 | External: 131 | Exist: Yes/No 132 | PackageManager: NPM/Maven/NuGet/RubyGems/etc 133 | 134 | # Pipeline Configuration 135 | Pipeline: 136 | Type: GithubActions/Jenkins/etc 137 | CODEOWNERS: Yes/No 138 | BranchProtection: Yes/No 139 | SignCommits: Yes/No 140 | PinActions: Yes/No 141 | 142 | # Network Information 143 | Network: 144 | Access: Public/Private 145 | 146 | # Data Flow Examples 147 | dataFlow: 148 | - name: UserAuthenticationFlow 149 | description: Handles user login and authentication. 150 | source: UserLoginInterface 151 | EncryptionTransit: Yes 152 | Authentication: 153 | Exist: Yes 154 | Type: JWT 155 | Authorization: read-write 156 | Protocol: HTTPS 157 | Communication: 158 | Type: RESTful API 159 | interactions: 160 | - from: UserLoginInterface 161 | to: AuthenticationService 162 | method: RESTful API 163 | protocol: HTTPS 164 | - from: AuthenticationService 165 | to: Database 166 | method: CredentialVerification 167 | protocol: JDBC/ODBC/DatabaseAPI 168 | servicesInvolved: [UserLoginInterface, AuthenticationService, Database] 169 | 170 | - name: OrderProcessingNotificationFlow 171 | description: Processes orders and sends notifications to users. 172 | source: OrderSubmissionInterface 173 | EncryptionTransit: Yes 174 | Authentication: 175 | Exist: Yes 176 | Type: API Keys 177 | Authorization: admin 178 | Protocol: HTTPS/AMQP/SMTP 179 | Communication: 180 | Type: REST APIs/Message Queues/WebSockets 181 | interactions: 182 | - from: OrderSubmissionInterface 183 | to: OrderProcessingService 184 | method: RESTful API 185 | protocol: HTTPS 186 | - from: OrderProcessingService 187 | to: NotificationService 188 | method: MessageQueue 189 | protocol: AMQP 190 | - from: NotificationService 191 | to: EmailService 192 | method: SMTP 193 | protocol: SMTP 194 | servicesInvolved: [OrderSubmissionInterface, OrderProcessingService, NotificationService, EmailService] 195 | 196 | - name: DataReportingAnalyticsFlow 197 | description: Aggregates data and generates analytics reports. 198 | source: DataCollectionService 199 | EncryptionTransit: Yes 200 | Authentication: 201 | Exist: Yes 202 | Type: OAuth 203 | Authorization: read-write 204 | Protocol: HTTPS/SQL 205 | Communication: 206 | Type: REST APIs/Batch Processing 207 | interactions: 208 | - from: DataCollectionService 209 | to: DataWarehouse 210 | method: BatchUpload 211 | protocol: HTTPS 212 | - from: DataWarehouse 213 | to: AnalyticsService 214 | method: DataQuery 215 | protocol: SQL 216 | - from: AnalyticsService 217 | to: ReportingTool 218 | method: RESTful API 219 | protocol: HTTPS 220 | servicesInvolved: [DataCollectionService, DataWarehouse, AnalyticsService, ReportingTool] 221 | 222 | - name: InventoryManagementFlow 223 | description: Manages inventory levels based on orders and supply chain updates. 224 | source: InventoryUpdateInterface 225 | EncryptionTransit: Yes 226 | Authentication: 227 | Exist: Yes 228 | Type: Client Certificates 229 | Authorization: admin 230 | Protocol: HTTPS 231 | Communication: 232 | Type: REST APIs/Direct Database Access 233 | interactions: 234 | - from: InventoryUpdateInterface 235 | to: InventoryService 236 | method: RESTful API 237 | protocol: HTTPS 238 | - from: InventoryService 239 | to: SupplierService 240 | method: RESTful API 241 | protocol: HTTPS 242 | - from: InventoryService 243 | to: Database 244 | method: UpdateQuery 245 | protocol: JDBC/ODBC/DatabaseAPI 246 | servicesInvolved: [InventoryUpdateInterface, InventoryService, SupplierService, Database] 247 | ``` 248 | 249 | 2. **Auth_service.yaml** 250 | ```yaml 251 | Version: '1.0' 252 | Date: 14.11.2023 253 | 254 | # Authentication Service Description 255 | Description: 256 | Name: AuthService 257 | Type: Service 258 | Criticality: Tier1 259 | 260 | # Service Functionality 261 | Functionality: Handles user authentication, including login and token generation. 262 | 263 | # Data Processing Details 264 | DataProcessed: 265 | Type: Confidential 266 | DataCategory: Auth 267 | EncryptionAtRest: Yes 268 | 269 | # Components Used by the Service 270 | Components: 271 | Internal: 272 | Exist: Yes 273 | Source: Private 274 | Note: Scoped Package Access 275 | External: 276 | Exist: Yes 277 | PackageManager: NPM 278 | 279 | # Pipeline Configuration 280 | Pipeline: 281 | Type: GithubActions 282 | CODEOWNERS: Yes 283 | BranchProtection: Yes 284 | SignCommits: Yes 285 | PinActions: Yes 286 | 287 | # Network Information 288 | Network: 289 | Access: Private 290 | 291 | # Authentication Service Data Flow 292 | dataFlow: # Removed the dash here 293 | - name: UserAuthenticationFlow 294 | description: Authenticates users and issues tokens. 295 | source: UserLoginInterface 296 | EncryptionTransit: Yes 297 | Authentication: 298 | Exist: Yes 299 | Type: JWT 300 | Authorization: read-write 301 | Protocol: HTTPS 302 | Communication: 303 | Type: RESTful API 304 | interactions: 305 | - from: UserLoginInterface 306 | to: AuthService 307 | method: RESTful API 308 | protocol: HTTPS 309 | - from: AuthService 310 | to: UserDatabase 311 | method: Query 312 | protocol: JDBC 313 | servicesInvolved: [UserLoginInterface, AuthService, UserDatabase] 314 | ``` 315 | 3. **Order_service.yaml** 316 | ```yaml 317 | Version: '1.0' 318 | Date: 14.11.2023 319 | 320 | # Order Processing Service Description 321 | Description: 322 | Name: OrderProcessingService 323 | Type: Service 324 | Criticality: Tier2 325 | 326 | # Service Functionality 327 | Functionality: Processes customer orders and manages the order database. 328 | 329 | # Data Processing Details 330 | DataProcessed: 331 | Type: Internal 332 | DataCategory: PCI 333 | EncryptionAtRest: Yes 334 | 335 | # Components Used by the Service 336 | Components: 337 | Internal: 338 | Exist: Yes 339 | Source: Private 340 | Note: Namespacing 341 | External: 342 | Exist: Yes 343 | PackageManager: Maven 344 | 345 | # Pipeline Configuration 346 | Pipeline: 347 | Type: Jenkins 348 | CODEOWNERS: No 349 | BranchProtection: Yes 350 | SignCommits: No 351 | PinActions: No 352 | 353 | # Network Information 354 | Network: 355 | Access: Public 356 | 357 | # Order Processing Service Data Flow 358 | dataFlow: 359 | - name: OrderProcessingFlow 360 | description: Processes customer orders and updates inventory. 361 | source: OrderSubmissionPortal 362 | EncryptionTransit: Yes 363 | Authentication: 364 | Exist: Yes 365 | Type: API Keys 366 | Authorization: admin 367 | Protocol: HTTPS 368 | Communication: 369 | Type: REST APIs 370 | interactions: 371 | - from: OrderSubmissionPortal 372 | to: OrderProcessingService 373 | method: RESTful API 374 | protocol: HTTPS 375 | - from: OrderProcessingService 376 | to: InventoryDatabase 377 | method: Update 378 | protocol: JDBC 379 | servicesInvolved: [OrderSubmissionPortal, OrderProcessingService, InventoryDatabase] 380 | ``` 381 | 382 | -------------------------------------------------------------------------------- /src/yaml_generator.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yevh/TaaC-AI/813677feba064038097a6f4d1ed0cd05b7e2a530/src/yaml_generator.png -------------------------------------------------------------------------------- /taac_yaml_generator.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from datetime import datetime 3 | from termcolor import colored 4 | 5 | class ServiceDescriptionGenerator: 6 | def __init__(self): 7 | self.service_description = {} 8 | 9 | def get_user_input(self, prompt, required=True, validation=None): 10 | while True: 11 | value = input(colored(prompt, 'blue', attrs=['bold'])).strip() 12 | if not value and required: 13 | print(colored("This field is required. Please provide a value.", 'red')) 14 | elif validation and not validation(value): 15 | print(colored("Invalid input. Please try again.", 'red')) 16 | else: 17 | return value 18 | 19 | def create_service_description(self): 20 | print(colored("\n# Version and Date", 'green', attrs=['bold'])) 21 | self.service_description['Version'] = self.get_user_input("Enter the service version (e.g., 1.0): ") 22 | self.service_description['Date'] = datetime.now().strftime("%d.%m.%Y") 23 | 24 | print(colored("\n# Service Description", 'green', attrs=['bold'])) 25 | self.service_description['Description'] = { 26 | 'Name': self.get_user_input("Enter the service name (e.g., AuthService): "), 27 | 'Type': self.get_user_input("Enter the service type (e.g., Authentication): "), 28 | 'Criticality': self.get_user_input("Enter the service criticality (Tier1/Tier2/Tier3): ", 29 | validation=lambda x: x in ['Tier1', 'Tier2', 'Tier3']) 30 | } 31 | 32 | print(colored("\n# Service Functionality", 'green', attrs=['bold'])) 33 | self.service_description['Functionality'] = self.get_user_input("Enter a brief description of the service functionality (e.g., Handles user authentication and authorization.): ") 34 | 35 | print(colored("\n# Data Processing Details", 'green', attrs=['bold'])) 36 | self.service_description['DataProcessed'] = { 37 | 'Type': self.get_user_input("Enter the type of data processed (Secret/Confidential/Internal/Public): ", 38 | validation=lambda x: x in ['Secret', 'Confidential', 'Internal', 'Public']), 39 | 'DataCategory': self.get_user_input("Enter the data category (Auth/PCI/PII/etc): "), 40 | 'EncryptionAtRest': self.get_user_input("Is data encrypted at rest? (Yes/No): ", 41 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize() 42 | } 43 | 44 | print(colored("\n# Components Used by the Service", 'green', attrs=['bold'])) 45 | self.service_description['Components'] = { 46 | 'Internal': { 47 | 'Exist': self.get_user_input("Do internal components exist? (Yes/No): ", 48 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize() 49 | }, 50 | 'External': { 51 | 'Exist': self.get_user_input("Do external components exist? (Yes/No): ", 52 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize() 53 | } 54 | } 55 | if self.service_description['Components']['Internal']['Exist'] == 'Yes': 56 | self.service_description['Components']['Internal']['Source'] = self.get_user_input("Enter the source of internal components (Private/Public): ") 57 | self.service_description['Components']['Internal']['Note'] = self.get_user_input("Enter any notes about internal components (e.g., Namespacing/Scoped Package Access): ", required=False) 58 | if self.service_description['Components']['External']['Exist'] == 'Yes': 59 | self.service_description['Components']['External']['PackageManager'] = self.get_user_input("Enter the package manager for external components (NPM/Maven/NuGet/RubyGems/etc): ") 60 | 61 | print(colored("\n# Pipeline Configuration", 'green', attrs=['bold'])) 62 | self.service_description['Pipeline'] = { 63 | 'Type': self.get_user_input("Enter the CI/CD pipeline type (GithubActions/Jenkins/etc): "), 64 | 'CODEOWNERS': self.get_user_input("Are CODEOWNERS used? (Yes/No): ", 65 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize(), 66 | 'BranchProtection': self.get_user_input("Is branch protection enabled? (Yes/No): ", 67 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize(), 68 | 'SignCommits': self.get_user_input("Are commits signed? (Yes/No): ", 69 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize(), 70 | 'PinActions': self.get_user_input("Are actions pinned? (Yes/No): ", 71 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize() 72 | } 73 | 74 | print(colored("\n# Network Information", 'green', attrs=['bold'])) 75 | self.service_description['Network'] = { 76 | 'Access': self.get_user_input("Enter the network access level (Public/Private): ", 77 | validation=lambda x: x.lower() in ['public', 'private']).capitalize() 78 | } 79 | 80 | print(colored("\n# Data Flow", 'green', attrs=['bold'])) 81 | self.service_description['dataFlow'] = [] 82 | 83 | print(colored("Illustrate the flow of data within the service, including sources, targets, and methods.", 'yellow')) 84 | print(colored("Example:", 'yellow')) 85 | print(colored(" - name: UserAuthenticationFlow", 'yellow')) 86 | print(colored(" description: Handles user login and authentication.", 'yellow')) 87 | print(colored(" source: UserLoginInterface", 'yellow')) 88 | print(colored(" EncryptionTransit: Yes", 'yellow')) 89 | print(colored(" Authentication:", 'yellow')) 90 | print(colored(" Exist: Yes", 'yellow')) 91 | print(colored(" Type: JWT", 'yellow')) 92 | print(colored(" Authorization: read-write", 'yellow')) 93 | print(colored(" Protocol: HTTPS", 'yellow')) 94 | print(colored(" Communication:", 'yellow')) 95 | print(colored(" Type: RESTful API", 'yellow')) 96 | print(colored(" interactions:", 'yellow')) 97 | print(colored(" - from: UserLoginInterface", 'yellow')) 98 | print(colored(" to: AuthService", 'yellow')) 99 | print(colored(" method: RESTful API", 'yellow')) 100 | print(colored(" protocol: HTTPS", 'yellow')) 101 | print(colored(" - from: AuthService", 'yellow')) 102 | print(colored(" to: UserDatabase", 'yellow')) 103 | print(colored(" method: Query", 'yellow')) 104 | print(colored(" protocol: JDBC", 'yellow')) 105 | print(colored(" servicesInvolved: [UserLoginInterface, AuthService, UserDatabase]", 'yellow')) 106 | 107 | while True: 108 | add_flow = self.get_user_input("\nDo you want to add a data flow? (Yes/No): ", 109 | validation=lambda x: x.lower() in ['yes', 'no']) 110 | if add_flow.lower() != 'yes': 111 | break 112 | 113 | data_flow = { 114 | 'name': self.get_user_input("Enter the name of the data flow: "), 115 | 'description': self.get_user_input("Enter a brief description of the data flow: "), 116 | 'source': self.get_user_input("Enter the source of the data flow: "), 117 | 'EncryptionTransit': self.get_user_input("Is data encrypted in transit? (Yes/No): ", 118 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize(), 119 | 'Authentication': { 120 | 'Exist': self.get_user_input("Does authentication exist for this data flow? (Yes/No): ", 121 | validation=lambda x: x.lower() in ['yes', 'no']).capitalize() 122 | }, 123 | 'Authorization': self.get_user_input("Enter the authorization level (read/write/admin/etc): "), 124 | 'Protocol': self.get_user_input("Enter the communication protocol (HTTPS/AMQP/etc): "), 125 | 'Communication': { 126 | 'Type': self.get_user_input("Enter the communication type (RESTful API/Message Queues/etc): ") 127 | }, 128 | 'interactions': [], 129 | 'servicesInvolved': [] 130 | } 131 | 132 | if data_flow['Authentication']['Exist'] == 'Yes': 133 | data_flow['Authentication']['Type'] = self.get_user_input("Enter the authentication type (JWT/API Keys/etc): ") 134 | 135 | while True: 136 | add_interaction = self.get_user_input("Do you want to add an interaction? (Yes/No): ", 137 | validation=lambda x: x.lower() in ['yes', 'no']) 138 | if add_interaction.lower() != 'yes': 139 | break 140 | 141 | interaction = { 142 | 'from': self.get_user_input("Enter the source of the interaction (e.g., UserLoginInterface): "), 143 | 'to': self.get_user_input("Enter the target of the interaction (e.g., AuthService): "), 144 | 'method': self.get_user_input("Enter the method of the interaction (e.g., RESTful API): "), 145 | 'protocol': self.get_user_input("Enter the protocol of the interaction (e.g., HTTPS): ") 146 | } 147 | data_flow['interactions'].append(interaction) 148 | data_flow['servicesInvolved'].extend([interaction['from'], interaction['to']]) 149 | 150 | data_flow['servicesInvolved'] = list(set(data_flow['servicesInvolved'])) 151 | self.service_description['dataFlow'].append(data_flow) 152 | 153 | return self.service_description 154 | 155 | def save_yaml_file(self, file_name): 156 | with open(file_name, 'w') as file: 157 | file.write(f"Version: '{self.service_description['Version']}'\n") 158 | file.write(f"Date: {self.service_description['Date']}\n\n") 159 | file.write("# Service Description\n") 160 | file.write(yaml.dump({'Description': self.service_description['Description']}, default_flow_style=False)) 161 | file.write("\n# Service Functionality\n") 162 | file.write(f"Functionality: {self.service_description['Functionality']}\n\n") 163 | file.write("# Data Processing Details\n") 164 | file.write(yaml.dump({'DataProcessed': self.service_description['DataProcessed']}, default_flow_style=False)) 165 | file.write("\n# Components Used by the Service\n") 166 | file.write(yaml.dump({'Components': self.service_description['Components']}, default_flow_style=False)) 167 | file.write("\n# Pipeline Configuration\n") 168 | file.write(yaml.dump({'Pipeline': self.service_description['Pipeline']}, default_flow_style=False)) 169 | file.write("\n# Network Information\n") 170 | file.write(yaml.dump({'Network': self.service_description['Network']}, default_flow_style=False)) 171 | file.write("\n# Data Flow\n") 172 | file.write(yaml.dump({'dataFlow': self.service_description['dataFlow']}, default_flow_style=False)) 173 | 174 | print(colored(f"\nService description saved to {file_name}", 'green')) 175 | 176 | def main(): 177 | print(colored("**Welcome to the Service Description Generator!**", 'cyan', attrs=['bold'])) 178 | print(colored("- This tool is designed to generate a valid service description for AI-driven Threat modeling-as-a-Code (TaaC-AI).", 'cyan')) 179 | print(colored("- TaaC-AI is available at https://github.com/yevh/TaaC-AI/", 'cyan')) 180 | print(colored("- This tool will guide you through creating a service description in YAML format.", 'cyan')) 181 | print(colored("- Please provide the requested information as prompted.", 'cyan')) 182 | 183 | generator = ServiceDescriptionGenerator() 184 | generator.create_service_description() 185 | 186 | file_name = generator.get_user_input("\nEnter the name of the YAML file to save the service description: ") 187 | if not file_name.endswith('.yaml'): 188 | file_name += '.yaml' 189 | 190 | generator.save_yaml_file(file_name) 191 | 192 | if __name__ == "__main__": 193 | main() 194 | -------------------------------------------------------------------------------- /template.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | {{ service.Description.Name }} Report 7 | 8 | 155 | 156 | 157 |
158 |
159 |

{{ service.Description.Name }} - Threat Modeling Report

160 | 161 |
162 | 163 |
164 | 165 |
166 |
Service Description
167 |
168 |

Name: {{ service.Description.Name }}

169 |

Type: {{ service.Description.Type }}

170 |

Criticality: {{ service.Description.Criticality }}

171 |
172 |
173 | 174 | 175 |
176 |
Service Functionality
177 |
178 |

{{ service.Functionality }}

179 |
180 |
181 | 182 | 183 |
184 |
Data Processing Details
185 |
186 |

Type: {{ service.DataProcessed.Type }}

187 |

Data Category: {{ service.DataProcessed.DataCategory }}

188 |

Encryption at Rest: {{ service.DataProcessed.EncryptionAtRest }}

189 |
190 |
191 | 192 | 193 |
194 |
Network Information
195 |
196 |

Access: {{ service.Network.Access }}

197 |
198 |
199 | 200 | 201 |
202 |
Components
203 |
204 |

Internal Components

205 |

Exist: {{ service.Components.Internal.Exist }}

206 |

Source: {{ service.Components.Internal.Source }}

207 |

External Components

208 |

Exist: {{ service.Components.External.Exist }}

209 |

Package Manager: {{ service.Components.External.PackageManager }}

210 |
211 |
212 | 213 | 214 |
215 |
Pipeline Configuration
216 |
217 |

Type: {{ service.Pipeline.Type }}

218 |

CODEOWNERS: {{ service.Pipeline.CODEOWNERS }}

219 |

Branch Protection: {{ service.Pipeline.BranchProtection }}

220 |

Sign Commits: {{ service.Pipeline.SignCommits }}

221 |

Pin Actions: {{ service.Pipeline.PinActions }}

222 |
223 |
224 |
225 | 226 | 227 |
228 |
229 |

Data Flow Information

230 |
231 |
232 |
233 |
234 |
235 | 236 | 237 |
238 |
Threats
239 |
240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | {{ threat_analysis_html | safe }} 254 | 255 |
TitleValidatorDescriptionCategoriesRemediationStatusActions
256 |
257 | 258 |
259 |
260 |
261 |
262 | 263 | 264 | 405 | 406 | 466 | 467 | 475 | 476 | 477 | 478 | --------------------------------------------------------------------------------