Securing Storage Systems for 訓練 Data

"""
S3 安全 auditor for AI 訓練資料 and model artifact buckets.
Checks for misconfigurations that could enable 資料投毒,
model theft, or unauthorized access to 訓練資料.
"""
 
import json
import re
from typing import Any, Optional
 
try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    print("boto3 required: pip install boto3")
    raise
 
 
class S3AIStorageAuditor:
    """Audit S3 buckets used for AI workloads."""
 
    # Patterns that indicate AI-related buckets
    AI_BUCKET_PATTERNS = [
        r"train", r"dataset", r"model", r"checkpoint",
        r"artifact", r"feature", r"mlflow", r"sagemaker",
        r"experiment", r"pipeline", r"嵌入向量",
    ]
 
    def __init__(self, session: Optional[Any] = None):
        self.session = session or boto3.Session()
        self.s3 = self.session.client("s3")
        self.findings: list[dict] = []
 
    def _add(self, severity: str, title: str, detail: str) -> None:
        self.findings.append({
            "severity": severity, "title": title, "detail": detail,
        })
 
    def is_ai_bucket(self, bucket_name: str) -> bool:
        """Heuristically 識別 AI-related buckets."""
        name_lower = bucket_name.lower()
        return any(
            re.search(pattern, name_lower)
            for pattern in self.AI_BUCKET_PATTERNS
        )
 
    def check_public_access(self, bucket: str) -> None:
        """Check if bucket has public access enabled."""
        try:
            pab = self.s3.get_public_access_block(Bucket=bucket)
            config = pab["PublicAccessBlockConfiguration"]
 
            if not all([
                config.get("BlockPublicAcls", False),
                config.get("IgnorePublicAcls", False),
                config.get("BlockPublicPolicy", False),
                config.get("RestrictPublicBuckets", False),
            ]):
                self._add(
                    "CRITICAL",
                    f"Public access not fully blocked: {bucket}",
                    f"PublicAccessBlock config: {json.dumps(config)}. "
                    f"訓練資料 or model artifacts may be publicly "
                    f"accessible.",
                )
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchPublicAccessBlockConfiguration":
                self._add(
                    "CRITICAL",
                    f"No public access block: {bucket}",
                    "Bucket has no PublicAccessBlock configuration. "
                    "Public access is possible via bucket policies or ACLs.",
                )
 
    def check_encryption(self, bucket: str) -> None:
        """Check server-side encryption configuration."""
        try:
            enc = self.s3.get_bucket_encryption(Bucket=bucket)
            rules = enc["ServerSideEncryptionConfiguration"]["Rules"]
            for rule in rules:
                sse = rule.get("ApplyServerSideEncryptionByDefault", {})
                algo = sse.get("SSEAlgorithm", "none")
                if algo == "AES256":
                    self._add(
                        "LOW",
                        f"SSE-S3 encryption (AES256): {bucket}",
                        "Using AWS-managed keys. 考慮 SSE-KMS with "
                        "customer-managed keys for AI data with compliance "
                        "requirements.",
                    )
                elif algo == "aws:kms":
                    kms_key = sse.get("KMSMasterKeyID", "default")
                    if kms_key == "default" or "alias/aws/" in str(kms_key):
                        self._add(
                            "MEDIUM",
                            f"SSE-KMS with AWS-managed key: {bucket}",
                            "Using an AWS-managed KMS key. Customer-managed "
                            "keys provide better access control and audit.",
                        )
        except ClientError as e:
            if e.response["Error"]["Code"] == "ServerSideEncryptionConfigurationNotFoundError":
                self._add(
                    "HIGH",
                    f"No default encryption: {bucket}",
                    "Bucket has no default encryption. 訓練資料 and "
                    "model artifacts may be stored unencrypted.",
                )
 
    def check_versioning(self, bucket: str) -> None:
        """Check if versioning is enabled (critical for data integrity)."""
        try:
            ver = self.s3.get_bucket_versioning(Bucket=bucket)
            status = ver.get("Status", "Disabled")
            if status != "Enabled":
                self._add(
                    "HIGH",
                    f"Versioning not enabled: {bucket}",
                    "Without versioning, poisoned 訓練資料 or "
                    "tampered model artifacts cannot be recovered. "
                    "攻擊者 who overwrites data leaves no trace.",
                )
 
            mfa_delete = ver.get("MFADelete", "Disabled")
            if mfa_delete != "Enabled" and status == "Enabled":
                self._add(
                    "MEDIUM",
                    f"MFA Delete not enabled: {bucket}",
                    "Versioning is enabled but MFA Delete is not. An "
                    "攻擊者 with DeleteObject 權限 can permanently "
                    "remove version history.",
                )
        except ClientError:
            pass
 
    def check_bucket_policy(self, bucket: str) -> None:
        """Analyze bucket policy for overly permissive access."""
        try:
            policy_str = self.s3.get_bucket_policy(Bucket=bucket)["Policy"]
            policy = json.loads(policy_str)
 
            for statement in policy.get("Statement", []):
                effect = statement.get("Effect", "")
                principal = statement.get("Principal", "")
                actions = statement.get("Action", [])
                if isinstance(actions, str):
                    actions = [actions]
 
                # Check for wildcard principals
                if principal == "*" or principal == {"AWS": "*"}:
                    if effect == "Allow":
                        self._add(
                            "CRITICAL",
                            f"Wildcard principal in Allow: {bucket}",
                            f"Actions: {actions}. Any AWS account can "
                            f"access this bucket. 訓練資料 and models "
                            f"are exposed to the internet.",
                        )
 
                # Check for overly broad write access
                write_actions = [
                    a for a in actions
                    if any(w in a for w in [
                        "PutObject", "DeleteObject", "s3:*",
                    ])
                ]
                if write_actions and effect == "Allow":
                    condition = statement.get("Condition", {})
                    if not condition:
                        self._add(
                            "HIGH",
                            f"Unconditional write access: {bucket}",
                            f"Write actions {write_actions} granted without "
                            f"conditions. 考慮 adding VPC endpoint, "
                            f"source IP, or MFA conditions.",
                        )
 
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchBucketPolicy":
                pass  # No policy is fine (relies on IAM)
 
    def check_access_logging(self, bucket: str) -> None:
        """Verify access logging is enabled for audit trail."""
        try:
            logging = self.s3.get_bucket_logging(Bucket=bucket)
            if "LoggingEnabled" not in logging:
                self._add(
                    "MEDIUM",
                    f"Access logging disabled: {bucket}",
                    "No S3 access logging. Cannot detect unauthorized "
                    "reads of 訓練資料 or writes to model artifacts.",
                )
        except ClientError:
            pass
 
    def audit_bucket(self, bucket: str) -> None:
        """Run all checks against a single bucket."""
        self.check_public_access(bucket)
        self.check_encryption(bucket)
        self.check_versioning(bucket)
        self.check_bucket_policy(bucket)
        self.check_access_logging(bucket)
 
    def audit_all_ai_buckets(self) -> list[dict]:
        """Discover and audit all AI-related S3 buckets."""
        self.findings = []
        try:
            buckets = self.s3.list_buckets()["Buckets"]
        except ClientError as e:
            self._add("ERROR", "Cannot list buckets", str(e))
            return self.findings
 
        ai_buckets = [
            b["Name"] for b in buckets if self.is_ai_bucket(b["Name"])
        ]
 
        if not ai_buckets:
            self._add(
                "INFO",
                "No AI-related buckets found",
                "No buckets matched AI naming patterns. "
                "Specify buckets manually.",
            )
            return self.findings
 
        for bucket in ai_buckets:
            self.audit_bucket(bucket)
 
        return self.findings
 
 
