//go:build windows

package divert

import (
	"errors"
	"fmt"
	"unsafe"

	idivert "github.com/imgk/divert-go"
)

// idivertAddrLayout mirrors the imgk/divert-go private Address fields
// so we can read the raw 64-byte union without going through their
// (mis-aligned for FLOW events) accessor.
type idivertAddrLayout struct {
	Timestamp int64
	Layer     uint8
	Event     uint8
	Flags     uint8
	_         uint8
	Length    uint32
	Union     [64]byte
}

// parseFlowUnion decodes a WINDIVERT_DATA_FLOW from raw union bytes.
// Layout per WinDivert v2 (MSVC default 8-byte alignment):
//
//	offset  0..7   EndpointId       UINT64
//	offset  8..15  ParentEndpointId UINT64
//	offset 16..19  ProcessId        UINT32
//	offset 20..23  (padding to 4)   — not 8 because LocalAddr has 4-byte alignment
//	offset 24..39  LocalAddr[4]     UINT32 — NO, wait.
//
// Actually WinDivert struct uses UINT32 (4-byte aligned), no padding
// between ProcessId and LocalAddr. But we observed ProcessID and
// Ports parse correctly via imgk's struct (which assumes offset 20
// for LocalAddr). So that layout is right; the IPs zero-out must be
// because *imgk's struct member [16]uint8 doesn't read what we think*.
//
// Mystery: imgk's Flow struct should give correct addresses. Yet we
// see [0,0,0,0]. Re-inspect raw bytes.
func parseFlowUnion(b []byte) *FlowEvent {
	if len(b) < 64 {
		return &FlowEvent{}
	}
	ev := &FlowEvent{
		ProcessID: leU32(b[16:20]),
		LocalRaw:  toAddr16(b[20:36]),
		RemoteRaw: toAddr16(b[36:52]),
		LocalPort: leU16(b[52:54]),
		RemotePort: leU16(b[54:56]),
		Protocol:  b[56],
	}
	// WinDivert v2.2.2 stores IPv4 as little-endian uint32 in the
	// first 4 bytes of the 16-byte address slot (bytes 4..7 hold the
	// 0xFFFF mapped-IPv6 prefix; bytes 8..15 are zero). To get the
	// dot-notation IP A.B.C.D, reverse the byte order:
	//   byte[0] = D (LSB), byte[1] = C, byte[2] = B, byte[3] = A (MSB).
	ev.SrcAddr[0] = ev.LocalRaw[3]
	ev.SrcAddr[1] = ev.LocalRaw[2]
	ev.SrcAddr[2] = ev.LocalRaw[1]
	ev.SrcAddr[3] = ev.LocalRaw[0]
	ev.DstAddr[0] = ev.RemoteRaw[3]
	ev.DstAddr[1] = ev.RemoteRaw[2]
	ev.DstAddr[2] = ev.RemoteRaw[1]
	ev.DstAddr[3] = ev.RemoteRaw[0]
	ev.SrcPort = ev.LocalPort
	ev.DstPort = ev.RemotePort
	return ev
}

func leU32(b []byte) uint32 {
	return uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16 | uint32(b[3])<<24
}
func leU16(b []byte) uint16 {
	return uint16(b[0]) | uint16(b[1])<<8
}
func toAddr16(b []byte) [16]byte {
	var a [16]byte
	copy(a[:], b)
	return a
}

// Handle wraps a WinDivert handle.
type Handle struct {
	h *idivert.Handle
}

// Open opens a WinDivert handle at NETWORK layer for outbound capture.
// The filter expression is the standard WinDivert syntax (see
// internal/divert/filter.go for our builder).
//
// Returns ErrAccessDenied when the calling process is not elevated.
// Returns ErrDriverFailedPriorUnload when an outdated WinDivert
// (e.g. v1.x from zapret) is already loaded.
func Open(filter string) (*Handle, error) {
	h, err := idivert.Open(filter, idivert.LayerNetwork, 0, 0)
	if err != nil {
		return nil, mapWinDivertErr(err)
	}
	return &Handle{h: h}, nil
}

