eval_all.py

import json
import re
import torch
import requests
import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*max_new_tokens.*max_length.*")

from unsloth import FastLanguageModel

BASE_URL = "https://kkaustav-cloud-config-auditor.hf.space"

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="/content/aws-auditor-sft-v5",
    max_seq_length=4096,
    dtype=None,
    load_in_4bit=True,
)
FastLanguageModel.for_inference(model)

# ── DO NOT TOUCH — scores 0.99 ────────────────────────────────────────────────
SYSTEM_DEFAULT = (
    "You are an expert AWS cloud security auditor. "
    "Return your audit as valid JSON ONLY — no markdown, no code blocks, no extra text.\n"
    "IMPORTANT: Your ENTIRE response must be ONE complete, closed JSON object. "
    "Close every array ] and every brace } before you stop.\n"
    '{"findings":["finding 1"],"severity":["CRITICAL"],'
    '"recommendations":["fix 1"],"config_patch":{"key":"value"}}'
)

# ── DO NOT TOUCH — scores 0.99 ────────────────────────────────────────────────
SYSTEM_LAMBDA_IAM = (
    "You are an expert AWS Lambda and IAM security auditor for a payment processing service.\n"
    "Analyze the provided Lambda configuration and report EVERY distinct security misconfiguration.\n\n"
    "CHECK EACH OF THESE SEPARATELY — each must be its own finding:\n"
    "1. IAM role: Is AdministratorAccess or any wildcard policy attached? (CRITICAL)\n"
    "2. Env var DB_PASSWORD: Is a plaintext database password stored here? (CRITICAL)\n"
    "3. Env var API_SECRET_KEY: Is a plaintext API key stored here? (CRITICAL)\n"
    "4. Env var STRIPE_SECRET: Is a plaintext Stripe/payment secret stored here? (CRITICAL)\n"
    "5. FunctionUrl.AuthType: Is it set to NONE (unauthenticated public access)? (HIGH)\n"
    "6. CORS AllowOrigins: Is it set to ['*'] (wildcard)? (HIGH)\n"
    "7. VpcConfig: Is it null (Lambda not isolated in a VPC)? (HIGH)\n"
    "8. ReservedConcurrency: Is it null (no DDoS/cost protection)? (MEDIUM)\n"
    "9. Tracing.Mode: Is it PassThrough (no X-Ray active tracing)? (LOW)\n"
    "10. CodeSigningConfigArn: Is it null (no code integrity verification)? (MEDIUM)\n"
    "11. Runtime: Is it python3.9 (deprecated/end-of-life runtime)? (MEDIUM)\n\n"
    "Rules:\n"
    "- List EACH env var secret as a SEPARATE finding — do not group them.\n"
    "- List ALL findings even if they seem minor.\n"
    "- Return valid JSON ONLY. No markdown. No code blocks.\n"
    "- Your ENTIRE response must be ONE complete, closed JSON object.\n"
    "- severity array must match findings array length.\n"
    '{"findings":["AdministratorAccess IAM policy attached","DB_PASSWORD stored as plaintext"],'
    '"severity":["CRITICAL","CRITICAL"],'
    '"recommendations":["Replace AdministratorAccess with least-privilege policy","Move DB_PASSWORD to AWS Secrets Manager"],'
    '"config_patch":{'
    '"ExecutionRole.AttachedPolicies":["arn:aws:iam::123456789012:policy/LambdaPaymentMinimalPolicy"],'
    '"EnvironmentVariables.DB_PASSWORD":"{{resolve:secretsmanager:payment-db-secret}}",'
    '"FunctionUrl.AuthType":"AWS_IAM",'
    '"VpcConfig":{"SubnetIds":["subnet-xxx"],"SecurityGroupIds":["sg-xxx"]},'
    '"ReservedConcurrency":100,'
    '"Tracing":{"Mode":"Active"},'
    '"CodeSigningConfigArn":"arn:aws:lambda:ap-south-1:123456789012:code-signing-config:csc-xxx"'
    "}}"
)

# ── REVERTED — back to generic prompt that scored 0.92 (13 findings) ─────────
# The precision prompt caused the model to DROP findings 12,13 (MultiRegion,AWSConfig)
# Generic prompt + repair_json recovers truncated output and gets 13 findings naturally
SYSTEM_RDS_CLOUDTRAIL = SYSTEM_DEFAULT