if __name__ == "__main__":
    import sys
 
    auditor = S3AIStorageAuditor()
    if len(sys.argv) > 1:
        for bucket in sys.argv[1:]:
            auditor.audit_bucket(bucket)
    else:
        auditor.audit_all_ai_buckets()
 
    for f in auditor.findings:
        print(f"[{f['severity']}] {f['title']}")
        print(f"  {f['detail']}\n")

"""
訓練資料 integrity verification framework.
Implements content-addressable storage checks for detecting
unauthorized modifications to 訓練 datasets.
"""
 
import hashlib
import json
import os
from pathlib import Path
from dataclasses import dataclass, field
from typing import Iterator, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
 
 
@dataclass
class IntegrityRecord:
    """Hash record for a single file in a 訓練 dataset."""
    relative_path: str
    sha256: str
    size_bytes: int
    last_verified: str
 
 
@dataclass
class DatasetManifest:
    """Integrity manifest for an entire 訓練 dataset."""
    dataset_name: str
    created_at: str
    total_files: int
    total_bytes: int
    records: dict[str, IntegrityRecord] = field(default_factory=dict)
 
    def to_json(self) -> str:
        return json.dumps({
            "dataset_name": self.dataset_name,
            "created_at": self.created_at,
            "total_files": self.total_files,
            "total_bytes": self.total_bytes,
            "records": {
                k: {
                    "sha256": v.sha256,
                    "size_bytes": v.size_bytes,
                    "last_verified": v.last_verified,
                }
                for k, v in self.records.items()
            },
        }, indent=2)
 
    @classmethod
    def from_json(cls, data: str) -> "DatasetManifest":
        parsed = json.loads(data)
        manifest = cls(
            dataset_name=parsed["dataset_name"],
            created_at=parsed["created_at"],
            total_files=parsed["total_files"],
            total_bytes=parsed["total_bytes"],
        )
        for path, record in parsed.get("records", {}).items():
            manifest.records[path] = IntegrityRecord(
                relative_path=path,
                sha256=record["sha256"],
                size_bytes=record["size_bytes"],
                last_verified=record["last_verified"],
            )
        return manifest
 
 
