// Command pwninit is a task-bootstrapper for winpwn, the analogue of the
// pwninit tool pwntools users reach for: point it at a challenge directory
// and it (1) prints the recon you'd otherwise run by hand -- arch,
// checksec, sections -- so a new task is legible in one command, and (2)
// scaffolds a minimal solve script (go.mod wired up via a replace directive
// + a bare Spawn/Interactive main.go) so `go run .` works from that
// directory immediately, instead of copying the workspace template by hand.
package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"winpwn"
)

// defaultLibPath is where the winpwn source lives on this machine. There's
// no published module to `go get`, so every generated go.mod needs a
// `replace winpwn => <path>` pointing at a real winpwn checkout; override
// with the WINPWN_HOME environment variable if it ever moves.
const defaultLibPath = `C:\tools\go_pwner`

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintf(os.Stderr, "pwninit: %v\n", err)
		os.Exit(1)
	}
}

func run(args []string) error {
	force := false
	targetArg := ""
	for _, a := range args {
		switch a {
		case "-force":
			force = true
		case "-h", "--help", "help":
			usage()
			return nil
		default:
			targetArg = a
		}
	}

	cwd, err := os.Getwd()
	if err != nil {
		return err
	}

	target, err := resolveTarget(cwd, targetArg)
	if err != nil {
		return err
	}
	targetBase := filepath.Base(target)
	fmt.Printf("[pwninit] target: %s\n\n", targetBase)

	if err := printRecon(target); err != nil {
		return err
	}

	libPath, err := resolveLibPath()
	if err != nil {
		return err
	}

	if err := scaffold(cwd, targetBase, libPath, force); err != nil {
		return err
	}

	fmt.Println("\n[pwninit] wrote go.mod + main.go -- next: go run .")
	return nil
}

func usage() {
	fmt.Fprint(os.Stderr, `pwninit -- scaffold a winpwn solve script for the challenge in this directory

Usage:
  pwninit                 auto-detect the single .exe/.dll in the current directory
  pwninit <target>        use this file explicitly
  pwninit -force          overwrite an existing go.mod/main.go in this directory

Prints checksec/sections recon immediately, then writes a go.mod (replace
winpwn => `+defaultLibPath+`, override via WINPWN_HOME) and a minimal
main.go (Spawn + Interactive, nothing else assumed) ready for "go run .".
`)
}

// resolveTarget returns explicit if set, otherwise the sole .exe/.dll in
// dir -- erroring with the full candidate list if that's ambiguous, the way
// a human would want to know *why* auto-detection refused to guess.
func resolveTarget(dir, explicit string) (string, error) {
	if explicit != "" {
		if _, err := os.Stat(explicit); err != nil {
			return "", fmt.Errorf("target %q: %w", explicit, err)
		}
		return explicit, nil
	}

	var candidates []string
	for _, pattern := range []string{"*.exe", "*.dll"} {
		matches, _ := filepath.Glob(filepath.Join(dir, pattern))
		candidates = append(candidates, matches...)
	}

	switch len(candidates) {
	case 0:
		return "", errors.New("no .exe/.dll found in this directory -- pass the target explicitly: pwninit <target>")
	case 1:
		return candidates[0], nil
	default:
		names := make([]string, len(candidates))
		for i, c := range candidates {
			names[i] = filepath.Base(c)
		}
		return "", fmt.Errorf("multiple binaries found (%s) -- pass the target explicitly: pwninit <target>",
			strings.Join(names, ", "))
	}
}

