From b27b8c984dd4e27cacbbdd024d9768fa6fb5f1dc Mon Sep 17 00:00:00 2001 From: rookit Date: Tue, 21 Apr 2026 16:40:15 +0800 Subject: [PATCH] tmp --- Lib/basemodule.py | 5 + ...ivilege-Escalation-via-AttachUserPolicy.py | 17 +- ...ege-Escalation-via-AttachUserPolicy_old.py | 308 ------------------ .../skills/asp-module-creator-zh/SKILL.md | 3 + PLUGINS/Redis/redis_stream_api.py | 17 + PLUGINS/SIRP/correlation.py | 7 +- 6 files changed, 40 insertions(+), 317 deletions(-) delete mode 100644 MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy_old.py diff --git a/Lib/basemodule.py b/Lib/basemodule.py index 523f15b..80b9a20 100644 --- a/Lib/basemodule.py +++ b/Lib/basemodule.py @@ -17,6 +17,11 @@ class BaseModule(BaseAPI): self.agent_state = None self.debug_message_id = None # 设置为非None以启用Debug模式 + def read_stream_head_ids(self, n): # 调试时使用，读取最近的n条消息 + redis_stream_api = RedisStreamAPI() + messages = redis_stream_api.read_stream_head_ids(self.module_name, n=n) + return messages + def read_stream_message(self) -> dict: """读取消息""" redis_stream_api = RedisStreamAPI() diff --git a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py index 5e296e5..b322514 100644 --- a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py +++ b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py @@ -111,8 +111,8 @@ class Module(BaseModule): correlation_uid = Correlation.generate_correlation_uid( rule_id=self.module_name, - time_window="24h", - keys=[principal_user, target_user, account_id], + time_window="5m", + keys=[account_id], timestamp=event_time_formatted ) @@ -172,13 +172,13 @@ class Module(BaseModule): alert_model.artifacts = None if alert_enrichments: - alert_model.alert_enrichments = alert_enrichments + alert_model.enrichments = alert_enrichments else: - alert_model.alert_enrichments = None + alert_model.enrichments = None # 保存告警 saved_alert_row_id = Alert.create(alert_model) - self.logger.info(f"Alert created: {saved_alert_row_id} with disposition {disposition}") + self.logger.info(f"Alert created: {saved_alert_row_id}") # 5. Case 处理 (Case Management) try: @@ -219,5 +219,8 @@ if __name__ == "__main__": os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") django.setup() module = Module() - module.debug_message_id = "1776309110386-0" - module.run() + + message_ids = module.read_stream_head_ids(100) + for message_id in message_ids: + module.debug_message_id = message_id + module.run() diff --git a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy_old.py b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy_old.py deleted file mode 100644 index 3ccbf34..0000000 --- a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy_old.py +++ /dev/null @@ -1,308 +0,0 @@ -import json -from datetime import datetime -from typing import Optional, Union, Dict, Any, List - -from langchain_core.messages import HumanMessage -from langgraph.graph import StateGraph, END -from langgraph.graph.state import CompiledStateGraph -from langgraph.types import Command -from pydantic import BaseModel, Field - -from Lib.api import get_current_time_str -from Lib.basemodule import LanggraphModule -from Lib.llmapi import BaseAgentState -from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.grouprule import Correlation, CorrelationConfig -from PLUGINS.SIRP.sirpapi import Alert, Case -from PLUGINS.SIRP.sirpcoremodel import ArtifactType, ArtifactRole, Severity, Impact, Disposition, AlertAction, Confidence, AlertAnalyticType, ProductCategory, \ - AlertPolicyType, AlertRiskLevel, AlertStatus, CasePriority, CaseStatus, ArtifactModel, AlertModel, CaseModel - - -class AnalyzeResult(BaseModel): - """Structure for extracting C2 communication analysis result""" - original_severity: Severity = Field(description="Original alert severity", default=Severity.UNKNOWN) - new_severity: Severity = Field(description="Recommended new severity level", default=Severity.UNKNOWN) - confidence: Confidence = Field(description="Confidence level assessment", default=Confidence.UNKNOWN) - analysis_rationale: str = Field(description="Analysis process and reasons", default=None) - attack_stage: Optional[Union[str, Dict[str, Any]]] = Field(description="e.g., 'T1071 - Application Layer Protocol', 'Command and Control'", default=None) - recommended_actions: Optional[Union[str, Dict[str, Any]]] = Field(description="e.g., 'Isolate host 10.1.1.5'", default=None) - - -class AgentState(BaseAgentState): - analyze_result: AnalyzeResult = None - - -class Module(LanggraphModule): - THREAD_NUM = 2 - - def __init__(self): - super().__init__() - self.init() - - def init(self): - def alert_preprocess_node(state: AgentState): - """ - Read one alert from Redis Stream and preprocess into AlertModel. - Extract artifacts and create AlertModel with proper structure. - """ - raw_message = self.read_stream_message() - if raw_message is None: - return Command(update={}, goto=END) - - alert_raw = raw_message - - event_name = alert_raw.get("eventName", "AttachUserPolicy") - event_time = alert_raw.get("eventTime", alert_raw.get("@timestamp", "")) - rule_name = "AWS IAM Privilege Escalation via AttachUserPolicy" - - user_identity = alert_raw.get("userIdentity", {}) - principal_user = user_identity.get("userName", "unknown-user") - principal_arn = user_identity.get("arn", "unknown-arn") - principal_id = user_identity.get("principalId", "unknown-id") - access_key_id = user_identity.get("accessKeyId", "unknown-key") - - request_params = alert_raw.get("requestParameters", {}) - target_user = request_params.get("userName", "unknown-target") - policy_arn = request_params.get("policyArn", "unknown-policy") - - source_ip = alert_raw.get("sourceIPAddress", "unknown-ip") - aws_region = alert_raw.get("awsRegion", "unknown-region") - account_id = alert_raw.get("recipientAccountId", alert_raw.get("cloud.account.id", "unknown-account")) - user_agent = alert_raw.get("userAgent", "unknown-agent") - event_id = alert_raw.get("eventID", "") - - risk_score = alert_raw.get("event.risk_score", alert_raw.get("risk_score", 100)) - log_level = alert_raw.get("log.level", "critical") - message = alert_raw.get("message", "") - - severity_map = { - "critical": Severity.CRITICAL, - "high": Severity.HIGH, - "medium": Severity.MEDIUM, - "low": Severity.LOW, - "informational": Severity.INFORMATIONAL, - } - severity = severity_map.get(log_level.lower(), Severity.CRITICAL) - - event_time_formatted = event_time if event_time else get_current_time_str() - - artifacts: List[ArtifactModel] = [ - ArtifactModel( - type=ArtifactType.USER, - role=ArtifactRole.ACTOR, - value=principal_user, - name="Principal User" - ), ArtifactModel( - type=ArtifactType.OTHER, - role=ArtifactRole.ACTOR, - value=principal_arn, - name="Principal ARN" - ), ArtifactModel( - type=ArtifactType.USER, - role=ArtifactRole.TARGET, - value=target_user, - name="Target User" - ), ArtifactModel( - type=ArtifactType.IP_ADDRESS, - role=ArtifactRole.ACTOR, - value=source_ip, - name="Source IP" - ), ArtifactModel( - type=ArtifactType.OTHER, - role=ArtifactRole.RELATED, - value=policy_arn, - name="Policy ARN" - ), ArtifactModel( - type=ArtifactType.OTHER, - role=ArtifactRole.RELATED, - value=account_id, - name="AWS Account ID" - ) - ] - - if access_key_id and access_key_id != "unknown-key": - artifacts.append(ArtifactModel( - type=ArtifactType.OTHER, - role=ArtifactRole.ACTOR, - value=access_key_id, - name="Access Key ID" - )) - - correlation_config = CorrelationConfig( - rule_id=self.module_name, - time_window="24h", - keys=[principal_user, target_user, account_id] - ) - group_rule = Correlation(config=correlation_config) - correlation_uid = group_rule.generate_correlation_uid(timestamp=event_time_formatted) - - alert_model = AlertModel( - title=f"AWS IAM Privilege Escalation: {principal_user} attached {policy_arn.split('/')[-1]} to {target_user}", - src_url=f"AWS CloudTrail - Event ID: {event_id}", - severity=severity, - status=AlertStatus.NEW, - status_detail="New alert received from AWS CloudTrail - awaiting analysis", - disposition=Disposition.DETECTED, - action=AlertAction.OBSERVED, - rule_id=self.module_name, - rule_name=rule_name, - source_uid=event_id, - correlation_uid=correlation_uid, - analytic_type=AlertAnalyticType.BEHAVIORAL, - analytic_name="AWS IAM Behavioral Anomaly Detection", - analytic_desc="Detects suspicious IAM policy attachment operations that indicate privilege escalation attempts or backdoor account creation", - analytic_state=None, - product_category=ProductCategory.CLOUD, - product_name="AWS CloudTrail", - product_vendor="Amazon AWS", - product_feature="CloudTrail Logging", - first_seen_time=event_time_formatted, - last_seen_time=event_time_formatted, - desc=message or f"IAM user {principal_user} attached policy {policy_arn} to user {target_user} in account {account_id}", - data_sources=["AWS CloudTrail"], - labels=["iam-privilege-escalation", "aws-cloudtrail", f"account-{account_id}"], - raw_data=json.dumps(alert_raw), - unmapped=json.dumps({ - "userAgent": user_agent, - "awsRegion": aws_region, - "requestID": alert_raw.get("requestID", ""), - "eventVersion": alert_raw.get("eventVersion", "") - }), - tactic="T1098.003 - AWS IAM Account Manipulation", - technique="Privilege Escalation", - sub_technique="Create IAM Access Keys", - mitigation="Enable IAM access analyzer, enforce MFA, use service control policies to restrict policy attachment", - policy_name="AWS IAM Access Policy", - policy_type=AlertPolicyType.IDENTITY_POLICY, - policy_desc="IAM identity-based policy that grants permissions to AWS services", - impact=Impact.CRITICAL if severity in [Severity.CRITICAL, Severity.HIGH] else Impact.HIGH, - confidence=Confidence.HIGH, - risk_level=AlertRiskLevel.CRITICAL if severity == Severity.CRITICAL else AlertRiskLevel.HIGH, - ) - - alert_model.artifacts = artifacts - - saved_alert_row_id = Alert.create(alert_model) - alert_model.row_id = saved_alert_row_id - - self.logger.debug(f"Alert created with Rowid: {saved_alert_row_id}") - - try: - existing_cases = Case.get_by_correlation_uid(correlation_uid, lazy_load=True) - - if existing_cases and len(existing_cases) > 0: - existing_case = existing_cases[0] - self.logger.debug(f"Found existing case with correlation_uid: {correlation_uid}, Case ID: {existing_case.row_id}") - - update_case = CaseModel(alerts=[*existing_case.alerts, saved_alert_row_id], row_id=existing_case.row_id) - Case.update(update_case) - - self.logger.debug(f"Alert {saved_alert_row_id} added to existing case {existing_case.row_id}") - - else: - self.logger.debug(f"No existing case found for correlation_uid: {correlation_uid}, creating new case") - - new_case = CaseModel( - title=f"AWS IAM Privilege Escalation: {principal_user} → {target_user}", - severity=severity, - impact=Impact.CRITICAL if severity in [Severity.CRITICAL, Severity.HIGH] else Impact.HIGH, - priority=CasePriority.CRITICAL if severity == Severity.CRITICAL else CasePriority.HIGH, - confidence=Confidence.HIGH, - status=CaseStatus.NEW, - description=f"AWS IAM privilege escalation detected: {principal_user} attached {policy_arn} to {target_user}", - category=ProductCategory.CLOUD, - tags=["iam-privilege-escalation", "aws-cloudtrail", f"account-{account_id}"], - correlation_uid=correlation_uid, - alerts=[saved_alert_row_id] - ) - - case_uid = Case.create(new_case) - self.logger.debug(f"New case created with UID: {case_uid}, Alert: {saved_alert_row_id}") - - except Exception as e: - self.logger.error(f"Error creating/updating case for correlation_uid {correlation_uid}: {str(e)}") - # return Command(update={}, goto=END) - return {"alert": alert_model} - - def alert_analyze_node(state: AgentState): - """ - Analyze the alert using AI (LLM) with structured few-shot examples for AWS IAM privilege escalation detection. - Leverages threat intelligence and cloud behavior patterns. - """ - system_prompt_template = self.load_system_prompt_template("senior_cloud_security_expert") - - current_date = datetime.now().strftime("%Y-%m-%d") - system_message = system_prompt_template.format(current_date=current_date) - - few_shot_examples = [ - ] - - alert = state.alert - messages = [ - system_message, - *few_shot_examples, - HumanMessage(content=alert.model_dump_json_for_ai()), - ] - - llm_api = LLMAPI() - llm = llm_api.get_model(tag=["fast", "structured_output"]) - llm_structured = llm.with_structured_output(AnalyzeResult) - response: AnalyzeResult = llm_structured.invoke(messages) - - return {"analyze_result": response} - - def alert_output_node(state: AgentState): - """ - Save analysis result to AlertModel and persist using SIRP API. - Updates severity, confidence, and enriches with AI-generated insights. - """ - alert_model: AlertModel = state.alert - analyze_result: AnalyzeResult = state.analyze_result - - alert_model.severity = analyze_result.new_severity - alert_model.confidence = analyze_result.confidence - alert_model.summary_ai = str(analyze_result.analysis_rationale) - - if analyze_result.recommended_actions: - alert_model.remediation = str(analyze_result.recommended_actions) - - labels = list(alert_model.labels) if alert_model.labels else [] - if analyze_result.attack_stage and analyze_result.confidence in [Confidence.HIGH, Confidence.MEDIUM]: - labels.append("confirmed-privilege-escalation") - if analyze_result.new_severity in [Severity.CRITICAL, Severity.HIGH]: - labels.append("high-priority-incident") - alert_model.labels = labels - - - - update_alert_row_id = Alert.update(alert_model) - - self.logger.info( - f"AWS IAM Privilege Escalation Alert saved with RowID: {update_alert_row_id}, Severity: {analyze_result.new_severity}, Confidence: {analyze_result.confidence}") - - return {} - - workflow = StateGraph(AgentState) - - workflow.add_node("alert_preprocess_node", alert_preprocess_node) - workflow.add_node("alert_analyze_node", alert_analyze_node) - workflow.add_node("alert_output_node", alert_output_node) - - workflow.set_entry_point("alert_preprocess_node") - workflow.add_edge("alert_preprocess_node", "alert_analyze_node") - workflow.add_edge("alert_analyze_node", "alert_output_node") - workflow.set_finish_point("alert_output_node") - - self.graph: CompiledStateGraph = workflow.compile(checkpointer=self.get_checkpointer()) - return True - - -if __name__ == "__main__": - import os - import django - - os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") - django.setup() - module = Module() - module.debug_message_id = "1772179651858-0" - module.run() diff --git a/PLUGINS/ClaudeCode/skills/asp-module-creator-zh/SKILL.md b/PLUGINS/ClaudeCode/skills/asp-module-creator-zh/SKILL.md index 9f70287..4970061 100644 --- a/PLUGINS/ClaudeCode/skills/asp-module-creator-zh/SKILL.md +++ b/PLUGINS/ClaudeCode/skills/asp-module-creator-zh/SKILL.md @@ -26,6 +26,7 @@ metadata: - 模块文件名必须与 SIEM rule 名称完全一致（含大小写）——Rule 名 = Redis Stream 名 = 文件名，三者强约束，任意一处不一致框架将无法路由告警。 - 编写代码前必须先获取 raw_alert 样本，不得凭空猜测字段结构。 +- 编写代码前必须读取 `PLUGINS/SIRP/sirpcoremodel.py`，所有 enum 值只能使用该文件中实际定义的值，不得凭记忆或推断自行发明。 - 所有模块必须继承 `BaseModule` 并实现 `run()` 方法。 - SIRP 数据层级：`Case → Alert → Artifact`（三级体系）。Artifact 是调查的最小原子实体（一个 IP、一个用户名），应尽量从 raw_alert 中提取；Alert 挂在 Case 下；同类告警通过 `correlation_uid` 聚合到同一个 Case。Enrichment 是独立于三级体系之外的横切附加层，可按需挂载到 Case / Alert / Artifact 任意一级。 - 参考实现：`MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py`。 @@ -85,6 +86,8 @@ metadata: ### Step 5 — 编写模块代码 +**前置动作：** 读取 `PLUGINS/SIRP/sirpcoremodel.py`，确认所有需要用到的 enum 的合法值，再开始写代码。 + 按以下结构生成 `MODULES/.py`： ```python diff --git a/PLUGINS/Redis/redis_stream_api.py b/PLUGINS/Redis/redis_stream_api.py index 8e6bf36..31e9565 100644 --- a/PLUGINS/Redis/redis_stream_api.py +++ b/PLUGINS/Redis/redis_stream_api.py @@ -124,6 +124,23 @@ class RedisStreamAPI(object): with ThreadPoolExecutor(max_workers=1) as executor: return executor.submit(_fetch).result(timeout=timeout) + def read_stream_head_ids(self, stream_name: str, n: int, timeout: Optional[float] = None) -> List[str]: + """ + 读取指定 stream 的前 n 条消息的 ID（非阻塞）. + :param stream_name: Stream 的名称. + :param n: 要读取的消息数量. + :param timeout: 超时时间（秒），None 表示不限制. + """ + + def _fetch(): + messages = self.redis_client.xrange(stream_name, min='-', max='+', count=n) + return [msg_id.decode() if isinstance(msg_id, bytes) else msg_id for msg_id, _ in messages] + + if timeout is None: + return _fetch() + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(_fetch).result(timeout=timeout) + def read_stream_message_by_id(self, stream_name: str, message_id: str, timeout: Optional[float] = None) -> dict: """ 读取指定 ID 的消息（精确匹配，非阻塞）. diff --git a/PLUGINS/SIRP/correlation.py b/PLUGINS/SIRP/correlation.py index ea1e033..5ba9c97 100644 --- a/PLUGINS/SIRP/correlation.py +++ b/PLUGINS/SIRP/correlation.py @@ -4,7 +4,7 @@ from typing import List, Union from typing import Literal -ValidTimeWindows = Literal['10m', '30m', '1h', '2h', '4h', '8h', '12h', '24h', '7d', '30d'] +ValidTimeWindows = Literal['5m','10m', '30m', '1h', '2h', '4h', '8h', '12h', '24h', '7d', '30d'] class Correlation(object): @@ -29,8 +29,11 @@ class Correlation(object): def generate_correlation_uid(cls, rule_id: str, time_window: ValidTimeWindows = "24h", - timestamp: datetime = datetime.now(timezone.utc), + timestamp: datetime = None, keys: List[Union[str, None]] = None) -> str: + if timestamp is None: + timestamp = datetime.now(timezone.utc) + keys = keys or [] time_bucket = cls._get_time_bucket(timestamp, time_window) key_parts = [rule_id, time_bucket]