from pypanther import LogType, Rule, RuleTest, Severity, panther_managed
from pypanther.helpers.panther_base_helpers import deep_get

wiz_alert_passthrough_tests: list[RuleTest] = [
    RuleTest(
        name="Open Alert",
        expected_result=True,
        log={
            "createdAt": "2024-06-04 02:28:06.763277000",
            "entitySnapshot": {
                "cloudProviderURL": "",
                "externalId": "someExternalId",
                "id": "12345",
                "name": "someName",
                "nativeType": "",
                "providerId": "someProviderId",
                "region": "",
                "resourceGroupExternalId": "",
                "subscriptionExternalId": "",
                "subscriptionName": "",
                "tags": {},
                "type": "DATA_FINDING",
            },
            "id": "54321",
            "notes": [],
            "projects": [
                {
                    "businessUnit": "",
                    "id": "45678",
                    "name": "Project 2",
                    "riskProfile": {"businessImpact": "MBI"},
                    "slug": "project-2",
                },
            ],
            "serviceTickets": [],
            "severity": "HIGH",
            "sourceRule": {
                "__typename": "Control",
                "controlDescription": "Alert Description",
                "id": "12345",
                "name": "Alert Name",
                "resolutionRecommendation": "Alert Resolution Recommendation",
                "securitySubCategories": [
                    {
                        "category": {"framework": {"name": "Wiz for Risk Assessment"}, "name": "High Profile Threats"},
                        "title": "High-profile vulnerability exploited in the wild",
                    },
                    {
                        "category": {"framework": {"name": "MITRE ATT&CK Matrix"}, "name": "TA0001 Initial Access"},
                        "title": "T1190 Exploit Public-Facing Application",
                    },
                ],
            },
            "status": "OPEN",
            "statusChangedAt": "2024-06-04 02:28:06.597355000",
            "type": "TOXIC_COMBINATION",
            "updatedAt": "2024-06-04 02:28:06.763277000",
        },
    ),
    RuleTest(
        name="Resolved Alert",
        expected_result=False,
        log={
            "createdAt": "2024-06-04 02:28:06.763277000",
            "entitySnapshot": {
                "cloudProviderURL": "",
                "externalId": "someExternalId",
                "id": "12345",
                "name": "someName",
                "nativeType": "",
                "providerId": "someProviderId",
                "region": "",
                "resourceGroupExternalId": "",
                "subscriptionExternalId": "",
                "subscriptionName": "",
                "tags": {},
                "type": "DATA_FINDING",
            },
            "id": "54321",
            "notes": [],
            "projects": [
                {
                    "businessUnit": "",
                    "id": "45678",
                    "name": "Project 2",
                    "riskProfile": {"businessImpact": "MBI"},
                    "slug": "project-2",
                },
            ],
            "serviceTickets": [],
            "severity": "HIGH",
            "sourceRule": {
                "__typename": "Control",
                "controlDescription": "Alert Description",
                "id": "12345",
                "name": "Alert Name",
                "resolutionRecommendation": "Alert Resolution Recommendation",
            },
            "status": "RESOLVED",
            "statusChangedAt": "2024-06-04 02:28:06.597355000",
            "type": "TOXIC_COMBINATION",
            "updatedAt": "2024-06-04 02:28:06.763277000",
        },
    ),
]


@panther_managed
class WizAlertPassthrough(Rule):
    id = "Wiz.Alert.Passthrough-prototype"
    default_description = "This rule enriches and contextualizes security alerts generated by Wiz."
    display_name = "Wiz Alert Passthrough Rule"
    default_runbook = "Review the Wiz alert details to determine what malicious behavior was detected, and whether or not it was blocked."
    default_reference = "https://www.wiz.io/product"
    default_severity = Severity.MEDIUM
    log_types = [LogType.WIZ_ISSUES]
    tests = wiz_alert_passthrough_tests

    def rule(self, event):
        return event.get("status") == "OPEN"

    def title(self, event):
        return f"[Wiz Alert]: {event.deep_get('sourceRule', 'name', default='ALERT_DESCRIPTION_NOT_FOUND')}"

    def severity(self, event):
        return event.get("severity")

    def dedup(self, event):
        return event.get("id")

    def description(self, event):
        return event.deep_get("sourceRule", "controlDescription", default="<DESCRIPTION_NOT_FOUND>")

    def runbook(self, event):
        return event.deep_get("sourceRule", "resolutionRecommendation", default="<RECOMMENDATION_NOT_FOUND>")

    def alert_context(self, event):
        security_subcategories = event.deep_get("sourceRule", "securitySubCategories", default=[{}])
        return {
            "id": event.get("id", "<ID_NOT_FOUND>"),
            "type": event.get("type", "<TYPE_NOT_FOUND>"),
            "entity_snapshot": event.get("entitySnapshot", {}),
            "mitre_attack_categories": [
                subcategory
                for subcategory in security_subcategories
                if deep_get(subcategory, "category", "framework", "name") == "MITRE ATT&CK Matrix"
            ],
        }