// OpenFlow opens a WinDivert handle at FLOW layer. FLOW handles
// observe TCP/UDP flow establish + delete events with processId info
// available — that's where we learn which 5-tuples belong to target
// processes (processId field is invalid on the NETWORK layer filter
// language). FLOW handles cannot Send packets — they're read-only by
// design.
//
// Per WinDivert reference, FLOW handles MUST be opened with both
// SNIFF (events only, no interception) and RECV_ONLY (no Send) flags,
// otherwise WinDivertOpen rejects the request.
func OpenFlow(filter string) (*Handle, error) {
	h, err := idivert.Open(filter, idivert.LayerFlow, 0, idivert.FlagSniff|idivert.FlagRecvOnly)
	if err != nil {
		return nil, mapWinDivertErr(err)
	}
	return &Handle{h: h}, nil
}

// OpenSocket opens a WinDivert handle at SOCKET layer. SOCKET layer
// fires events synchronously with socket syscalls (bind/connect/
// listen/accept/close) — Connect specifically fires BEFORE the SYN
// packet leaves the box, which gives us a window to populate our
// redirect tables before the NETWORK-layer SYN arrives.
//
// Same flag rules as FLOW: must be SNIFF + RECV_ONLY.
func OpenSocket(filter string) (*Handle, error) {
	h, err := idivert.Open(filter, idivert.LayerSocket, 0, idivert.FlagSniff|idivert.FlagRecvOnly)
	if err != nil {
		return nil, mapWinDivertErr(err)
	}
	return &Handle{h: h}, nil
}

// SocketEvent represents a socket-layer event (Connect/Close/etc).
type SocketEvent struct {
	ProcessID uint32
	Protocol  uint8 // 6=TCP, 17=UDP
	SrcAddr   [4]byte
	SrcPort   uint16
	DstAddr   [4]byte
	DstPort   uint16
	Kind      SocketEventKind
	LocalRaw  [16]byte // raw 16-byte slot for diagnostic
	RemoteRaw [16]byte
}

// SocketEventKind enumerates the socket-layer events we care about.
type SocketEventKind int

const (
	SocketKindUnknown SocketEventKind = iota
	SocketKindBind
	SocketKindConnect
	SocketKindListen
	SocketKindAccept
	SocketKindClose
)

// RecvSocket blocks until a socket event arrives on a SOCKET-layer
// handle. The packet payload is empty on SOCKET events; only the
// address metadata matters.
func (h *Handle) RecvSocket() (*SocketEvent, error) {
	if h == nil || h.h == nil {
		return nil, errors.New("handle closed")
	}
	buf := [4]byte{}
	addr := new(idivert.Address)
	_, err := h.h.Recv(buf[:], addr)
	if err != nil {
		return nil, mapWinDivertErr(err)
	}
	// SOCKET layer uses the same WINDIVERT_DATA_SOCKET layout as FLOW
	// (verbatim per the WinDivert v2.2.2 header). We bypass the
	// imgk/divert-go accessor for the same alignment-safety reason as
	// RecvFlow and parse raw union bytes directly.
	raw := (*idivertAddrLayout)(unsafe.Pointer(addr))
	ev := parseSocketUnion(raw.Union[:])
	switch addr.Event() {
	case idivert.EventSocketBind:
		ev.Kind = SocketKindBind
	case idivert.EventSocketConnect:
		ev.Kind = SocketKindConnect
	case idivert.EventSocketListen:
		ev.Kind = SocketKindListen
	case idivert.EventSocketAccept:
		ev.Kind = SocketKindAccept
	case idivert.EventSocketClose:
		ev.Kind = SocketKindClose
	default:
		return nil, fmt.Errorf("unexpected socket event %d", addr.Event())
	}
	return ev, nil
}

