server/app.py

#!/usr/bin/env python3

from fastapi import FastAPI, Query, Request
from fastapi.responses import JSONResponse
import uvicorn
import json

app = FastAPI(title="AWS Security Auditor OpenEnv", version="1.0.0")

TASKS = {
    "easy_security_group": {
        "description": (
            "You are auditing an EC2 Security Group attached to a public-facing web server. "
            "Identify ALL security misconfigurations in the inbound and outbound rules."
        ),
        "config": json.dumps({
            "SecurityGroup": {
                "GroupId": "sg-0abc123def456789",
                "GroupName": "web-server-sg",
                "VpcId": "vpc-0123456789abcdef0",
                "IpPermissions": [
                    {"IpProtocol": "tcp", "FromPort": 22, "ToPort": 22, "IpRanges": [{"CidrIp": "0.0.0.0/0", "Description": ""}]},
                    {"IpProtocol": "tcp", "FromPort": 3389, "ToPort": 3389, "IpRanges": [{"CidrIp": "0.0.0.0/0", "Description": ""}]},
                    {"IpProtocol": "tcp", "FromPort": 80, "ToPort": 80, "IpRanges": [{"CidrIp": "0.0.0.0/0"}]},
                    {"IpProtocol": "tcp", "FromPort": 443, "ToPort": 443, "IpRanges": [{"CidrIp": "0.0.0.0/0"}]},
                    {"IpProtocol": "-1", "FromPort": -1, "ToPort": -1, "IpRanges": [{"CidrIp": "0.0.0.0/0"}]},
                ],
                "IpPermissionsEgress": [
                    {"IpProtocol": "-1", "FromPort": -1, "ToPort": -1, "IpRanges": [{"CidrIp": "0.0.0.0/0"}]}
                ],
                "Tags": []
            }
        }, indent=2),
        "expected_findings": [
            ["port 22", "ssh", "22/tcp", "ssh access", "ssh open"],
            ["port 3389", "rdp", "3389", "remote desktop"],
            ["all traffic", "all protocol", "protocol -1", "all ports", "unrestricted inbound", "all inbound"],
            ["unrestricted egress", "unrestricted outbound", "outbound unrestricted", "all outbound", "egress allows all", "0.0.0.0/0 egress", "outbound allow all"],
            ["no tag", "missing tag", "untagged", "empty tag", "tags are empty", "no tags", "tag not configured"],
        ]
    },

    "medium_s3_policy": {
        "description": (
            "Audit this S3 bucket configuration including its bucket policy, ACL settings, "
            "and security controls. Find all misconfigurations."
        ),
        "config": json.dumps({
            "BucketName": "prod-data-backup-20240315",
            "BucketACL": "public-read",
            "PublicAccessBlockConfiguration": {
                "BlockPublicAcls": False,
                "IgnorePublicAcls": False,
                "BlockPublicPolicy": False,
                "RestrictPublicBuckets": False
            },
            "VersioningConfiguration": {"Status": "Suspended"},
            "ServerSideEncryptionConfiguration": {"Rules": []},
            "LoggingConfiguration": {},
            "BucketPolicy": {
                "Version": "2012-10-17",
                "Statement": [{
                    "Sid": "PublicRead",
                    "Effect": "Allow",
                    "Principal": "*",
                    "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"],
                    "Resource": "arn:aws:s3:::prod-data-backup-20240315/*"
                }]
            },
            "MFADelete": "Disabled",
            "LifecycleConfiguration": {}
        }, indent=2),
        "expected_findings": [
            ["public-read", "public acl", "acl is public", "bucket acl", "publicly readable"],
            ["blockpublicacls", "block public acl", "public access block", "public access is not blocked", "block public access"],
            ["encryption", "sse", "server-side encryption", "no encryption", "not encrypted", "encryption not enabled"],
            ["versioning suspended", "versioning is disabled", "no versioning", "versioning not enabled", "versioning is not"],
            ["access logging", "logging is not", "no logging", "server access log", "logging not enabled", "logging disabled"],
            ["principal", "s3:putobject", "s3:deleteobject", "allows public write", "unauthenticated", "public write access"],
            ["mfa delete", "mfa-delete", "mfa is disabled", "no mfa", "mfa not enabled"],
        ]
    },

    "hard_iam_vpc": {
        "description": (
            "Audit this combined IAM policy and VPC configuration. "
            "Identify all privilege escalation risks, IAM misconfigurations, and network security issues."
        ),
        "config": json.dumps({
            "IAMPolicy": {
                "PolicyName": "DevOpsFullAccess",
                "PolicyDocument": {
                    "Version": "2012-10-17",
                    "Statement": [
                        {"Effect": "Allow", "Action": "*", "Resource": "*", "Condition": {}},
                        {"Effect": "Allow", "Action": ["iam:PassRole", "iam:CreateAccessKey", "iam:AttachUserPolicy"], "Resource": "*"}
                    ]
                }
            },
            "VPC": {
                "VpcId": "vpc-0deadbeef12345678",
                "CidrBlock": "10.0.0.0/16",
                "FlowLogs": [],
                "DefaultSecurityGroup": {
                    "GroupId": "sg-default",
                    "IpPermissions": [{"IpProtocol": "-1", "IpRanges": [{"CidrIp": "0.0.0.0/0"}]}],
                    "IpPermissionsEgress": [{"IpProtocol": "-1", "IpRanges": [{"CidrIp": "0.0.0.0/0"}]}]
                },
                "Subnets": [
                    {"SubnetId": "subnet-public-1", "MapPublicIpOnLaunch": True, "CidrBlock": "10.0.1.0/24"}
                ]
            },
            "CloudTrail": {"Enabled": False}
        }, indent=2),
        "expected_findings": [
            ["action *", "wildcard", "all actions", "full admin", "administrator", "action: \"*\"", "allows all actions"],
            ["resource *", "all resources", "resource: \"*\"", "any resource"],
            ["passrole", "iam:passrole", "privilege escalation", "pass role"],
            ["createaccesskey", "iam:createaccesskey", "create access key", "access key creation"],
            ["no mfa", "mfa not enforced", "missing mfa", "mfa condition", "without mfa", "mfa is not required"],
            ["flow log", "vpc flow", "no flow log", "flow logs not", "flow logs are not", "flow logs disabled"],
            ["default security group", "default sg", "default security group allows"],
            ["public ip", "auto-assign", "map public", "automatically assign", "public subnet", "assigns public"],
            ["cloudtrail", "no cloudtrail", "cloudtrail is not", "cloudtrail disabled", "cloudtrail not enabled"],
        ]
    },

    "medium_lambda_iam": {
        "description": (
            "Audit this AWS Lambda function configuration and its attached IAM execution role. "
            "Identify all security misconfigurations including over-privileged roles and secrets management issues."
        ),
        "config": json.dumps({
            "LambdaFunction": {
                "FunctionName": "process-payment-handler",
                "Runtime": "python3.9",
                "Timeout": 900,
                "MemorySize": 128,
                "VpcConfig": {},
                "Environment": {
                    "Variables": {
                        "DB_PASSWORD": "SuperSecret123!",
                        "API_KEY": "sk-live-abcdef1234567890",
                        "STRIPE_SECRET": "sk_live_abcdef",
                        "ENV": "production"
                    }
                },
                "DeadLetterConfig": {},
                "ReservedConcurrentExecutions": None,
                "TracingConfig": {"Mode": "PassThrough"},
                "ResourcePolicy": None
            },
            "ExecutionRole": {
                "RoleName": "lambda-execution-role",
                "AttachedPolicies": [
                    {"PolicyArn": "arn:aws:iam::aws:policy/AdministratorAccess"}
                ],
                "InlinePolicies": []
            }
        }, indent=2),
        "expected_findings": [
            ["administratoraccess", "administrator access", "admin policy", "overprivileged", "full admin"],
            ["db_password", "plaintext", "hardcoded", "plaintext secret", "environment variable secret", "secrets manager", "ssm parameter"],
            ["api_key", "api key", "plaintext api", "hardcoded api"],
            ["stripe_secret", "stripe", "payment secret", "hardcoded payment"],
            ["no vpc", "vpc config empty", "not in vpc", "lambda vpc", "no network isolation"],
            ["dead letter", "dlq", "no dlq", "dead letter queue"],
            ["timeout 900", "max timeout", "timeout too high", "900 second"],
            ["tracing", "x-ray", "passthrough", "no tracing", "tracing disabled"],
        ]
    },

    "hard_rds_cloudtrail": {
        "description": (
            "Audit this RDS database instance and CloudTrail configuration. "
            "Identify all data security, compliance, and audit logging gaps."
        ),
        "config": json.dumps({
            "RDSInstance": {
                "DBInstanceIdentifier": "prod-mysql-primary",
                "Engine": "mysql",
                "EngineVersion": "5.7.38",
                "PubliclyAccessible": True,
                "StorageEncrypted": False,
                "BackupRetentionPeriod": 0,
                "MultiAZ": False,
                "DeletionProtection": False,
                "AutoMinorVersionUpgrade": False,
                "IAMDatabaseAuthenticationEnabled": False,
                "EnabledCloudwatchLogsExports": [],
                "VpcSecurityGroups": [{"VpcSecurityGroupId": "sg-rds-open", "Status": "active"}],
                "StorageType": "gp2",
                "EnhancedMonitoringEnabled": False
            },
            "CloudTrail": {
                "TrailARN": "arn:aws:cloudtrail:us-east-1:123456789:trail/main-trail",
                "IsMultiRegionTrail": False,
                "IncludeGlobalServiceEvents": False,
                "LogFileValidationEnabled": False,
                "S3BucketName": "cloudtrail-logs-public",
                "S3BucketACL": "public-read-write",
                "CloudWatchLogsLogGroupArn": None,
                "KMSKeyId": None,
                "IsOrganizationTrail": False
            }
        }, indent=2),
        # ── ONLY THIS BLOCK CHANGED ──────────────────────────────────────
        "expected_findings": [
            ["publicly accessible", "rds public", "public database", "public exposure",
             "publicly exposed", "database is public", "rds instance is publicly"],
            ["storage encrypted", "not encrypted", "encryption disabled", "storageencrypted false",
             "no storage encryption", "unencrypted storage", "encryption not enabled"],
            ["backup", "backupretentionperiod 0", "no backup", "retention period 0",
             "backup retention", "automated backup", "retention is 0"],
            ["multiaz", "multi-az", "no multi-az", "single az", "no high availability",
             "single availability zone", "not multi-az", "multi az not"],
            ["deletion protection", "deletionprotection false", "no deletion protection",
             "deletion protection disabled", "deletion protection not enabled"],
            ["auto minor version", "autominorversionupgrade", "version upgrade", "patch",
             "automatic minor version", "minor version upgrade disabled", "auto upgrade disabled"],
            ["iam database auth", "iam authentication", "password authentication only",
             "iam db auth", "database authentication", "no iam auth", "iam database authentication"],
            ["cloudwatch log", "no log export", "logs not enabled", "audit log",
             "log exports", "cloudwatch logs export", "no cloudwatch log", "logs are not exported"],
            ["multi-region trail", "ismultiregiontrail false", "single region", "not multi-region",
             "single-region", "not configured for multiple regions", "trail not multi", "multi region trail"],
            ["log file validation", "logfilevalidation", "validation disabled",
             "log validation", "file validation", "integrity validation", "log file integrity",
             "log file validation is not enabled", "validation not enabled"],
            ["s3 bucket public", "cloudtrail bucket public", "public-read-write", "bucket acl",
             "s3 bucket acl", "public bucket", "publicly accessible bucket", "bucket is public"],
            ["cloudwatch logs", "cloudwatchlogsloggrouparn", "no cloudwatch integration",
             "no cloudwatch logs group", "cloudwatch log group", "not integrated with cloudwatch",
             "cloudwatch logs log group", "cloudwatch integration not"],
            ["kms", "no kms", "kms encryption", "trail encryption",
             "not encrypted with kms", "kms key not", "encryption key", "kms key is not"],
        ]
        # ── END CHANGE ───────────────────────────────────────────────────
    }
}

