#!/usr/bin/env python3
"""AWS cloud incident containment agent.

Automates incident containment procedures in AWS environments including
EC2 instance isolation, IAM credential revocation, security group lockdown,
S3 bucket access restriction, and forensic snapshot creation using boto3.
"""
import argparse
import json
import os
import sys
from datetime import datetime, timezone

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    print("[!] 'boto3' library required: pip install boto3", file=sys.stderr)
    sys.exit(1)


def get_session(profile=None, region=None):
    """Create a boto3 session."""
    kwargs = {}
    if profile:
        kwargs["profile_name"] = profile
    if region:
        kwargs["region_name"] = region
    return boto3.Session(**kwargs)


def isolate_ec2_instance(session, instance_id, vpc_id=None):
    """Isolate an EC2 instance by replacing security groups with a deny-all SG."""
    ec2 = session.client("ec2")
    findings = []
    print(f"[*] Isolating EC2 instance: {instance_id}")

    # Get current instance details
    try:
        resp = ec2.describe_instances(InstanceIds=[instance_id])
        instance = resp["Reservations"][0]["Instances"][0]
        current_sgs = [sg["GroupId"] for sg in instance.get("SecurityGroups", [])]
        instance_vpc = instance.get("VpcId", vpc_id)
        findings.append({"action": "describe_instance", "status": "OK",
                         "detail": f"Current SGs: {current_sgs}"})
    except ClientError as e:
        findings.append({"action": "describe_instance", "status": "FAIL",
                         "severity": "CRITICAL", "detail": str(e)})
        return findings

    # Create or find isolation security group
    isolation_sg_name = f"incident-isolation-{instance_id[:8]}"
    isolation_sg_id = None
    try:
        existing = ec2.describe_security_groups(
            Filters=[{"Name": "group-name", "Values": [isolation_sg_name]},
                     {"Name": "vpc-id", "Values": [instance_vpc]}]
        )
        if existing["SecurityGroups"]:
            isolation_sg_id = existing["SecurityGroups"][0]["GroupId"]
        else:
            resp = ec2.create_security_group(
                GroupName=isolation_sg_name,
                Description=f"Incident isolation SG for {instance_id}",
                VpcId=instance_vpc,
            )
            isolation_sg_id = resp["GroupId"]
            # Revoke default egress rule
            ec2.revoke_security_group_egress(
                GroupId=isolation_sg_id,
                IpPermissions=[{"IpProtocol": "-1", "IpRanges": [{"CidrIp": "0.0.0.0/0"}]}],
            )
        findings.append({"action": "create_isolation_sg", "status": "OK",
                         "detail": f"SG: {isolation_sg_id} (no ingress, no egress)"})
    except ClientError as e:
        findings.append({"action": "create_isolation_sg", "status": "FAIL",
                         "severity": "HIGH", "detail": str(e)})
        return findings

    # Replace security groups with isolation SG
    try:
        ec2.modify_instance_attribute(
            InstanceId=instance_id,
            Groups=[isolation_sg_id],
        )
        findings.append({"action": "apply_isolation_sg", "status": "OK",
                         "detail": f"Replaced {current_sgs} with [{isolation_sg_id}]"})
    except ClientError as e:
        findings.append({"action": "apply_isolation_sg", "status": "FAIL",
                         "severity": "CRITICAL", "detail": str(e)})

    # Tag instance as contained
    try:
        ec2.create_tags(
            Resources=[instance_id],
            Tags=[
                {"Key": "IncidentStatus", "Value": "CONTAINED"},
                {"Key": "ContainmentTime", "Value": datetime.now(timezone.utc).isoformat()},
                {"Key": "OriginalSecurityGroups", "Value": ",".join(current_sgs)},
            ],
        )
        findings.append({"action": "tag_instance", "status": "OK"})
    except ClientError as e:
        findings.append({"action": "tag_instance", "status": "FAIL", "detail": str(e)})

    return findings