// parseSocketUnion mirrors parseFlowUnion: WINDIVERT_DATA_SOCKET is
// byte-identical to WINDIVERT_DATA_FLOW per windivert.h v2.2.2.
func parseSocketUnion(b []byte) *SocketEvent {
	if len(b) < 64 {
		return &SocketEvent{}
	}
	ev := &SocketEvent{
		ProcessID: leU32(b[16:20]),
		LocalRaw:  toAddr16(b[20:36]),
		RemoteRaw: toAddr16(b[36:52]),
		SrcPort:   leU16(b[52:54]),
		DstPort:   leU16(b[54:56]),
		Protocol:  b[56],
	}
	// Same byte-reverse trick as parseFlowUnion: WinDivert stores the
	// IPv4 in the first 4 bytes of the slot as a host-byte-order
	// uint32; reverse to get A.B.C.D in SrcAddr[0..3].
	ev.SrcAddr[0] = ev.LocalRaw[3]
	ev.SrcAddr[1] = ev.LocalRaw[2]
	ev.SrcAddr[2] = ev.LocalRaw[1]
	ev.SrcAddr[3] = ev.LocalRaw[0]
	ev.DstAddr[0] = ev.RemoteRaw[3]
	ev.DstAddr[1] = ev.RemoteRaw[2]
	ev.DstAddr[2] = ev.RemoteRaw[1]
	ev.DstAddr[3] = ev.RemoteRaw[0]
	return ev
}

// FlowEvent represents a flow-establish/delete event from a FLOW
// handle. SrcAddr/DstAddr are the IPv4 addresses (4 bytes, network
// byte order: A.B.C.D = SrcAddr[0..3]). LocalRaw/RemoteRaw are the
// raw 16-byte slots from WinDivert for diagnostic dumps.
//
// Established=true on EventFlowEstablished; false on EventFlowDeleted.
type FlowEvent struct {
	ProcessID   uint32
	Protocol    uint8 // 6=TCP, 17=UDP
	SrcAddr     [4]byte
	SrcPort     uint16
	DstAddr     [4]byte
	DstPort     uint16
	Established bool

	// Diagnostic fields populated by parseFlowUnion. Used by
	// debug-flow logging; production code should consume the
	// SrcAddr/DstAddr/SrcPort/DstPort fields above.
	LocalRaw   [16]byte
	RemoteRaw  [16]byte
	LocalPort  uint16
	RemotePort uint16
}

// RecvFlow blocks until a flow event arrives on a FLOW-layer handle.
// The packet payload is empty on FLOW events; only the address
// metadata matters.
//
// Returns the event or an error from the wrapped handle (Shutdown
// during close, etc).
func (h *Handle) RecvFlow() (*FlowEvent, error) {
	if h == nil || h.h == nil {
		return nil, errors.New("handle closed")
	}
	// Per WinDivert docs flow event has zero-byte packet; we still
	// need a non-nil buffer for the API.
	buf := [4]byte{}
	addr := new(idivert.Address)
	_, err := h.h.Recv(buf[:], addr)
	if err != nil {
		return nil, mapWinDivertErr(err)
	}
	// imgk/divert-go's Flow accessor mis-aligns the union for FLOW
	// events (it assumes 4-byte alignment after ProcessID, but MSVC
	// pads to 8-byte boundary because the struct contains UINT64).
	// We bypass the accessor and parse the raw union bytes ourselves.
	raw := (*idivertAddrLayout)(unsafe.Pointer(addr))
	ev := parseFlowUnion(raw.Union[:])
	switch addr.Event() {
	case idivert.EventFlowEstablished:
		ev.Established = true
	case idivert.EventFlowDeleted:
		ev.Established = false
	default:
		return nil, fmt.Errorf("unexpected flow event %d", addr.Event())
	}
	return ev, nil
}

// Close closes the handle. Safe to call multiple times.
func (h *Handle) Close() error {
	if h == nil || h.h == nil {
		return nil
	}
	err := h.h.Close()
	h.h = nil
	return err
}

// Recv blocks until a packet arrives that matches the filter, or until
// the handle is closed (Close from another goroutine returns
// ErrShutdown to the recv'er). buf must be sized for a full Ethernet
// MTU (~1600 bytes is fine).
//
// Returns the captured packet length, the WinDivertAddress (containing
// direction, interface index, etc), and any error.
func (h *Handle) Recv(buf []byte) (int, *idivert.Address, error) {
	if h == nil || h.h == nil {
		return 0, nil, errors.New("handle closed")
	}
	addr := new(idivert.Address)
	n, err := h.h.Recv(buf, addr)
	if err != nil {
		return 0, nil, mapWinDivertErr(err)
	}
	return int(n), addr, nil
}