state = {
    "task": None,
    "step": 0,
    "last_reward": 0.0,
    "last_feedback": None,
    "done": False,
}

def score_response(task_name, findings, severity, recommendations, config_patch):
    if task_name not in TASKS:
        return 0.0

    expected = TASKS[task_name]["expected_findings"]
    agent_text = " ".join(str(x).lower() for x in (findings + recommendations))

    matched = 0
    for keyword_group in expected:
        if any(kw.lower() in agent_text for kw in keyword_group):
            matched += 1

    raw_score = matched / len(expected)

    high_count = sum(1 for s in severity if str(s).upper() == "HIGH")
    if high_count >= 2 and raw_score > 0.5:
        raw_score = min(1.0, raw_score + 0.05)

    if isinstance(config_patch, dict) and len(config_patch) > 0:
        raw_score = min(1.0, raw_score + 0.03)

    return round(raw_score, 3)

def build_feedback(task_name, findings, recommendations):
    expected = TASKS[task_name]["expected_findings"]
    agent_text = " ".join(str(x).lower() for x in (findings + recommendations))
    missed = []

    for keyword_group in expected:
        if not any(kw.lower() in agent_text for kw in keyword_group):
            missed.append(keyword_group[0])

    if not missed:
        return "All key findings identified. Excellent audit."

    return (
        "Missed findings related to: "
        + ", ".join(missed[:4])
        + ". Re-examine the full configuration carefully and use more explicit AWS wording."
    )