// printRecon prints the "understand this task in one command" block: arch,
// checksec, and section permissions/entropy.
func printRecon(target string) error {
	pf, err := winpwn.OpenPE(target)
	if err != nil {
		return err
	}
	defer pf.Close()

	is64, err := pf.Is64Bit()
	if err != nil {
		return err
	}
	base, _ := pf.ImageBase()
	entry, _ := pf.EntryPoint()

	arch := "x86"
	if is64 {
		arch = "x64"
	}
	fmt.Printf("Arch:       %s\n", arch)
	fmt.Printf("ImageBase:  0x%X\n", base)
	fmt.Printf("EntryPoint: 0x%X\n", entry)

	r, err := pf.Checksec()
	if err != nil {
		return err
	}
	fmt.Println()
	fmt.Printf("ASLR:           %s\n", yesNo(r.ASLR))
	fmt.Printf("HighEntropyVA:  %s\n", yesNo(r.HighEntropyVA))
	fmt.Printf("DEP/NX:         %s\n", yesNo(r.DEP))
	fmt.Printf("CFG:            %s\n", yesNo(r.CFG))
	if r.SEHApplicable {
		fmt.Printf("SafeSEH:        %s\n", yesNo(r.SafeSEH))
	} else {
		fmt.Printf("SafeSEH:        N/A (x64 uses table-based SEH)\n")
	}
	fmt.Printf("GS (heuristic): %s\n", yesNo(r.GSHeuristic))
	fmt.Printf("Authenticode:   %s\n", yesNo(r.AuthenticodeSigned))

	fmt.Println("\n--- sections ---")
	for _, sec := range pf.Sections() {
		perm := ""
		if sec.IsReadable() {
			perm += "R"
		} else {
			perm += "-"
		}
		if sec.IsWritable() {
			perm += "W"
		} else {
			perm += "-"
		}
		if sec.IsExecutable() {
			perm += "X"
		} else {
			perm += "-"
		}
		entropy, _ := sec.Entropy()
		fmt.Printf("  %-8s %s  VA=0x%-8X Offset=0x%-8X Size=0x%-8X entropy=%.2f\n",
			sec.Name, perm, sec.VirtualAddress, sec.Offset, sec.VirtualSize, entropy)
	}

	if libs, err := pf.ImportedLibs(); err == nil {
		fmt.Println("\n--- imported libs (live image base -- stable until next reboot) ---")
		for _, lib := range libs {
			if lib.Err != nil {
				fmt.Printf("  %-24s (failed to load: %v)\n", lib.Name, lib.Err)
				continue
			}
			fmt.Printf("  %-24s 0x%016X\n", lib.Name, lib.Base)
		}
	}

	return nil
}

func yesNo(b bool) string {
	if b {
		return "Yes"
	}
	return "No"
}

// resolveLibPath finds a real winpwn checkout to point the generated
// go.mod's replace directive at, honoring WINPWN_HOME over the hardcoded
// default so this still works if the library ever moves.
func resolveLibPath() (string, error) {
	path := os.Getenv("WINPWN_HOME")
	if path == "" {
		path = defaultLibPath
	}
	goModPath := filepath.Join(path, "go.mod")
	data, err := os.ReadFile(goModPath)
	if err != nil {
		return "", fmt.Errorf("can't find winpwn at %q (%w) -- set WINPWN_HOME to override", path, err)
	}
	if !strings.Contains(string(data), "module winpwn") {
		return "", fmt.Errorf("%q doesn't look like the winpwn module (go.mod has no \"module winpwn\")", path)
	}
	abs, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	return filepath.ToSlash(abs), nil
}

// scaffold writes go.mod + main.go into dir and runs `go mod tidy` (with
// GOPROXY=off: winpwn's transitive deps are already in the local module
// cache from building winpwn itself, so this never needs network access)
// to fill in the indirect requires/go.sum instead of hardcoding version
// strings that would drift the moment winpwn's own go.mod changes.
func scaffold(dir, targetBase, libPath string, force bool) error {
	goModPath := filepath.Join(dir, "go.mod")
	mainGoPath := filepath.Join(dir, "main.go")

	if !force {
		for _, p := range []string{goModPath, mainGoPath} {
			if _, err := os.Stat(p); err == nil {
				return fmt.Errorf("%s already exists -- pass -force to overwrite", p)
			}
		}
	}

	moduleName := sanitizeModuleName(filepath.Base(dir))

	goMod := fmt.Sprintf("module %s\n\ngo 1.26.2\n\nrequire winpwn v0.0.0\n\nreplace winpwn => %s\n",
		moduleName, libPath)
	if err := os.WriteFile(goModPath, []byte(goMod), 0644); err != nil {
		return err
	}

	if err := os.WriteFile(mainGoPath, []byte(mainGoSkeleton(targetBase)), 0644); err != nil {
		return err
	}

	cmd := exec.Command("go", "mod", "tidy")
	cmd.Dir = dir
	cmd.Env = append(os.Environ(), "GOPROXY=off")
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("go mod tidy: %w\n%s", err, out)
	}
	return nil
}

// sanitizeModuleName turns a directory name into something `go build`
// accepts as a module path -- Go module paths reject spaces and most
// punctuation, and CTF task directories are rarely named with that in mind.
func sanitizeModuleName(name string) string {
	var sb strings.Builder
	for _, r := range name {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '_', r == '-':
			sb.WriteRune(r)
		default:
			sb.WriteRune('_')
		}
	}
	if sb.Len() == 0 {
		return "solve"
	}
	return sb.String()
}

func mainGoSkeleton(targetBase string) string {
	return fmt.Sprintf(`package main

import (
	"log"

	"winpwn"
)

const target = %q

func main() {
	tube, err := winpwn.Spawn(target)
	if err != nil {
		log.Fatalf("Spawn: %%v", err)
	}

	// TODO: exploit here

	tube.Interactive()
}
`, targetBase)
}