feat: initial project setup and core functionality

- Add basic MITRE ATT&CK layer generation
- Configure logging and error handling
- Setup project structure and dependencies
- Add .gitignore for Python/Pipenv
- Implement core LLM service

.gitignore

@@ -0,0 +1,55 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+env/
+ENV/
+
+# Pipenv
+Pipfile
+Pipfile.lock
+.pipenv/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Logs
+*.log
+logs/
+
+# Local development settings
+.env
+.env.local
+.env.*.local
+
+# Chainlit
+.chainlit/
+chainlit.md
+
+# Project specific
+responses/
+public/responses/
+test_output_*.json 

Dockerfile

@@ -4,6 +4,8 @@ RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 
+RUN --mount=type=secret,id=GEMINI_API_KEY,mode=0444,required=true
+
 WORKDIR /app
 
 COPY --chown=user ./requirements.txt requirements.txt

app.py

@@ -1,28 +1,53 @@
 import chainlit as cl
+from models.chat_state import get_initial_state, add_user_message
+from workflows.chat_workflow import chainlit_app
+
+# Chainlit message handler
+@cl.on_message
+async def process_user_message(message: cl.Message):
+    """Process user message and generate response"""
+    # Get state from session
+    state = cl.user_session.get("state")
+    
+    # Add user message to state
+    state = add_user_message(state, message.content)
+    
+    # Generate response
+    state = await chainlit_app.ainvoke(state)
+    
+    # Update state in session
+    cl.user_session.set("state", state)
 
 @cl.on_chat_start
 async def start():
     """
     This function is called when a new chat starts.
-    You can use it to initialize any variables or models.
-    """
-    # Initialize any variables or models here
-    await cl.Message(
-        content="Welcome to the chatbot! How can I help you today?"
-    ).send()
-
-@cl.on_message
-async def main(message: str):
-    """
-    This function is called every time a user inputs a message.
-    Args:
-        message: The user's message
-    """
-    # Process the user's message here
-    response = f"You said: {message.content}"
+    Initialize the chat state and store it in the session
+    """  
+    # Initialize state
+    state = get_initial_state()
     
-    # Send a response back to the user
-    await cl.Message(
-        content=response
-    ).send()
+    # Store state in session
+    cl.user_session.set("state", state)
+
+
 
+@cl.set_chat_profiles
+async def chat_profile():
+    return [
+        cl.ChatProfile(
+            name="mitre_attck_navigator_layer_writer",
+            markdown_description="MITRE ATT&CK Navigatorのlayerを生成するチャットボットです。",
+            icon="public/icon.jpg",
+            starters=[
+                cl.Starter(
+                    label="Lockbitの攻撃シナリオ",
+                    message="Lockbitの攻撃シナリオを生成してください"
+                ),
+                cl.Starter(
+                    label="ランサムウェアギャングが使っているテクニックの頻度",
+                    message="最近のランサムウェアギャングが使っているテクニックと頻度を色の濃さで表現してください"
+                ),
+            ]
+        )
+    ]

models/attack_schema.py

@@ -0,0 +1,150 @@
+"""
+MITRE ATT&CK Navigator JSON Schema
+"""
+
+from pydantic import BaseModel, Field, validator, StringConstraints
+from typing import Optional, Dict, Any, List, Annotated
+import re
+from datetime import datetime
+
+# カスタム型の定義
+TechniqueID = Annotated[str, StringConstraints(pattern=r'^T\d{4}(\.\d{3})?$')]
+TacticName = Annotated[str, StringConstraints(pattern=r'^[a-z-]+$')]
+HexColor = Annotated[str, StringConstraints(pattern=r'^#[0-9A-Fa-f]{6}$')]
+Domain = Annotated[str, StringConstraints(pattern=r'^(enterprise-attack|mobile-attack|ics-attack)$')]
+
+class Technique(BaseModel):
+    """Technique model for ATT&CK Navigator"""
+    techniqueID: TechniqueID = Field(description="The MITRE ATT&CK technique ID (e.g., T1234 or T1234.001)")
+    tactic: Optional[TacticName] = Field(description="The tactic this technique belongs to (lowercase with hyphens)")
+    score: Optional[int] = Field(description="Score between 0 and 100", ge=0, le=100)
+    color: Optional[HexColor] = Field(description="Hex color code (e.g., #FF0000)")
+    comment: Optional[str] = Field(description="Comment about the technique", max_length=1000)
+    enabled: Optional[bool] = Field(description="Whether the technique is enabled", default=True)
+    metadata: Optional[List[Dict[str, str]]] = Field(description="Additional metadata")
+    links: Optional[List[Dict[str, str]]] = Field(description="Related links")
+    showSubtechniques: Optional[bool] = Field(description="Whether to show subtechniques", default=True)
+
+    @validator('links')
+    def validate_links(cls, v):
+        if v is not None:
+            for link in v:
+                if not all(k in link for k in ['label', 'url']):
+                    raise ValueError("Each link must have 'label' and 'url' fields")
+        return v
+
+class VersionInfo(BaseModel):
+    """バージョン情報のサブスキーマ"""
+    attack: str = Field(
+        ...,
+        description="""ATT&CKのバージョン。ユーザーが指定したバージョンを使用すること。
+        例: ユーザーが「ATT&CKバージョンは16で」と指定した場合、"16"と設定する。
+        指定がない場合は、DEFAULT_LAYER_SETTINGSの値を使用する。"""
+    )
+    navigator: str = Field(
+        ...,
+        description="""Navigatorのバージョン。ユーザーが指定したバージョンを使用すること。
+        指定がない場合は、DEFAULT_LAYER_SETTINGSの値を使用する。"""
+    )
+    layer: str = Field(
+        ...,
+        description="""レイヤーのバージョン。ユーザーが指定したバージョンを使用すること。
+        指定がない場合は、DEFAULT_LAYER_SETTINGSの値を使用する。"""
+    )
+
+class AttackLayer(BaseModel):
+    """MITRE ATT&CK Navigator Layer"""
+    name: str = Field(..., description="レイヤーの名前")
+    description: str = Field(..., description="レイヤーの説明")
+    domain: Domain = Field(..., description="ATT&CKドメイン")
+    versions: VersionInfo = Field(..., description="バージョン情報")
+    filters: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "platforms": [
+                "Windows",
+                "Linux",
+                "macOS",
+                "Containers",
+                "IaaS",
+                "Network"
+            ]
+        },
+        description="フィルター設定"
+    )
+    sorting: int = Field(default=0, description="ソート順")
+    layout: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "layout": "side",
+            "showName": True,
+            "showID": False,
+            "showAggregateScores": True,
+            "countUnscored": True,
+            "aggregateFunction": "average"
+        },
+        description="レイアウト設定"
+    )
+    hideDisabled: bool = Field(default=False, description="無効なテクニックを非表示")
+    techniques: List[Technique] = Field(..., description="テクニックのリスト")
+    gradient: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "colors": ["#ff6666", "#ffe766", "#8ec843"],
+            "minValue": 0,
+            "maxValue": 100
+        },
+        description="スコアのグラデーション設定"
+    )
+    legendItems: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="凡例アイテム"
+    )
+    showTacticRowBackground: bool = Field(default=False, description="タクティック行の背景を表示")
+    tacticRowBackground: str = Field(default="#dddddd", description="タクティック行の背景色")
+    selectTechniquesAcrossTactics: bool = Field(default=True, description="タクティック間でテクニックを選択")
+    selectSubtechniquesWithParent: bool = Field(default=False, description="サブテクニックを親と共に選択")
+    selectVisibleTechniques: bool = Field(default=False, description="表示されているテクニックのみを選択")
+    metadata: List[Dict[str, str]] = Field(
+        default_factory=lambda: [
+            {"name": "作成者", "value": "AIアシスタント"},
+            {"name": "作成日", "value": datetime.now().strftime("%Y-%m-%d")}
+        ],
+        description="メタデータ"
+    )
+
+# デフォルトのレイヤー設定
+DEFAULT_LAYER_SETTINGS = {
+    "versions": {
+        "attack": "16.1",
+        "navigator": "5.1.0",
+        "layer": "4.5"
+    },
+    "layout": {
+        "layout": "side",
+        "showName": True,
+        "showID": False,
+        "showAggregateScores": True,
+        "countUnscored": True,
+        "aggregateFunction": "average"
+    },
+    "gradient": {
+        "colors": ["#ff6666", "#ffe766", "#8ec843"],
+        "minValue": 0,
+        "maxValue": 100
+    }
+}
+
+# プロンプトテンプレート
+ATTACK_PROMPT = """
+あなたはMITRE ATT&CKフレームワークの専門家です。
+以下のシナリオに基づいて、MITRE ATT&CK NavigatorのレイヤーJSONを生成してください。
+
+回答形式:
+{format_instructions}
+
+注意:
+- 有効なMITRE ATT&CKテクニックIDを使用すること
+- 適切なタクティクスを含めること
+- スコアは0-100の範囲で設定すること
+- 有効な16進数カラーコードを使用すること
+- 意味のあるコメントを含めること
+- 適切なメタデータを設定すること
+""" 

models/chat_state.py

@@ -0,0 +1,48 @@
+from typing import List, Optional, Union, Dict, Any
+from langgraph.graph.message import MessagesState
+from langchain_core.messages import AIMessage, HumanMessage
+from pydantic import BaseModel
+
+class AttackState(MessagesState, total=False):
+    """State for the ATT&CK Navigator workflow"""
+    attack_json: Optional[Dict[str, Any]] = None
+    scenario: Optional[str] = None
+    is_valid_context: Optional[bool] = None
+    extracted_user_scenario: Optional[str] = None
+    extracted_user_layer_operation: Optional[str] = None
+
+def get_initial_state() -> AttackState:
+    """Get the initial state for the workflow"""
+    return AttackState(
+        messages=[],
+        attack_json=None,
+        scenario=None,
+        is_valid_context=None,
+        extracted_user_scenario=None,
+        extracted_user_layer_operation=None
+    )
+
+def add_user_message(state: AttackState, content: str) -> AttackState:
+    """Add a user message to the state"""
+    state['messages'].append(HumanMessage(content=content))
+    return state
+
+def add_ai_message(state: AttackState, content: str) -> AttackState:
+    """Add an AI message to the state"""
+    state['messages'].append(AIMessage(content=content))
+    return state
+
+def set_attack_json(state: AttackState, attack_json: Dict[str, Any]) -> AttackState:
+    """Set the ATT&CK JSON in the state"""
+    state['attack_json'] = attack_json
+    return state
+
+def set_scenario(state: AttackState, scenario: str) -> AttackState:
+    """Set the scenario text in the state"""
+    state['scenario'] = scenario
+    return state
+
+def set_valid_context(state: AttackState, is_valid: bool) -> AttackState:
+    """Set the context validity in the state"""
+    state['is_valid_context'] = is_valid
+    return state 

requirements.txt

@@ -1,2 +1,7 @@
 chainlit
 websockets
+langgraph
+google-generativeai
+langchain
+langchain-google-genai
+uvicorn

services/file_service.py

@@ -0,0 +1,30 @@
+import os
+import chainlit as cl
+import uuid
+from datetime import datetime
+# Ensure public/jsons directory exists
+os.makedirs("public/jsons", exist_ok=True)
+
+def save_response(content: str) -> tuple[str, str]:
+    """
+    Save response content to a file and return filename and filepath
+    """
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    file_uuid = str(uuid.uuid4())
+    filename = f"layer_{file_uuid}_{timestamp}.json"
+    filepath = os.path.join("public/jsons", filename)
+    
+    with open(filepath, "w", encoding="utf-8") as f:
+        f.write(content)
+    
+    return filename, filepath
+
+def create_file_element(filename: str, filepath: str) -> cl.File:
+    """
+    Create a Chainlit File element for the saved response
+    """
+    return cl.File(
+        name=filename,
+        path=filepath,
+        display="inline",
+    ) 

services/llm_service.py

@@ -0,0 +1,190 @@
+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 

workflows/chat_workflow.py

@@ -0,0 +1,376 @@
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import MessagesState
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from services.llm_service import LLMService
+from services.file_service import save_response, create_file_element
+from models.chat_state import AttackState, get_initial_state
+import chainlit as cl
+import json
+import logging
+import os
+from urllib.parse import quote
+# Initialize services
+llm_service = LLMService()
+logger = logging.getLogger(__name__)
+
+CHAINLIT_URL = os.environ.get("SPACE_HOST")
+if not CHAINLIT_URL:
+    CHAINLIT_URL = "http://localhost:8080"
+
+@cl.step(name="コンテキスト評価", type="evaluation")
+async def evaluate_context_node(state: AttackState) -> AttackState:
+    """Node for evaluating if the user input is valid for ATT&CK context"""
+    msg = cl.Message(content="")
+    
+    # Get the last user message
+    user_messages = [msg for msg in state['messages'] if isinstance(msg, HumanMessage)]
+    user_message = user_messages[-1].content if user_messages else ""
+    
+    try:
+        evaluation_result = await llm_service.evaluate_context(user_message)
+        state['is_valid_context'] = evaluation_result.is_valid
+        state['extracted_user_scenario'] = evaluation_result.extracted_scenario
+        state['extracted_user_layer_operation'] = evaluation_result.extracted_layer_operation
+        
+        if state['is_valid_context']:
+            response_text = "入力はATT&CKフレームワークのコンテキストに合致します。シナリオの評価を続けます。"
+        else:
+            response_text = "申し訳ありませんが、この入力はサイバー攻撃の分析やATT&CKフレームワークのレイヤーに関する指示として認識できませんでした。適切な指示を入力してください。"
+        
+        await msg.stream_token(response_text)
+        await msg.send()
+        state['messages'].append(AIMessage(content=response_text))
+        
+    except Exception as e:
+        error_msg = f"コンテキスト評価中にエラーが発生しました: {str(e)}"
+        await msg.stream_token(error_msg)
+        await msg.send()
+        state['messages'].append(AIMessage(content=error_msg))
+        state['is_valid_context'] = False
+    
+    return state
+
+@cl.step(name="シナリオ更新", type="update")
+async def update_scenario_node(state: AttackState) -> AttackState:
+    """Node for updating the scenario based on user input"""
+    msg = cl.Message(content="")
+    
+    # Get the last user message
+    user_message = state.get('extracted_user_scenario', None)
+    current_scenario = state.get('scenario', None)
+    if not user_message and not current_scenario:
+        raise ValueError("シナリオの更新に必要な情報がありません。")
+        
+    try:
+        updated_scenario = await llm_service.generate_scenario(user_message, current_scenario)
+        state['scenario'] = updated_scenario
+        
+        message = "新しいシナリオを作成しました。" if not current_scenario else "シナリオを更新しました。"
+        await msg.stream_token(message)
+        await msg.send()
+        state['messages'].append(AIMessage(content=message))
+        
+    except Exception as e:
+        error_msg = f"シナリオの{'作成' if not current_scenario else '更新'}中にエラーが発生しました: {str(e)}"
+        await msg.stream_token(error_msg)
+        await msg.send()
+        state['messages'].append(AIMessage(content=error_msg))
+    
+    return state
+
+@cl.step(name="JSON生成/更新", type="generation", language="json")
+async def generate_json_node(state: AttackState) -> AttackState:
+    """Node for generating or updating ATT&CK Navigator JSON"""
+    user_message = state.get('extracted_user_layer_operation')
+    current_scenario = state.get('scenario')
+    existing_json = state.get('attack_json')
+    
+    try:
+        json_content = await llm_service.generate_attack_json(user_message, current_scenario, existing_json)
+        
+        # Save JSON to file
+        filename, filepath = save_response(json_content)
+        file_element = create_file_element(filename, filepath)
+        
+        json_url = CHAINLIT_URL + "/" + filepath
+        json_url = quote(json_url)
+
+        # Prepare and send the response message
+        response = "MITRE ATT&CK Navigatorレイヤーを更新しました。" if existing_json else "MITRE ATT&CK Navigatorレイヤーを生成しました。"
+        response += " ファイルをダウンロードしてインポートできます。"
+        response += f"ATT&CK Navigator : https://mitre-attack.github.io/attack-navigator//#layerURL={json_url}"
+        
+        msg = cl.Message(content=response, elements=[file_element])
+        await msg.send()
+        
+        # Update state
+        state['messages'].append(AIMessage(content=response))
+        state['attack_json'] = json.loads(json_content)
+        
+    except Exception as e:
+        error_msg = f"ATT&CK Navigatorレイヤーの生成中にエラーが発生しました: {str(e)}"
+        msg = cl.Message(content=error_msg)
+        await msg.send()
+        state['messages'].append(AIMessage(content=error_msg))
+    
+    return state
+
+async def display_state_node(state: AttackState) -> AttackState:
+    """Node for displaying the current state before ending"""
+    async with cl.Step(name="状態表示", type="display") as step:
+        # コンテキストの評価結果
+        if state.get('is_valid_context') is not None:
+            status = "有効" if state['is_valid_context'] else "無効"
+            step.input = "コンテキスト評価"
+            step.output = f"コンテキストの評価結果: {status}"
+        
+        # 現在のシナリオ
+        if state.get('scenario'):
+            async with cl.Step(name="シナリオ情報", type="display") as scenario_step:
+                scenario_step.input = "現在のシナリオ"
+                scenario_step.output = state['scenario']
+        
+        # JSONの状態
+        if state.get('attack_json'):
+            async with cl.Step(name="ATT&CK情報", type="display") as attack_step:
+                techniques = state['attack_json'].get('techniques', [])
+                technique_count = len(techniques)
+                attack_step.input = "登録済みテクニック"
+                
+                if technique_count > 0:
+                    technique_list = "\n".join([
+                        f"- {t.get('techniqueID', 'Unknown')}: {t.get('comment', '説明なし')}"
+                        for t in techniques[:5]
+                    ])
+                    if technique_count > 5:
+                        technique_list += f"\n... 他 {technique_count - 5} 件"
+                    attack_step.output = f"登録済みテクニック数: {technique_count}\n\n{technique_list}"
+                else:
+                    attack_step.output = "登録済みテクニックはありません"
+    msg = cl.Message(content="処理が完了しました。")
+    await msg.send()
+    return state
+
+# Create the graph
+workflow = StateGraph(AttackState)
+
+# Add nodes
+workflow.add_node("evaluate_context", evaluate_context_node)
+workflow.add_node("update_scenario", update_scenario_node)
+workflow.add_node("generate_json", generate_json_node)
+workflow.add_node("display_state", display_state_node)
+
+# Add edges
+workflow.add_edge(START, "evaluate_context")
+workflow.add_conditional_edges(
+    "evaluate_context",
+    lambda state: state.get('is_valid_context', False),
+    {
+        True: "update_scenario",
+        False: "display_state"
+    }
+)
+workflow.add_edge("update_scenario", "generate_json")
+workflow.add_edge("generate_json", "display_state")
+workflow.add_edge("display_state", END)
+
+# Compile the graph
+chainlit_app = workflow.compile()
+
+async def test_evaluate_context_node(state: AttackState) -> AttackState:
+    """テスト用のコンテキスト評価ノード"""
+    user_messages = [msg for msg in state['messages'] if isinstance(msg, HumanMessage)]
+    user_message = user_messages[-1].content if user_messages else ""
+    
+    try:
+        evaluation_result = await llm_service.evaluate_context(user_message)
+        state['is_valid_context'] = evaluation_result.is_valid
+        state['extracted_user_scenario'] = evaluation_result.extracted_scenario
+        state['extracted_user_layer_operation'] = evaluation_result.extracted_layer_operation
+        response_text = "入力はATT&CKフレームワークのコンテキストに合致します。シナリオの評価を続けます。" if state['is_valid_context'] else "申し訳ありませんが、この入力はサイバー攻撃の分析やATT&CKフレームワークのレイヤーに関する指示として認識できませんでした。適切な指示を入力してください。"
+        state['messages'].append(AIMessage(content=response_text))
+    except Exception as e:
+        error_msg = f"コンテキスト評価中にエラーが発生しました: {str(e)}"
+        state['messages'].append(AIMessage(content=error_msg))
+        state['is_valid_context'] = False
+    
+    return state
+
+async def test_update_scenario_node(state: AttackState) -> AttackState:
+    """テスト用のシナリオ更新ノード"""
+    # Get the last user message
+    user_message = state.get('extracted_user_scenario')
+    current_scenario = state.get('scenario')
+    
+    try:
+        updated_scenario = await llm_service.generate_scenario(user_message, current_scenario)
+        state['scenario'] = updated_scenario
+        message = "新しいシナリオを作成しました。" if not current_scenario else "シナリオを更新しました。"
+        state['messages'].append(AIMessage(content=message))
+    except Exception as e:
+        error_msg = f"シナリオの{'作成' if not current_scenario else '更新'}中にエラーが発生しました: {str(e)}"
+        state['messages'].append(AIMessage(content=error_msg))
+    
+    return state
+
+async def test_generate_json_node(state: AttackState) -> AttackState:
+    """テスト用のJSON生成ノード"""
+    user_message = state.get('extracted_user_layer_operation')
+    current_scenario = state.get('scenario')
+    existing_json = state.get('attack_json')
+    
+    try:
+        json_content = await llm_service.generate_attack_json(user_message, current_scenario, existing_json)
+        response = "MITRE ATT&CK Navigatorレイヤーを更新しました。" if existing_json else "MITRE ATT&CK Navigatorレイヤーを生成しました。"
+        response += " ファイルをダウンロードしてインポートできます。"
+        state['messages'].append(AIMessage(content=response))
+        state['attack_json'] = json.loads(json_content)
+    except Exception as e:
+        error_msg = f"ATT&CK Navigatorレイヤーの生成中にエラーが発生しました: {str(e)}"
+        state['messages'].append(AIMessage(content=error_msg))
+    
+    return state
+
+async def test_display_state_node(state: AttackState) -> AttackState:
+    """テスト用の状態表示ノード"""
+    summary = []
+    
+    if state.get('is_valid_context') is not None:
+        status = "有効" if state['is_valid_context'] else "無効"
+        summary.append(f"コンテキストの評価結果: {status}")
+    
+    if state.get('scenario'):
+        summary.append(f"現在のシナリオ:\n{state['scenario']}")
+    
+    if state.get('attack_json'):
+        techniques = state['attack_json'].get('techniques', [])
+        technique_count = len(techniques)
+        summary.append(f"登録済みテクニック数: {technique_count}")
+        
+        if technique_count > 0:
+            technique_list = "\n".join([
+                f"- {t.get('techniqueID', 'Unknown')}: {t.get('comment', '説明なし')}"
+                for t in techniques[:5]
+            ])
+            if technique_count > 5:
+                technique_list += f"\n... 他 {technique_count - 5} 件"
+            summary.append(f"\n登録済みテクニック:\n{technique_list}")
+    
+    if summary:
+        state_summary = "\n\n".join(summary)
+        state['messages'].append(AIMessage(content=f"現在の状態:\n{state_summary}"))
+    
+    return state
+
+async def main():
+    """テスト用のメイン関数"""
+    try:
+        # 初期状態の作成
+        initial_state = get_initial_state()
+        
+        # テスト用の既存シナリオ
+        existing_scenario = """
+        標的システムへの不正アクセスシナリオ
+
+        概要:
+        攻撃者は、標的のシステムに不正アクセスを試み、機密情報を窃取します。
+
+        攻撃フェーズ:
+        1. 初期アクセス
+           - パスワードスプレー攻撃による認証情報の取得
+           - 有効なアカウントの特定
+
+        2. 実行
+           - 取得した認証情報を使用してシステムにログイン
+           - 不正なコマンドの実行
+
+        3. 権限昇格
+           - 管理者権限の取得
+           - システム設定の変更
+
+        4. 防御回避
+           - ログの削除
+           - 攻撃痕跡の隠蔽
+        """
+        
+        # テスト用の既存JSON
+        existing_json = {
+            "name": "Test Layer",
+            "versions": {
+                "attack": "16.0",
+                "navigator": "4.9.0",
+                "layer": "4.5"
+            },
+            "domain": "enterprise-attack",
+            "description": "Test layer for development",
+            "filters": {
+                "platforms": ["Windows", "Linux", "macOS"]
+            },
+            "gradient": {
+                "colors": ["#ffffff", "#ff6666"],
+                "minValue": 0,
+                "maxValue": 100
+            },
+            "techniques": [
+                {
+                    "techniqueID": "T1110",
+                    "score": 50,
+                    "color": "#ff6666",
+                    "comment": "パスワードスプレー攻撃による認証情報の取得",
+                    "enabled": True
+                },
+                {
+                    "techniqueID": "T1078",
+                    "score": 50,
+                    "color": "#ff6666",
+                    "comment": "有効なアカウントを使用した不正アクセス",
+                    "enabled": True
+                }
+            ]
+        }
+        
+        # 初期状態に既存データを設定
+        initial_state['scenario'] = existing_scenario
+        initial_state['attack_json'] = existing_json
+        
+        # テスト用のユーザーメッセージ
+        test_message = """
+        以下の攻撃シナリオを分析してください：
+        
+        攻撃者は、標的のシステムに不正アクセスを試みます。
+        まず、パスワードスプレー攻撃を実行して、有効なアカウントの認証情報を取得���ます。
+        取得した認証情報を使用して、システムにログインし、機密情報を窃取します。
+        最後に、攻撃の痕跡を隠蔽するために、ログを削除します。
+        """
+        
+        test_message = "ATTACKのバージョンを16にして、テクニックは青にして。"
+
+        # ユーザーメッセージを状態に追加
+        initial_state['messages'].append(HumanMessage(content=test_message))
+        
+        # テスト用ワークフローの実行
+        state = await test_evaluate_context_node(initial_state)
+        
+        if state.get('is_valid_context', False):
+            # シナリオ更新
+            state = await test_update_scenario_node(state)
+            
+            # JSON生成
+            state = await test_generate_json_node(state)
+        
+        # 状態表示
+        state = await test_display_state_node(state)
+        
+        # 結果の表示
+        for msg in state['messages']:
+            role = "User" if isinstance(msg, HumanMessage) else "Assistant"
+            print(f"\n{role}:")
+            print(msg.content)
+        
+    except Exception as e:
+        print(f"エラーが発生しました: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())