Trend Cloud One File Storage SecurityからGuardDuty Malware Protection for S3への移行を試してみた

import boto3
import json
import os
import logging
from datetime import datetime, timezone
from botocore.exceptions import ClientError

# ロガーの設定
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 環境変数から設定を取得
SNS_TOPIC_ARN = os.environ.get('SNS_TOPIC_ARN', '')
TAG_KEY = os.environ.get('TAG_KEY', 'GuardDutyMalwareScanStatus')
TAG_VALUE_THREAT = os.environ.get('TAG_VALUE_THREAT', 'THREATS_FOUND')
DELETE_ORIGINAL = os.environ.get('DELETE_ORIGINAL', 'true').lower() == 'true'

# AWSクライアントの初期化
s3_client = boto3.client('s3')
sns_client = boto3.client('sns') if SNS_TOPIC_ARN else None

def send_notification(message, subject):
    """SNS通知を送信"""
    if not sns_client or not SNS_TOPIC_ARN:
        logger.info("SNS通知はスキップされました（SNS_TOPIC_ARNが未設定）")
        return

    try:
        sns_client.publish(
            TopicArn=SNS_TOPIC_ARN,
            Subject=subject,
            Message=message
        )
        logger.info(f"SNS通知を送信しました: {subject}")
    except ClientError as e:
        logger.error(f"SNS通知の送信に失敗: {e}")

def get_object_metadata(bucket_name, object_key):
    """オブジェクトのメタデータを取得"""
    try:
        response = s3_client.head_object(Bucket=bucket_name, Key=object_key)
        return {
            'ContentType': response.get('ContentType', 'unknown'),
            'ContentLength': response.get('ContentLength', 0),
            'LastModified': response.get('LastModified', '').isoformat() if response.get('LastModified') else '',
            'ETag': response.get('ETag', ''),
            'ServerSideEncryption': response.get('ServerSideEncryption', 'None'),
            'StorageClass': response.get('StorageClass', 'STANDARD')
        }
    except ClientError as e:
        logger.error(f"メタデータ取得エラー: {e}")
        return {}

def quarantine_object(bucket_name, object_key, metadata, quarantine_bucket):
    """オブジェクトを隔離バケットに移動"""
    timestamp = datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')
    quarantine_key = f"{timestamp}/{bucket_name}/{object_key}"

    try:
        # メタデータをタグとして準備
        tagging = {
            'TagSet': [
                {'Key': 'OriginalBucket', 'Value': bucket_name},
                {'Key': 'OriginalKey', 'Value': object_key},
                {'Key': 'QuarantineDate', 'Value': timestamp},
                {'Key': 'MalwareStatus', 'Value': TAG_VALUE_THREAT},
                {'Key': 'FileSize', 'Value': str(metadata.get('ContentLength', 0))}
            ]
        }

        # サイズ制限（タグの値は256文字まで）
        for tag in tagging['TagSet']:
            if len(tag['Value']) > 256:
                tag['Value'] = tag['Value'][:253] + '...'

        # オブジェクトをコピー
        copy_source = {'Bucket': bucket_name, 'Key': object_key}
        s3_client.copy_object(
            CopySource=copy_source,
            Bucket=quarantine_bucket,
            Key=quarantine_key,
            TaggingDirective='REPLACE',
            Tagging=f"OriginalBucket={bucket_name}&OriginalKey={object_key}&QuarantineDate={timestamp}"
        )

        logger.info(f"オブジェクトを隔離しました: {bucket_name}/{object_key} -> {quarantine_bucket}/{quarantine_key}")

        # 元のオブジェクトを削除（設定による）
        if DELETE_ORIGINAL:
            s3_client.delete_object(Bucket=bucket_name, Key=object_key)
            logger.info(f"元のオブジェクトを削除しました: {bucket_name}/{object_key}")
        else:
            logger.info(f"元のオブジェクトは保持されます: {bucket_name}/{object_key}")

        return True, quarantine_key

    except ClientError as e:
        logger.error(f"隔離処理エラー: {e}")
        return False, None