def create_forensic_snapshot(session, instance_id):
    """Create EBS snapshots for forensic preservation."""
    ec2 = session.client("ec2")
    findings = []
    print(f"[*] Creating forensic snapshots for: {instance_id}")

    try:
        resp = ec2.describe_instances(InstanceIds=[instance_id])
        instance = resp["Reservations"][0]["Instances"][0]
        volumes = []
        for mapping in instance.get("BlockDeviceMappings", []):
            vol_id = mapping.get("Ebs", {}).get("VolumeId")
            if vol_id:
                volumes.append((mapping["DeviceName"], vol_id))
    except ClientError as e:
        findings.append({"action": "describe_volumes", "status": "FAIL", "detail": str(e)})
        return findings

    for device_name, vol_id in volumes:
        try:
            snap = ec2.create_snapshot(
                VolumeId=vol_id,
                Description=f"Forensic snapshot - {instance_id} {device_name} - "
                            f"{datetime.now(timezone.utc).isoformat()}",
                TagSpecifications=[{
                    "ResourceType": "snapshot",
                    "Tags": [
                        {"Key": "Purpose", "Value": "Forensic-Preservation"},
                        {"Key": "SourceInstance", "Value": instance_id},
                        {"Key": "SourceVolume", "Value": vol_id},
                        {"Key": "CreatedBy", "Value": "incident-containment-agent"},
                    ],
                }],
            )
            findings.append({
                "action": "create_snapshot",
                "status": "OK",
                "detail": f"{vol_id} ({device_name}) -> {snap['SnapshotId']}",
            })
        except ClientError as e:
            findings.append({"action": "create_snapshot", "status": "FAIL",
                             "detail": f"{vol_id}: {e}"})

    return findings


def revoke_iam_credentials(session, username):
    """Revoke all IAM credentials for a compromised user."""
    iam = session.client("iam")
    findings = []
    print(f"[*] Revoking credentials for IAM user: {username}")

    # Deactivate access keys
    try:
        keys = iam.list_access_keys(UserName=username)
        for key in keys.get("AccessKeyMetadata", []):
            key_id = key["AccessKeyId"]
            iam.update_access_key(
                UserName=username, AccessKeyId=key_id, Status="Inactive"
            )
            findings.append({"action": "deactivate_access_key", "status": "OK",
                             "detail": f"Key {key_id[:8]}... deactivated"})
    except ClientError as e:
        findings.append({"action": "deactivate_access_keys", "status": "FAIL", "detail": str(e)})

    # Invalidate console session by attaching deny-all inline policy
    deny_policy = json.dumps({
        "Version": "2012-10-17",
        "Statement": [{"Effect": "Deny", "Action": "*", "Resource": "*"}],
    })
    try:
        iam.put_user_policy(
            UserName=username,
            PolicyName="IncidentDenyAll",
            PolicyDocument=deny_policy,
        )
        findings.append({"action": "attach_deny_policy", "status": "OK",
                         "detail": "Deny-all policy attached"})
    except ClientError as e:
        findings.append({"action": "attach_deny_policy", "status": "FAIL", "detail": str(e)})

    # Delete login profile (console access)
    try:
        iam.delete_login_profile(UserName=username)
        findings.append({"action": "delete_console_access", "status": "OK"})
    except iam.exceptions.NoSuchEntityException:
        findings.append({"action": "delete_console_access", "status": "SKIP",
                         "detail": "No console access configured"})
    except ClientError as e:
        findings.append({"action": "delete_console_access", "status": "FAIL", "detail": str(e)})

    return findings


def restrict_s3_bucket(session, bucket_name):
    """Restrict S3 bucket access during incident containment."""
    s3 = session.client("s3")
    findings = []
    print(f"[*] Restricting S3 bucket: {bucket_name}")

    # Block public access
    try:
        s3.put_public_access_block(
            Bucket=bucket_name,
            PublicAccessBlockConfiguration={
                "BlockPublicAcls": True,
                "IgnorePublicAcls": True,
                "BlockPublicPolicy": True,
                "RestrictPublicBuckets": True,
            },
        )
        findings.append({"action": "block_public_access", "status": "OK"})
    except ClientError as e:
        findings.append({"action": "block_public_access", "status": "FAIL", "detail": str(e)})

    # Enable versioning (preserve evidence)
    try:
        s3.put_bucket_versioning(
            Bucket=bucket_name,
            VersioningConfiguration={"Status": "Enabled"},
        )
        findings.append({"action": "enable_versioning", "status": "OK"})
    except ClientError as e:
        findings.append({"action": "enable_versioning", "status": "FAIL", "detail": str(e)})

    return findings