@app.get("/health")
def health():
    return {
        "status": "ok",
        "environment": "aws-security-auditor",
        "version": "1.0.0",
    }

@app.post("/reset")
async def reset(task: str = Query(...)):
    if task not in TASKS:
        return JSONResponse(
            status_code=400,
            content={"error": f"Unknown task: {task}. Valid tasks: {list(TASKS.keys())}"},
        )

    state["task"] = task
    state["step"] = 0
    state["last_reward"] = 0.0
    state["last_feedback"] = None
    state["done"] = False

    return {
        "observation": {
            "task_description": TASKS[task]["description"],
            "config": TASKS[task]["config"],
            "feedback": None,
            "last_reward": 0.0,
            "step": 0,
        }
    }

@app.post("/step")
async def step(request: Request):
    body = await request.json()

    findings = body.get("findings", [])
    severity = body.get("severity", [])
    recommendations = body.get("recommendations", [])
    config_patch = body.get("config_patch", {})

    task_name = state.get("task")
    if not task_name:
        return JSONResponse(status_code=400, content={"error": "Call /reset first."})

    state["step"] += 1
    reward = score_response(task_name, findings, severity, recommendations, config_patch)
    feedback = build_feedback(task_name, findings, recommendations)

    state["last_reward"] = reward
    state["last_feedback"] = feedback
    done = reward >= 0.9 or state["step"] >= 5
    state["done"] = done

    return {
        "reward": reward,
        "done": done,
        "observation": {
            "task_description": TASKS[task_name]["description"],
            "config": TASKS[task_name]["config"],
            "feedback": feedback,
            "last_reward": reward,
            "step": state["step"],
        },
    }

@app.get("/state")
def get_state():
    return state

@app.get("/schema")
def schema():
    return {
        "reset": {
            "input": {"task": "string"},
            "output": {"observation": "object"},
        },
        "step": {
            "input": {
                "findings": "list[str]",
                "severity": "list[str]",
                "recommendations": "list[str]",
                "config_patch": "object",
            },
            "output": {
                "reward": "float",
                "done": "bool",
                "observation": "object",
            },
        },
    }

@app.get("/metadata")
def metadata():
    return {
        "name": "aws-security-auditor",
        "version": "1.0.0",
        "tasks": list(TASKS.keys()),
        "profile": "openenv-http/1.x",
        "success_threshold": 0.55,
        "max_steps": 5,
    }

@app.get("/mcp")
def mcp():
    return {
        "protocol": "openenv-http/1.x",
        "endpoints": ["/reset", "/step", "/state", "/health", "/schema", "/metadata"],
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)