# ── hard_iam_vpc — NEW: explicitly lists ALL 17 config fields in TWO sections ─
# Root cause of 0.50: model was generating IAM findings (10) but SKIPPING all VPC
# findings entirely. Fix: structure prompt with a clear VPC section header and
# explicitly name every VPC field from the actual config
SYSTEM_IAM_VPC = (
    "You are an expert AWS IAM and VPC security auditor.\n"
    "The config has TWO sections: IAMRole and VPC. You MUST audit BOTH sections.\n"
    "Report EVERY misconfiguration as a SEPARATE finding.\n\n"
    "=== SECTION 1: IAMRole findings — check each separately ===\n"
    "1. AssumeRolePolicyDocument has Principal.Service='*' — any service can assume this role (CRITICAL)\n"
    "2. InlinePolicies has Action='*' — wildcard action grants full AWS access (CRITICAL)\n"
    "3. InlinePolicies has Resource='*' — no resource-level restriction (CRITICAL)\n"
    "4. Inline policy 'AllAccess-temp' used — should be a least-privilege managed policy (MEDIUM)\n"
    "5. MFAEnabled=false — no MFA requirement on the role (HIGH)\n"
    "6. PasswordPolicy.MinimumPasswordLength=6 — minimum must be at least 14 (MEDIUM)\n"
    "7. PasswordPolicy.RequireUppercaseCharacters=false — uppercase not enforced (MEDIUM)\n"
    "8. PasswordPolicy.RequireNumbers=false — numbers not required (MEDIUM)\n"
    "9. PasswordPolicy.RequireSymbols=false — symbols not required (MEDIUM)\n"
    "10. PasswordPolicy.MaxPasswordAge=0 — passwords never expire (MEDIUM)\n"
    "11. PasswordPolicy.PasswordReusePrevention=0 — old passwords can be reused (MEDIUM)\n\n"
    "=== SECTION 2: VPC findings — check each separately ===\n"
    "12. VPC.FlowLogsEnabled=false — VPC Flow Logs are disabled, no traffic visibility (HIGH)\n"
    "13. VPC.CloudTrailEnabled=false — CloudTrail is disabled, no API audit logging (CRITICAL)\n"
    "14. VPC.GuardDutyEnabled=false — GuardDuty threat detection is not enabled (HIGH)\n"
    "15. NetworkACL inbound rule: Protocol='-1', CidrBlock='0.0.0.0/0', Action='allow' — all inbound traffic permitted (HIGH)\n"
    "16. NetworkACL outbound rule: Protocol='-1', CidrBlock='0.0.0.0/0', Action='allow' — all outbound traffic permitted (MEDIUM)\n"
    "17. Subnet MapPublicIpOnLaunch=true — instances get public IPs automatically on launch (HIGH)\n\n"
    "CRITICAL RULES:\n"
    "- You MUST include findings from BOTH the IAMRole section AND the VPC section.\n"
    "- Do NOT stop after IAM findings — continue to VPC findings 12-17.\n"
    "- Each numbered item = one separate finding string.\n"
    "- Do NOT group multiple items into one finding.\n"
    "- Respond with valid JSON ONLY. No markdown. No code blocks.\n"
    "- Your ENTIRE response must be ONE complete, closed JSON object.\n"
    "- severity array length must EXACTLY match findings array length.\n\n"
    "Example output structure (use ALL 17 findings):\n"
    '{"findings":['
    '"IAM role trust policy allows any AWS service to assume it — Principal.Service is wildcard (*)",'
    '"Inline policy AllAccess-temp grants Action=* — full administrative access",'
    '"Inline policy AllAccess-temp has Resource=* — no resource-level restriction",'
    '"Inline policy used instead of least-privilege managed policy (AllAccess-temp)",'
    '"MFAEnabled=false — IAM role has no MFA requirement",'
    '"PasswordPolicy MinimumPasswordLength=6 — below the 14-character minimum",'
    '"PasswordPolicy RequireUppercaseCharacters=false — uppercase not required",'
    '"PasswordPolicy RequireNumbers=false — numbers not required",'
    '"PasswordPolicy RequireSymbols=false — symbols not required",'
    '"PasswordPolicy MaxPasswordAge=0 — passwords never expire",'
    '"PasswordPolicy PasswordReusePrevention=0 — no password reuse prevention",'
    '"VPC FlowLogsEnabled=false — no network traffic logging",'
    '"VPC CloudTrailEnabled=false — no API audit logging",'
    '"VPC GuardDutyEnabled=false — no threat detection",'
    '"NetworkACL inbound allows all traffic (Protocol=-1, 0.0.0.0/0)",'
    '"NetworkACL outbound allows all traffic (Protocol=-1, 0.0.0.0/0)",'
    '"Subnet MapPublicIpOnLaunch=true — public IPs auto-assigned on launch"'
    "],"
    '"severity":["CRITICAL","CRITICAL","CRITICAL","MEDIUM","HIGH","MEDIUM","MEDIUM","MEDIUM","MEDIUM","MEDIUM","MEDIUM","HIGH","CRITICAL","HIGH","HIGH","MEDIUM","HIGH"],'
    '"recommendations":['
    '"Restrict Principal.Service to ecs-tasks.amazonaws.com",'
    '"Replace wildcard Action with specific required permissions",'
    '"Replace wildcard Resource with specific ARNs",'
    '"Replace inline policy with least-privilege managed policy",'
    '"Enable MFA for all IAM role assumptions",'
    '"Set MinimumPasswordLength to 14",'
    '"Enable RequireUppercaseCharacters",'
    '"Enable RequireNumbers",'
    '"Enable RequireSymbols",'
    '"Set MaxPasswordAge to 90 days",'
    '"Set PasswordReusePrevention to 24",'
    '"Enable VPC Flow Logs",'
    '"Enable CloudTrail in all regions",'
    '"Enable GuardDuty",'
    '"Restrict NACL inbound to known CIDR blocks",'
    '"Restrict NACL outbound to required destinations",'
    '"Disable MapPublicIpOnLaunch on all production subnets"'
    "],"
    '"config_patch":{'
    '"IAMRole.AssumeRolePolicyDocument.Statement[0].Principal":{"Service":"ecs-tasks.amazonaws.com"},'
    '"VPC.FlowLogsEnabled":true,'
    '"VPC.GuardDutyEnabled":true,'
    '"VPC.CloudTrailEnabled":true,'
    '"Subnets[0].MapPublicIpOnLaunch":false'
    "}}"
)