def format_summary(all_actions):
    """Print containment summary."""
    print(f"\n{'='*60}")
    print(f"  Cloud Incident Containment Report")
    print(f"{'='*60}")

    success = sum(1 for a in all_actions if a.get("status") == "OK")
    failed = sum(1 for a in all_actions if a.get("status") == "FAIL")
    print(f"  Actions  : {len(all_actions)}")
    print(f"  Success  : {success}")
    print(f"  Failed   : {failed}")

    print(f"\n  Actions Taken:")
    for a in all_actions:
        icon = "OK" if a["status"] == "OK" else "!!" if a["status"] == "FAIL" else "--"
        print(f"    [{icon}] {a['action']:30s} {a.get('detail', '')[:50]}")


def main():
    parser = argparse.ArgumentParser(
        description="AWS cloud incident containment agent"
    )
    sub = parser.add_subparsers(dest="command")

    p_iso = sub.add_parser("isolate", help="Isolate EC2 instance")
    p_iso.add_argument("--instance-id", required=True)

    p_snap = sub.add_parser("snapshot", help="Create forensic snapshots")
    p_snap.add_argument("--instance-id", required=True)

    p_iam = sub.add_parser("revoke-iam", help="Revoke IAM user credentials")
    p_iam.add_argument("--username", required=True)

    p_s3 = sub.add_parser("restrict-s3", help="Restrict S3 bucket")
    p_s3.add_argument("--bucket", required=True)

    p_full = sub.add_parser("full-contain", help="Full containment: isolate + snapshot + IAM")
    p_full.add_argument("--instance-id", required=True)
    p_full.add_argument("--username", help="IAM user to revoke")
    p_full.add_argument("--bucket", help="S3 bucket to restrict")

    parser.add_argument("--profile", help="AWS CLI profile")
    parser.add_argument("--region", help="AWS region")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    session = get_session(args.profile, args.region)
    all_actions = []

    if args.command == "isolate":
        all_actions.extend(isolate_ec2_instance(session, args.instance_id))
    elif args.command == "snapshot":
        all_actions.extend(create_forensic_snapshot(session, args.instance_id))
    elif args.command == "revoke-iam":
        all_actions.extend(revoke_iam_credentials(session, args.username))
    elif args.command == "restrict-s3":
        all_actions.extend(restrict_s3_bucket(session, args.bucket))
    elif args.command == "full-contain":
        all_actions.extend(isolate_ec2_instance(session, args.instance_id))
        all_actions.extend(create_forensic_snapshot(session, args.instance_id))
        if args.username:
            all_actions.extend(revoke_iam_credentials(session, args.username))
        if args.bucket:
            all_actions.extend(restrict_s3_bucket(session, args.bucket))

    format_summary(all_actions)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "AWS Incident Containment",
        "command": args.command,
        "actions": all_actions,
        "success_count": sum(1 for a in all_actions if a["status"] == "OK"),
        "fail_count": sum(1 for a in all_actions if a["status"] == "FAIL"),
    }

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()

"""
Cloud Incident Containment Automation Script
Provides containment procedures for AWS, Azure, and GCP environments.
"""

import json
import os
from datetime import datetime, timezone
from pathlib import Path