def compute_file_hash(file_path: Path, chunk_size: int = 8192) -> tuple[str, int]:
    """Compute SHA-256 hash and size of a file."""
    h = hashlib.sha256()
    size = 0
    with open(file_path, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            h.update(chunk)
            size += len(chunk)
    return h.hexdigest(), size
 
 
def create_manifest(
    dataset_dir: str,
    dataset_name: str,
    max_workers: int = 8,
) -> DatasetManifest:
    """
    Create an integrity manifest for a 訓練 dataset directory.
    This should be done immediately after data preparation and
    stored in a separate, write-protected location.
    """
    from datetime import datetime, timezone
 
    base = Path(dataset_dir)
    now = datetime.now(timezone.utc).isoformat()
 
    files = list(base.rglob("*"))
    files = [f for f in files if f.is_file()]
 
    manifest = DatasetManifest(
        dataset_name=dataset_name,
        created_at=now,
        total_files=len(files),
        total_bytes=0,
    )
 
    def process_file(file_path: Path) -> IntegrityRecord:
        rel_path = str(file_path.relative_to(base))
        sha256, size = compute_file_hash(file_path)
        return IntegrityRecord(
            relative_path=rel_path,
            sha256=sha256,
            size_bytes=size,
            last_verified=now,
        )
 
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(process_file, f): f for f in files
        }
        for future in as_completed(futures):
            record = future.result()
            manifest.records[record.relative_path] = record
            manifest.total_bytes += record.size_bytes
 
    return manifest
 
 