def lambda_handler(event, context):
    """メインハンドラー"""
    try:
        # イベント情報をログ出力
        logger.info(f"受信イベント: {json.dumps(event)}")

        # EventBridgeイベントの解析
        if 'detail' not in event:
            logger.error("イベントに'detail'フィールドがありません")
            return {
                'statusCode': 400,
                'body': json.dumps('Invalid event format')
            }

        detail = event['detail']

        # バケット名とオブジェクトキーの取得
        bucket_name = detail.get('bucket', {}).get('name')
        object_key = detail.get('object', {}).get('key')

        if not bucket_name or not object_key:
            logger.error(f"必須パラメータが不足: bucket={bucket_name}, key={object_key}")
            return {
                'statusCode': 400,
                'body': json.dumps('Missing required parameters')
            }

        # EventBridgeから隔離バケットを取得（入力トランスフォーマーで追加される）
        quarantine_bucket = detail.get('quarantine_bucket', '')

        if not quarantine_bucket:
            # フォールバック：命名規則で自動決定
            quarantine_bucket = f"{bucket_name}-quarantine"
            logger.warning(f"隔離バケットが指定されていないため、命名規則を使用: {quarantine_bucket}")

        logger.info(f"処理対象: {bucket_name}/{object_key} -> 隔離先: {quarantine_bucket}")

        # タグの確認
        try:
            response = s3_client.get_object_tagging(
                Bucket=bucket_name,
                Key=object_key
            )
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                logger.warning(f"オブジェクトが見つかりません: {bucket_name}/{object_key}")
                return {
                    'statusCode': 404,
                    'body': json.dumps('Object not found')
                }
            raise

        tags = {tag['Key']: tag['Value'] for tag in response.get('TagSet', [])}
        logger.info(f"オブジェクトタグ: {tags}")

        # マルウェアが検出されたか確認
        if tags.get(TAG_KEY) == TAG_VALUE_THREAT:
            logger.warning(f"マルウェアを検出: {bucket_name}/{object_key}")

            # メタデータ取得
            metadata = get_object_metadata(bucket_name, object_key)

            # 隔離処理
            success, quarantine_key = quarantine_object(bucket_name, object_key, metadata, quarantine_bucket)

            if success:
                # 通知の送信
                notification_message = f"""
マルウェアが検出され、隔離されました。

検出情報:
- 元のバケット: {bucket_name}
- 元のオブジェクト: {object_key}
- ファイルサイズ: {metadata.get('ContentLength', 0)} bytes
- 検出時刻: {datetime.now(timezone.utc).isoformat()}
- 隔離先: {quarantine_bucket}/{quarantine_key}
- 元ファイル削除: {'はい' if DELETE_ORIGINAL else 'いいえ'}

詳細はCloudWatchログをご確認ください。
                """

                send_notification(
                    notification_message,
                    f"[GuardDuty] マルウェア検出通知: {object_key}"
                )

                return {
                    'statusCode': 200,
                    'body': json.dumps({
                        'message': 'Malware quarantined successfully',
                        'bucket': bucket_name,
                        'key': object_key,
                        'quarantine_location': f"{quarantine_bucket}/{quarantine_key}"
                    })
                }
            else:
                # 隔離失敗時の通知
                error_message = f"""
マルウェアが検出されましたが、隔離に失敗しました。

検出情報:
- バケット: {bucket_name}
- オブジェクト: {object_key}
- 検出時刻: {datetime.now(timezone.utc).isoformat()}

CloudWatchログで詳細を確認し、手動で対応してください。
                """

                send_notification(
                    error_message,
                    f"[GuardDuty] マルウェア隔離失敗: {object_key}"
                )

                return {
                    'statusCode': 500,
                    'body': json.dumps('Quarantine failed')
                }
        else:
            logger.info(f"マルウェアは検出されませんでした: {bucket_name}/{object_key}")
            return {
                'statusCode': 200,
                'body': json.dumps('No malware detected')
            }

    except Exception as e:
        logger.error(f"予期しないエラー: {str(e)}", exc_info=True)

        # エラー通知
        if sns_client and SNS_TOPIC_ARN:
            try:
                error_notification = f"""
Lambda関数で予期しないエラーが発生しました。

エラー内容: {str(e)}
関数名: {context.function_name}
リクエストID: {context.aws_request_id}

CloudWatchログで詳細を確認してください。
                """
                send_notification(
                    error_notification,
                    "[GuardDuty] Lambda処理エラー"
                )
            except:
                pass

        return {
            'statusCode': 500,
            'body': json.dumps(f'Internal error: {str(e)}')
        }