| 1 | from enum import Enum |
| 2 | from typing import Any |
| 3 | from typing import Dict |
| 4 | from typing import List |
| 5 | from typing import Literal |
| 6 | from typing import Optional |
| 7 | |
| 8 | from fastapi import HTTPException |
| 9 | from pydantic import BaseModel |
| 10 | from pydantic import ConfigDict |
| 11 | from pydantic import Field |
| 12 | from pydantic import field_validator |
| 13 | from pydantic import model_validator |
| 14 | |
| 15 | |
| 16 | class ActiveResponsesSupported(Enum): |
| 17 | WINDOWS_FIREWALL = "Block or unblock any outbound traffic to the defined IP address via the Windows Firewall" |
| 18 | # Add more active responses here as needed |
| 19 | |
| 20 | |
| 21 | class ActiveResponse(BaseModel): |
| 22 | name: str |
| 23 | description: str |
| 24 | |
| 25 | |
| 26 | class ActiveResponsesSupportedResponse(BaseModel): |
| 27 | supported_active_responses: List[ActiveResponse] |
| 28 | success: bool |
| 29 | message: str |
| 30 | |
| 31 | |
| 32 | class ActiveResponseDetails(BaseModel): |
| 33 | name: str |
| 34 | description: str |
| 35 | markdown_content: str |
| 36 | # TODO[pydantic]: The following keys were removed: `json_encoders`. |
| 37 | # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information. |
| 38 | model_config = ConfigDict(json_encoders={str: lambda v: v.encode("utf-8", "ignore").decode("utf-8")}) |
| 39 | |
| 40 | |
| 41 | class ActiveResponseDetailsResponse(BaseModel): |
| 42 | success: bool |
| 43 | message: str |
| 44 | active_response: ActiveResponseDetails |
| 45 | |
| 46 | |
| 47 | # ! Invoke Active Response ! # |
| 48 | class AlertAction(str, Enum): |
| 49 | unblock = "unblock" |
| 50 | block = "block" |
| 51 | sysmon_config_reload = "sysmon_config_reload" |
| 52 | |
| 53 | |
| 54 | class BaseModelWithEnum(BaseModel): |
| 55 | model_config = ConfigDict(use_enum_values=True) |
| 56 | |
| 57 | |
| 58 | class WindowsFirewallAlert(BaseModelWithEnum): |
| 59 | action: AlertAction |
| 60 | ip: str |
| 61 | |
| 62 | |
| 63 | class LinuxFirewallAlert(BaseModelWithEnum): |
| 64 | action: AlertAction |
| 65 | ip: str |
| 66 | |
| 67 | |
| 68 | class SysmonConfigReloadAlert(BaseModelWithEnum): |
| 69 | action: Literal[AlertAction.sysmon_config_reload] = AlertAction.sysmon_config_reload |
| 70 | |
| 71 | |
| 72 | class ActiveResponseCommand(str, Enum): |
| 73 | windows_firewall = "windows_firewall" |
| 74 | linux_firewall = "linux_firewall" |
| 75 | sysmon_config_reload = "sysmon_config_reload" |
| 76 | |
| 77 | @classmethod |
| 78 | def _missing_(cls, value): |
| 79 | for member in cls: |
| 80 | if member.name == value: |
| 81 | return member |
| 82 | |
| 83 | for active_response in ActiveResponsesSupported: |
| 84 | if active_response.name.lower() == value.lower(): |
| 85 | return cls[f"{value}0"] |
| 86 | |
| 87 | raise HTTPException( |
| 88 | status_code=400, |
| 89 | detail=f"Invalid command: {value}, must be one of {', '.join([member.name for member in cls])}", |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | class ParamsModel(BaseModel): |
| 94 | wait_for_complete: bool |
| 95 | agents_list: Optional[List[str]] = None |
| 96 | |
| 97 | @field_validator("agents_list", mode="before") |
| 98 | @classmethod |
| 99 | def check_agents_list(cls, v): |
| 100 | if v == ["*"]: |
| 101 | return [] |
| 102 | return v |
| 103 | |
| 104 | |
| 105 | class InvokeActiveResponseRequest(BaseModel): |
| 106 | endpoint: Literal["/active-response"] = "/active-response" |
| 107 | arguments: list[str] = Field(default_factory=list) |
| 108 | command: ActiveResponseCommand |
| 109 | custom: Literal[True] = True |
| 110 | alert: Dict[str, Any] |
| 111 | params: ParamsModel |
| 112 | |
| 113 | @model_validator(mode="before") |
| 114 | @classmethod |
| 115 | def create_alert(cls, values): |
| 116 | command = values.get("command") |
| 117 | alert = values.get("alert") |
| 118 | if command == ActiveResponseCommand.windows_firewall: |
| 119 | values["alert"] = WindowsFirewallAlert(**alert) |
| 120 | elif command == ActiveResponseCommand.linux_firewall: |
| 121 | values["alert"] = LinuxFirewallAlert(**alert) |
| 122 | elif command == ActiveResponseCommand.sysmon_config_reload: |
| 123 | values["alert"] = SysmonConfigReloadAlert(**alert) |
| 124 | else: |
| 125 | raise HTTPException(status_code=400, detail="Invalid command for alert") |
| 126 | |
| 127 | return values |
| 128 | |
| 129 | model_config = ConfigDict( |
| 130 | json_schema_extra={ |
| 131 | "example": { |
| 132 | "endpoint": "/active-response", |
| 133 | "arguments": [], |
| 134 | "command": "windows_firewall", |
| 135 | "custom": True, |
| 136 | "alert": {"action": "block", "ip": "1.1.1.1"}, |
| 137 | "params": {"wait_for_complete": True, "agents_list": ["032"]}, |
| 138 | }, |
| 139 | }, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | class InvokeActiveResponseResponse(BaseModel): |
| 144 | success: bool |
| 145 | message: str |