SYSTEM_MAP = {
    "medium_lambda_iam":   SYSTEM_LAMBDA_IAM,
    "hard_rds_cloudtrail": SYSTEM_RDS_CLOUDTRAIL,
    "hard_iam_vpc":        SYSTEM_IAM_VPC,
}

# ── Robust JSON repair (unchanged) ───────────────────────────────────────────
def repair_json(text):
    text = text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text)
    text = text.strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    start = text.find("{")
    if start == -1:
        return None
    text = text[start:]

    in_string = False
    escape_next = False
    stack = []

    for ch in text:
        if escape_next:
            escape_next = False
            continue
        if ch == "\\" and in_string:
            escape_next = True
            continue
        if ch == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch in "{[":
            stack.append(ch)
        elif ch == "]" and stack and stack[-1] == "[":
            stack.pop()
        elif ch == "}" and stack and stack[-1] == "{":
            stack.pop()

    repair = ""
    if in_string:
        repair += '"'
    for opener in reversed(stack):
        repair += "]" if opener == "[" else "}"

    try:
        return json.loads(text + repair)
    except json.JSONDecodeError:
        pass

    findings_match = re.findall(r'"findings"\s*:\s*\[([^\]]*)', text + repair, re.DOTALL)
    if findings_match:
        items = re.findall(r'"((?:[^"\\]|\\.)*)"', findings_match[0])
        if items:
            return {"findings": items, "severity": [], "recommendations": [], "config_patch": {}}
    return None


# ── Run a single task+run ─────────────────────────────────────────────────────
def run_task(task, run_num, debug=False):
    try:
        reset_resp = requests.post(f"{BASE_URL}/reset?task={task}", timeout=30).json()
    except Exception as e:
        print(f"  {task} run {run_num}: RESET ERROR — {e}")
        return 0.0

    obs = reset_resp.get("observation", {})
    config_text = obs["config"] if isinstance(obs, dict) and "config" in obs else str(obs)

    system_prompt = SYSTEM_MAP.get(task, SYSTEM_DEFAULT)

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user",   "content": f"Audit this AWS config:\n{config_text}"}
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs,
            attention_mask=torch.ones_like(inputs),
            max_new_tokens=3072,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    response_text = tokenizer.decode(
        outputs[0][inputs.shape[1]:], skip_special_tokens=True
    ).strip()

    parsed = repair_json(response_text)
    if parsed is None:
        print(f"  {task} run {run_num}: PARSE FAIL — raw: {response_text[:120]}")
        return 0.0

    if debug:
        print(f"\n  [DEBUG] {task} findings ({len(parsed.get('findings', []))} total):")
        for i, f in enumerate(parsed.get("findings", [])):
            print(f"    {i+1}. {f}")

    try:
        step = requests.post(f"{BASE_URL}/step", json=parsed, timeout=30).json()
        if debug:
            print(f"\n  [DEBUG] scorer response: {json.dumps(step, indent=2)[:800]}")
    except Exception as e:
        print(f"  {task} run {run_num}: STEP ERROR — {e}")
        return 0.0

    reward   = step.get("reward", 0.0)
    findings = len(parsed.get("findings", []))
    print(f"  {task} run {run_num}: {reward:.4f}  ({findings} findings)")
    return reward


tasks = [
    "easy_security_group",
    "medium_s3_policy",
    "medium_lambda_iam",
    "hard_rds_cloudtrail",
    "hard_iam_vpc",
]

DEBUG_TASKS = {"hard_iam_vpc"}

print("=" * 50)
all_scores = {}
for task in tasks:
    scores = [run_task(task, i, debug=(task in DEBUG_TASKS and i == 1)) for i in range(1, 4)]
    avg = sum(scores) / 3
    all_scores[task] = avg
    print(f"  >>> {task} AVG: {avg:.4f}\n")

print("=" * 50)
print(f"{'FINAL SCORES':^50}")
print("=" * 50)
for task, score in all_scores.items():
    bar = "█" * int(score * 20)
    print(f"{task:<30} {score:.4f}  {bar}")