models/attack_schema.py

"""
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進数カラーコードを使用すること
- 意味のあるコメントを含めること
- 適切なメタデータを設定すること
"""