from langchain_google_genai import ChatGoogleGenerativeAI from langchain_core.messages import SystemMessage, HumanMessage, AIMessage from langchain_core.output_parsers import PydanticOutputParser, JsonOutputParser from langchain.output_parsers import RetryOutputParser, OutputFixingParser from langchain_core.prompts import PromptTemplate from pydantic import BaseModel, Field, field_validator from typing import Optional, Dict, Any, List, Union import json import logging from models.attack_schema import AttackLayer, Technique, DEFAULT_LAYER_SETTINGS # ロガーの設定 logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) class EvaluationResult(BaseModel): is_valid: bool = Field(description="コンテキスト評価結果（TrueまたはFalse）") extracted_scenario: Optional[str] = Field(None, description="シナリオの抽出結果") extracted_layer_operation: Optional[str] = Field(None, description="レイヤー操作の抽出結果") class LLMService: """Service for handling LLM interactions""" def __init__(self, model_name: str = "gemini-2.0-flash", temperature: float = 0.7): self.llm = ChatGoogleGenerativeAI( model=model_name, temperature=temperature, ) self.attack_evaluation_parser = PydanticOutputParser(pydantic_object=EvaluationResult) self.attack_evaluation_fixing_parser = OutputFixingParser.from_llm( parser=self.attack_evaluation_parser, llm=self.llm ) self.layer_parser = PydanticOutputParser(pydantic_object=AttackLayer) self.layer_fixing_parser = OutputFixingParser.from_llm( parser=self.layer_parser, llm=self.llm ) self.json_parser = JsonOutputParser(pydantic_object=AttackLayer) self.json_retry_parser = RetryOutputParser.from_llm( parser=self.json_parser, llm=self.llm ) async def evaluate_context(self, user_input: str) -> EvaluationResult: """Evaluate if the user input is valid for ATT&CK context""" system_prompt = """ 以下のユーザ入力が、サイバー攻撃の分析やMITRE ATT&CKフレームワークのレイヤーに対する指示として適切かどうかを評価し、指示を分割してください。 評価基準: 1. サイバー攻撃のシナリオや分析に関する内容か 2. ATT&CKフレームワークのレイヤーに対する指示や更新要求か 3. 攻撃テクニックやタクティックに関する質問や説明か 回答形式: {format_instructions} 注意: - 評価結果（is_valid）は必ずTrueまたはFalseのいずれかのみを出力してください。 - True: コンテキストに合致 - False: コンテキストに合致しない - シナリオの抽出結果（extracted_scenario）は、ユーザ入力の文書からサイバー攻撃分析シナリオに関する部分を抽出してください。なければ省力可能。 - レイヤー操作の抽出結果（extracted_layer_operation）は、ユーザ入力の文書からMITRE ATT&CKフレームワークのレイヤー操作に関する部分を抽出してください。なければ省略可能。 """ logger.info(f"Evaluating context for: {user_input}") try: response = await self.llm.ainvoke([ SystemMessage(content=system_prompt.format( format_instructions=self.attack_evaluation_parser.get_format_instructions() )), HumanMessage(content=user_input) ]) result = self.attack_evaluation_fixing_parser.parse(response.content) if not result.is_valid: logger.info("Invalid context detected") return result except Exception as e: logger.error(f"Error in evaluate_context: {str(e)}") raise async def generate_scenario(self, user_input: str, current_scenario: Optional[str] = None) -> str: """Generate or update a scenario based on user input""" if not user_input: user_input = "シナリオ更新について指示はありません" if not current_scenario: system_prompt = """ ユーザ入力から想起されるサイバー攻撃分析シナリオを作成してください。 要件: 1. シナリオには以下の構造を含めてください： - タイトル - 概要 - 攻撃フェーズ 2. 各攻撃フェーズは、MITRE ATT&CKフレームワークのTacticsに沿ったものにしてください 3. シナリオは具体的で分析可能な内容にしてください 4. 実際の攻撃事例や既知の脅威アクターの戦術を参考にしてください 5. ユーザの具体的な要件や指示に従ってください """ human_prompt = f"以下の要件に従って、サイバー攻撃分析シナリオを作成してください：\n\n{user_input}" else: system_prompt = """ 既存のシナリオを更新してください。 要件: 1. ユーザの指示に従ってシナリオを更新 2. 既存の情報は保持しつつ、新しい情報を追加または修正 3. シナリオの一貫性を維持 4. シナリオの構造（タイトル、概要、攻撃フェーズなど）を維持 5. ユーザの具体的な要件や指示に従ってください """ human_prompt = f""" 現在のシナリオ: {current_scenario} 以下の要件に従って、シナリオを更新してください： {user_input} """ try: response = await self.llm.ainvoke([ SystemMessage(content=system_prompt), HumanMessage(content=human_prompt) ]) return response.content.strip() except Exception as e: logger.error(f"Error in generate_scenario: {str(e)}") raise async def generate_attack_json(self, layer_operation: str = None, scenario: str = None, existing_json: Optional[Dict] = None) -> str: """Generate or update ATT&CK Navigator JSON based on scenario""" if not scenario: raise ValueError("シナリオが指定されていません") if not layer_operation: layer_operation = "レイヤー操作に関して指定はありません" try: system_prompt = """ MITRE ATT&CK Navigatorのレイヤーを生成または更新してください。 要件: 1. シナリオの各攻撃フェーズに対応するテクニックを特定 2. テクニックにはコメントとして、シナリオ内での使用方法を説明 3. テクニックの選択は具体的な根拠に基づくこと 4. 既存のJSONがある場合は、その情報を考慮して更新 5. versionsオブジェクトには必ず以下の3つのキーを含めること： - attack: ATT&CKのバージョン - navigator: Navigatorのバージョン - layer: レイヤーのバージョン DEFAULT_LAYER_SETTINGS: {DEFAULT_LAYER_SETTINGS} バージョン指定のルール: 1. ユーザーが特定のバージョンを指定した場合（例：「ATT&CKバージョンは16で」）、その値を優先して使用 2. ユーザーが指定していない場合は、デフォルトのバージョン設定を使用 3. 既存のJSONがある場合は、そのバージョン情報を保持 """ human_prompt = f""" レイヤー操作: {layer_operation} シナリオ: {scenario} {'既存のJSON:' + json.dumps(existing_json, indent=2) if existing_json else '新規作成'} """ response = await self.llm.ainvoke([ SystemMessage(content=system_prompt), HumanMessage(content=human_prompt) ]) try: result = self.layer_fixing_parser.parse(response.content) # バージョン情報が正しくない場合は、デフォルトのバージョン情報を使用 if not all(k in result.versions for k in ['attack', 'navigator', 'layer']): result.versions = DEFAULT_LAYER_SETTINGS['versions'] return json.dumps(result.model_dump(), indent=2, ensure_ascii=False) except Exception as e: logger.error(f"JSON parsing error: {str(e)}") raise except Exception as e: logger.error(f"Error in generate_attack_json: {str(e)}") raise