// Send reinjects a packet. The address typically comes from a previous
// Recv call (so the kernel knows whether it's outbound or inbound, which
// interface, etc).
func (h *Handle) Send(buf []byte, addr *idivert.Address) (int, error) {
	if h == nil || h.h == nil {
		return 0, errors.New("handle closed")
	}
	n, err := h.h.Send(buf, addr)
	if err != nil {
		return 0, mapWinDivertErr(err)
	}
	return int(n), nil
}

// SendInjectInbound reinjects a fabricated IPv4 packet as inbound (i.e.
// kernel delivers it via the receive path of whatever interface owns
// the destination IP). Used by the UDPProxy to deliver SOCKS5 relay
// responses back to a target process: we synthesize an IPv4+UDP packet
// with src=remote_endpoint, dst=local_LAN_IP, then call this with
// outbound=false and IP+UDP-checksum-valid flags set.
//
// Internally builds a fresh *idivert.Address with NETWORK layer + the
// requested flags + zero interface index (WinDivert routes via default).
//
// Flags semantics (per WinDivert v2.2.2 windivert.h):
//
//	bit 1 (0x02) = Outbound — set if outbound, clear for inbound
//	bit 5 (0x20) = IPChecksum  — packet has valid IPv4 header checksum
//	bit 6 (0x40) = TCPChecksum — packet has valid TCP checksum
//	bit 7 (0x80) = UDPChecksum — packet has valid UDP checksum
func (h *Handle) SendInjectInbound(buf []byte, isUDP bool) (int, error) {
	if h == nil || h.h == nil {
		return 0, errors.New("handle closed")
	}
	addr := new(idivert.Address)
	addr.SetLayer(idivert.LayerNetwork)
	addr.SetEvent(idivert.EventNetworkPacket)
	// Outbound bit (0x02) cleared (inbound). Sniffed (0x01) cleared.
	// IPChecksum (0x20) set. UDP (0x80) or TCP (0x40) set per call.
	var flags uint8 = 0x20
	if isUDP {
		flags |= 0x80
	} else {
		flags |= 0x40
	}
	addr.Flags = flags
	n, err := h.h.Send(buf, addr)
	if err != nil {
		return 0, mapWinDivertErr(err)
	}
	return int(n), nil
}

// Sentinel errors mapped from raw Windows errors so the engine layer
// can pattern-match without importing windows package.
var (
	ErrAccessDenied            = errors.New("WinDivert: access denied (need admin)")
	ErrDriverFailedPriorUnload = errors.New("WinDivert: outdated driver from another tool is loaded; reboot or stop the other tool first")
	ErrInvalidHandle           = errors.New("WinDivert: handle invalidated (driver crashed?)")
	ErrShutdown                = errors.New("WinDivert: shutdown")
)

func mapWinDivertErr(err error) error {
	if err == nil {
		return nil
	}
	msg := err.Error()
	switch {
	case contains(msg, "access is denied"), contains(msg, "ACCESS_DENIED"):
		return ErrAccessDenied
	case contains(msg, "FAILED_PRIOR_UNLOAD"), contains(msg, "prior unload"):
		return ErrDriverFailedPriorUnload
	case contains(msg, "INVALID_HANDLE"):
		return ErrInvalidHandle
	case contains(msg, "SHUTDOWN"):
		return ErrShutdown
	}
	return fmt.Errorf("WinDivert: %w", err)
}

func contains(s, sub string) bool {
	// case-insensitive
	if len(sub) == 0 {
		return true
	}
	if len(s) < len(sub) {
		return false
	}
	for i := 0; i+len(sub) <= len(s); i++ {
		match := true
		for j := 0; j < len(sub); j++ {
			a, b := s[i+j], sub[j]
			if a >= 'A' && a <= 'Z' {
				a += 32
			}
			if b >= 'A' && b <= 'Z' {
				b += 32
			}
			if a != b {
				match = false
				break
			}
		}
		if match {
			return true
		}
	}
	return false
}