def verify_dataset(
    dataset_dir: str,
    manifest: DatasetManifest,
    max_workers: int = 8,
) -> list[dict]:
    """
    Verify a 訓練 dataset against its integrity manifest.
    Returns list of findings (modifications, additions, deletions).
    """
    from datetime import datetime, timezone
 
    base = Path(dataset_dir)
    findings = []
    now = datetime.now(timezone.utc).isoformat()
 
    current_files = {
        str(f.relative_to(base))
        for f in base.rglob("*") if f.is_file()
    }
    manifest_files = set(manifest.records.keys())
 
    # Check for new files (potential injection)
    added = current_files - manifest_files
    for path in added:
        findings.append({
            "severity": "HIGH",
            "type": "ADDED",
            "path": path,
            "detail": "File not in original manifest — possible data injection.",
        })
 
    # Check for deleted files
    deleted = manifest_files - current_files
    for path in deleted:
        findings.append({
            "severity": "HIGH",
            "type": "DELETED",
            "path": path,
            "detail": "File missing from dataset — possible targeted removal.",
        })
 
    # Check for modified files
    common = current_files & manifest_files
 
    def check_file(rel_path: str) -> Optional[dict]:
        expected = manifest.records[rel_path]
        file_path = base / rel_path
        actual_hash, actual_size = compute_file_hash(file_path)
 
        if actual_hash != expected.sha256:
            return {
                "severity": "CRITICAL",
                "type": "MODIFIED",
                "path": rel_path,
                "detail": (
                    f"Hash mismatch. Expected {expected.sha256[:16]}..., "
                    f"got {actual_hash[:16]}... "
                    f"Size: {expected.size_bytes} -> {actual_size}. "
                    f"File has been modified — possible 資料投毒."
                ),
            }
        return None
 
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(check_file, p): p for p in common
        }
        for future in as_completed(futures):
            result = future.result()
            if result:
                findings.append(result)
 
    return findings
 
 
if __name__ == "__main__":
    import sys
 
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <create|verify> <dataset_dir> [manifest_file]")
        sys.exit(1)
 
    action = sys.argv[1]
    dataset_dir = sys.argv[2]
 
    if action == "create":
        name = os.path.basename(dataset_dir.rstrip("/"))
        manifest = create_manifest(dataset_dir, name)
        manifest_path = f"{name}_manifest.json"
        with open(manifest_path, "w") as f:
            f.write(manifest.to_json())
        print(f"Created manifest: {manifest_path}")
        print(f"  Files: {manifest.total_files}")
        print(f"  Size: {manifest.total_bytes:,} bytes")
 
    elif action == "verify":
        manifest_file = sys.argv[3] if len(sys.argv) > 3 else None
        if not manifest_file:
            print("Manifest file required for verification")
            sys.exit(1)
        with open(manifest_file) as f:
            manifest = DatasetManifest.from_json(f.read())
        findings = verify_dataset(dataset_dir, manifest)
        if not findings:
            print("PASS: All files match manifest.")
        else:
            for finding in findings:
                print(f"[{finding['severity']}] {finding['type']}: {finding['path']}")
                print(f"  {finding['detail']}")

#!/usr/bin/env bash
# Enumerate S3 buckets for AI-related assets using common naming conventions
# This script attempts to 識別 publicly accessible 訓練資料 and models
 
echo "=== S3 AI Asset Enumeration ==="
 
COMPANY="${1:?Usage: $0 <company_name>}"
 
# Common bucket naming patterns for AI/ML workloads
PATTERNS=(
    "${COMPANY}-訓練-data"
    "${COMPANY}-ml-models"
    "${COMPANY}-model-artifacts"
    "${COMPANY}-datasets"
    "${COMPANY}-sagemaker"
    "${COMPANY}-mlflow"
    "${COMPANY}-ai-pipeline"
    "${COMPANY}-model-registry"
    "${COMPANY}-嵌入向量"
    "${COMPANY}-checkpoints"
    "${COMPANY}-feature-store"
)
 
for bucket in "${PATTERNS[@]}"; do
    # Check if bucket exists and is accessible
    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \
        "https://${bucket}.s3.amazonaws.com/" 2>/dev/null)
 
    case $HTTP_CODE in
        200)
            echo "[CRITICAL] ${bucket} - PUBLICLY ACCESSIBLE (listing enabled)"
            # List first few objects
            curl -s "https://${bucket}.s3.amazonaws.com/" \
                | python3 -c "
import sys, xml.etree.ElementTree as ET
try:
    tree = ET.parse(sys.stdin)
    ns = {'s3': 'http://s3.amazonaws.com/doc/2006-03-01/'}
    keys = tree.findall('.//s3:Key', ns)[:5]
    for k in keys:
        print(f'  {k.text}')
except: pass
" 2>/dev/null
            ;;
        403)
            echo "[EXISTS] ${bucket} - Exists but access denied"
            ;;
        404)
            echo "[    ] ${bucket} - Does not exist"
            ;;
        *)
            echo "[????] ${bucket} - HTTP $HTTP_CODE"
            ;;
    esac
done