resource "aws_iam_role" "step_function_role" {
  name = "guardduty-malware-response-role-${var.environment}"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "states.amazonaws.com"
        }
      }
    ]
  })
}

resource "aws_iam_policy" "step_function_policy" {
  name = "guardduty-malware-response-policy-${var.environment}"

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect = "Allow"
        Action = [
          "s3:GetObject",
          "s3:DeleteObject",
          "s3:GetObjectTagging",
          "s3:ListBucket"
        ]
        Resource = "*"
      },
      {
        Effect = "Allow"
        Action = [
          "s3:PutObject",
          "s3:PutObjectTagging"
        ]
        Resource = "arn:aws:s3:::${module.quarantine_bucket.s3_bucket_name_output}/*"
      },
      {
        Effect = "Allow"
        Action = [
          "sns:Publish"
        ]
        Resource = aws_sns_topic.guardduty_notifications.arn
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "attach" {
  role       = aws_iam_role.step_function_role.name
  policy_arn = aws_iam_policy.step_function_policy.arn
}

resource "aws_sfn_state_machine" "sfn_state_machine" {
  name     = "guardduty-malware-response-${var.environment}"
  role_arn = aws_iam_role.step_function_role.arn

  definition = <<EOF
{
  "Comment": "Move S3 object to quarantine based on GuardDuty Malware Protection Scan Result",
  "StartAt": "Check Scan Result Status",
  "States": {
    "Check Scan Result Status": {
      "Type": "Choice",
      "Choices": [
        {
          "And": [
            {
              "Variable": "$.detail.scanResultDetails.scanResultStatus",
              "StringEquals": "THREATS_FOUND"
            },
            {
              "Variable": "$.detail.scanStatus",
              "StringEquals": "COMPLETED"
            }
          ],
          "Next": "Notify Admin"
        }
      ],
      "Default": "Finish"
    },
    "Notify Admin": {
      "Type": "Task",
      "Resource": "arn:aws:states:::sns:publish",
      "Parameters": {
        "TopicArn": "${aws_sns_topic.guardduty_notifications.arn}",
        "Message.$": "States.JsonToString($)",
        "Subject": "Malware Detected in S3"
      },
      "ResultPath": "$.snsResult",
      "Next": "Copy Object to Quarantine"
    },
    "Copy Object to Quarantine": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:s3:copyObject",
      "Parameters": {
        "Bucket": "${module.quarantine_bucket.s3_bucket_name_output}",
        "CopySource.$": "States.Format('{}/{}', $.detail.s3ObjectDetails.bucketName, $.detail.s3ObjectDetails.objectKey)",
        "Key.$": "States.Format('{}/{}', $.detail.s3ObjectDetails.bucketName, $.detail.s3ObjectDetails.objectKey)"
      },
      "ResultPath": "$.copyResult",
      "Next": "Delete Object"
    },
    "Delete Object": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:s3:deleteObject",
      "Parameters": {
        "Bucket.$": "$.detail.s3ObjectDetails.bucketName",
        "Key.$": "$.detail.s3ObjectDetails.objectKey"
      },
      "ResultPath": "$.deleteResult",
      "Next": "Finish"
    },
    "Finish": {
      "Type": "Pass",
      "End": true
    }
  }
}
EOF
}