class CloudContainmentManager:
    """Manages cloud incident containment across AWS, Azure, and GCP."""

    def __init__(self, output_dir="cloud_containment_results"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.actions_log = []
        self.contained_resources = []

    def log_action(self, platform, action, resource, status, details=""):
        """Log a containment action for audit trail."""
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "platform": platform,
            "action": action,
            "resource": resource,
            "status": status,
            "details": details,
        }
        self.actions_log.append(entry)
        print(f"[{platform}] {action}: {resource} -> {status}")

    def generate_aws_containment_commands(self, incident_type, resource_details):
        """Generate AWS CLI containment commands based on incident type."""
        commands = []

        if incident_type == "credential_compromise":
            user = resource_details.get("username", "UNKNOWN")
            commands.extend([
                {
                    "step": 1,
                    "description": "Disable all access keys",
                    "command": f"aws iam list-access-keys --user-name {user} --query 'AccessKeyMetadata[].AccessKeyId' --output text | xargs -I {{}} aws iam update-access-key --user-name {user} --access-key-id {{}} --status Inactive",
                },
                {
                    "step": 2,
                    "description": "Attach deny-all policy",
                    "command": f'aws iam put-user-policy --user-name {user} --policy-name EmergencyDenyAll --policy-document \'{{"Version":"2012-10-17","Statement":[{{"Effect":"Deny","Action":"*","Resource":"*"}}]}}\'',
                },
                {
                    "step": 3,
                    "description": "Delete console password",
                    "command": f"aws iam delete-login-profile --user-name {user}",
                },
                {
                    "step": 4,
                    "description": "Deactivate MFA devices",
                    "command": f"aws iam list-mfa-devices --user-name {user}",
                },
            ])

        elif incident_type == "ec2_compromise":
            instance_id = resource_details.get("instance_id", "i-UNKNOWN")
            vpc_id = resource_details.get("vpc_id", "vpc-UNKNOWN")
            volume_ids = resource_details.get("volume_ids", [])

            commands.extend([
                {
                    "step": 1,
                    "description": "Create forensic snapshots of all volumes",
                    "command": " && ".join([
                        f'aws ec2 create-snapshot --volume-id {vol} --description "Forensic-IR-$(date +%Y%m%d)" --tag-specifications \'ResourceType=snapshot,Tags=[{{Key=IR-Case,Value=active}}]\''
                        for vol in volume_ids
                    ]) if volume_ids else f"aws ec2 describe-instance-attribute --instance-id {instance_id} --attribute blockDeviceMapping",
                },
                {
                    "step": 2,
                    "description": "Create quarantine security group",
                    "command": f"aws ec2 create-security-group --group-name quarantine-$(date +%s) --description 'IR Quarantine' --vpc-id {vpc_id}",
                },
                {
                    "step": 3,
                    "description": "Remove default egress rule from quarantine SG",
                    "command": "aws ec2 revoke-security-group-egress --group-id SG_ID --ip-permissions '[{\"IpProtocol\":\"-1\",\"FromPort\":-1,\"ToPort\":-1,\"IpRanges\":[{\"CidrIp\":\"0.0.0.0/0\"}]}]'",
                },
                {
                    "step": 4,
                    "description": "Apply quarantine SG to instance",
                    "command": f"aws ec2 modify-instance-attribute --instance-id {instance_id} --groups SG_ID",
                },
                {
                    "step": 5,
                    "description": "Tag instance as contained",
                    "command": f"aws ec2 create-tags --resources {instance_id} --tags Key=IR-Status,Value=Contained Key=IR-Date,Value=$(date +%Y%m%d)",
                },
            ])

        elif incident_type == "s3_exposure":
            bucket = resource_details.get("bucket_name", "UNKNOWN")
            commands.extend([
                {
                    "step": 1,
                    "description": "Block all public access",
                    "command": f"aws s3api put-public-access-block --bucket {bucket} --public-access-block-configuration BlockPublicAcls=true,IgnorePublicAcls=true,BlockPublicPolicy=true,RestrictPublicBuckets=true",
                },
                {
                    "step": 2,
                    "description": "Enable versioning for evidence",
                    "command": f"aws s3api put-bucket-versioning --bucket {bucket} --versioning-configuration Status=Enabled",
                },
                {
                    "step": 3,
                    "description": "Enable server access logging",
                    "command": f"aws s3api put-bucket-logging --bucket {bucket} --bucket-logging-status '{{\"LoggingEnabled\":{{\"TargetBucket\":\"ir-logs-bucket\",\"TargetPrefix\":\"{bucket}/\"}}}}'",
                },
            ])

        elif incident_type == "lambda_compromise":
            function_name = resource_details.get("function_name", "UNKNOWN")
            commands.extend([
                {
                    "step": 1,
                    "description": "Stop function invocations",
                    "command": f"aws lambda put-function-concurrency --function-name {function_name} --reserved-concurrent-executions 0",
                },
                {
                    "step": 2,
                    "description": "Remove event source mappings",
                    "command": f"aws lambda list-event-source-mappings --function-name {function_name}",
                },
                {
                    "step": 3,
                    "description": "Capture function configuration",
                    "command": f"aws lambda get-function --function-name {function_name} > lambda_evidence_{function_name}.json",
                },
            ])

        return commands

    def generate_azure_containment_commands(self, incident_type, resource_details):
        """Generate Azure CLI/PowerShell containment commands."""
        commands = []

        if incident_type == "identity_compromise":
            user_id = resource_details.get("user_id", "UNKNOWN")
            commands.extend([
                {
                    "step": 1,
                    "description": "Revoke all refresh tokens",
                    "command": f"az rest --method POST --url 'https://graph.microsoft.com/v1.0/users/{user_id}/revokeSignInSessions'",
                },
                {
                    "step": 2,
                    "description": "Disable user account",
                    "command": f"az ad user update --id {user_id} --account-enabled false",
                },
                {
                    "step": 3,
                    "description": "Force password reset",
                    "command": f"az ad user update --id {user_id} --force-change-password-next-sign-in true",
                },
            ])

        elif incident_type == "vm_compromise":
            vm_name = resource_details.get("vm_name", "UNKNOWN")
            rg = resource_details.get("resource_group", "UNKNOWN")
            commands.extend([
                {
                    "step": 1,
                    "description": "Create disk snapshot",
                    "command": f"az snapshot create -g {rg} -n forensic-snap-$(date +%s) --source $(az vm show -g {rg} -n {vm_name} --query 'storageProfile.osDisk.managedDisk.id' -o tsv)",
                },
                {
                    "step": 2,
                    "description": "Create quarantine NSG",
                    "command": f"az network nsg create -g {rg} -n quarantine-nsg",
                },
                {
                    "step": 3,
                    "description": "Add deny-all inbound rule",
                    "command": f"az network nsg rule create -g {rg} --nsg-name quarantine-nsg -n DenyAllInbound --priority 100 --direction Inbound --access Deny --protocol '*' --source-address-prefix '*' --destination-address-prefix '*'",
                },
                {
                    "step": 4,
                    "description": "Add deny-all outbound rule",
                    "command": f"az network nsg rule create -g {rg} --nsg-name quarantine-nsg -n DenyAllOutbound --priority 100 --direction Outbound --access Deny --protocol '*' --source-address-prefix '*' --destination-address-prefix '*'",
                },
            ])

        return commands

    def generate_containment_playbook(self, case_id, platform, incident_type, resource_details):
        """Generate complete containment playbook."""
        if platform == "aws":
            commands = self.generate_aws_containment_commands(incident_type, resource_details)
        elif platform == "azure":
            commands = self.generate_azure_containment_commands(incident_type, resource_details)
        else:
            commands = []

        playbook = {
            "case_id": case_id,
            "platform": platform,
            "incident_type": incident_type,
            "generated": datetime.now(timezone.utc).isoformat(),
            "resource_details": resource_details,
            "pre_containment_checklist": [
                "Verify incident is confirmed (not false positive)",
                "Notify incident commander",
                "Create forensic snapshots/backups BEFORE containment",
                "Document current state of affected resources",
                "Identify break-glass credentials for IR team",
            ],
            "containment_commands": commands,
            "post_containment_checklist": [
                "Verify containment is effective (test connectivity)",
                "Confirm forensic evidence is preserved",
                "Update incident ticket with containment actions",
                "Notify stakeholders of containment status",
                "Begin eradication planning",
            ],
        }

        playbook_file = self.output_dir / f"containment_playbook_{case_id}.json"
        with open(playbook_file, "w") as f:
            json.dump(playbook, f, indent=2)

        print(f"[+] Containment playbook generated: {playbook_file}")
        print(f"    Platform: {platform}")
        print(f"    Type: {incident_type}")
        print(f"    Steps: {len(commands)}")
        return playbook

    def generate_report(self):
        """Generate containment actions report."""
        report = {
            "title": "Cloud Incident Containment Report",
            "generated": datetime.now(timezone.utc).isoformat(),
            "actions_taken": self.actions_log,
            "resources_contained": self.contained_resources,
            "total_actions": len(self.actions_log),
        }

        report_file = self.output_dir / "containment_report.json"
        with open(report_file, "w") as f:
            json.dump(report, f, indent=2)

        print(f"[+] Containment report: {report_file}")
        return report


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Cloud Incident Containment Tool")
    parser.add_argument("--platform", choices=["aws", "azure", "gcp"], required=True)
    parser.add_argument("--incident-type", choices=[
        "credential_compromise", "ec2_compromise", "s3_exposure",
        "lambda_compromise", "identity_compromise", "vm_compromise",
    ], required=True)
    parser.add_argument("--case-id", default="IR-2025-001")
    parser.add_argument("--resource-json", help="JSON file with resource details")
    parser.add_argument("-o", "--output", default="cloud_containment_results")

    args = parser.parse_args()

    resource_details = {}
    if args.resource_json:
        with open(args.resource_json) as f:
            resource_details = json.load(f)

    manager = CloudContainmentManager(output_dir=args.output)
    manager.generate_containment_playbook(
        args.case_id, args.platform, args.incident_type, resource_details
    )


if __name__ == "__main__":
    main()