package exploit

import (
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/scope"
	"github.com/google/uuid"
)

// Executor safely runs exploitation commands with scope validation and cleanup.
type Executor struct {
	scopeDef *scope.ScopeDefinition
	cleanup  CleanupRegisterer
	dryRun   bool
	safeMode bool
	confirm  ConfirmFunc

	// allowedExecutables, when non-empty, restricts the binaries the
	// executor is willing to launch to a known tool surface. Empty means
	// no allowlist gate (preserves behaviour for tests / callers that
	// don't opt in via WithAllowedExecutables). See #43.
	allowedExecutables map[string]struct{}
}

// ConfirmFunc is the human-in-the-loop hook for assist mode (4.6.4).
// Called BEFORE every executed step (after scope + safe-mode pass).
// Return false to skip the step (treated as a benign no-op, not a
// failure). Return an error to abort the whole campaign.
//
// CLI wires a TTY-backed implementation in cli/scan.go; agents pass
// nil to opt out.
type ConfirmFunc func(step pipeline.AttackStep) (proceed bool, err error)

// CleanupRegisterer is the interface for registering cleanup actions.
type CleanupRegisterer interface {
	Register(ctx context.Context, campaignID uuid.UUID, command, target string) error
}

// NewExecutor creates a new command executor.
func NewExecutor(scopeDef *scope.ScopeDefinition, cleanup CleanupRegisterer, dryRun bool) *Executor {
	return &Executor{
		scopeDef: scopeDef,
		cleanup:  cleanup,
		dryRun:   dryRun,
	}
}

// WithSafeMode enables the safe-mode allowlist on a new executor.
// Use via: exploit.NewExecutor(scope, cleanup, dryRun).WithSafeMode(true).
// Safe-mode rejects commands containing destructive tokens (rm, DROP,
// TRUNCATE, kill, chmod, chown, shutdown) with an actionable error —
// the researcher can still run them if they really want by dropping the
// flag, but the swarm itself will never fire them.
func (e *Executor) WithSafeMode(on bool) *Executor {
	e.safeMode = on
	return e
}

// AllowedToolNames returns the executable names this executor will accept,
// so callers building an attack plan (which happens before Execute is ever
// called) can tell the LLM what's actually usable. Empty slice means no
// allowlist gate is configured.
func (e *Executor) AllowedToolNames() []string {
	names := make([]string, 0, len(e.allowedExecutables)+len(builtinVerbs))
	for n := range e.allowedExecutables {
		names = append(names, n)
	}
	// Built-in verbs are always usable by the planner, even when no binary
	// allowlist is configured — they are the primitive for API-logic attacks.
	for n := range builtinVerbs {
		names = append(names, n)
	}
	return names
}

// WithAllowedExecutables restricts the binaries the executor will
// launch to the given set. The engine derives this list from
// tools.Coordinator.RegisteredToolNames(), so adding a new tool
// adapter automatically widens the allowlist — no hand-maintained
// list to drift. An empty / nil slice disables the gate entirely
// (used by tests). See #43.
func (e *Executor) WithAllowedExecutables(names []string) *Executor {
	if len(names) == 0 {
		e.allowedExecutables = nil
		return e
	}
	e.allowedExecutables = make(map[string]struct{}, len(names))
	for _, n := range names {
		e.allowedExecutables[strings.ToLower(n)] = struct{}{}
	}
	return e
}

// WithConfirm wires assist mode (4.6.4). When non-nil, the executor
// pauses before every command and asks the function to approve. A
// false return skips the step; an error aborts the campaign.
func (e *Executor) WithConfirm(fn ConfirmFunc) *Executor {
	e.confirm = fn
	return e
}

// destructiveTokens is the deny-list safe-mode uses. Case-insensitive
// whole-word match at parse time (parseCommand turns the command into
// []string so 'search' won't match 'sh' for example).
var destructiveTokens = map[string]struct{}{
	"rm":       {},
	"rmdir":    {},
	"shred":    {},
	"kill":     {},
	"killall":  {},
	"chmod":    {},
	"chown":    {},
	"shutdown": {},
	"reboot":   {},
	"halt":     {},
	"mkfs":     {},
	"dd":       {},
	// SQL DDL keywords show up in sqlmap/manual curl payloads; safe mode
	// blocks these to prevent accidental data loss during PoC validation.
	"drop":     {},
	"truncate": {},
	"delete":   {},
}

// Execute runs an attack step with full safety checks. It is the single-step
// entry point; chain variables ({{name}} substitution and --capture) are not
// available here — use ExecuteChain for multi-step attack paths.
func (e *Executor) Execute(ctx context.Context, step pipeline.AttackStep, campaignID uuid.UUID) (*pipeline.ExecutionResult, error) {
	return e.executeWithVars(ctx, step, campaignID, nil)
}

// ExecuteChain runs an ordered attack path, threading a variable store across
// steps: before each step, {{name}} placeholders in its command are replaced
// with values captured by earlier steps; after each step, any --capture
// directives extract response values into the store. This is what makes a
// stateful chain — authenticate, capture the session token, replay a
// privileged request as another user — actually work. Each step still passes
// the full scope / safe-mode / allowlist / assist gate.
//
// A step that errors does not abort the chain: its result is recorded and the
// run continues, matching the swarm's "keep planning around the gap" behavior.
func (e *Executor) ExecuteChain(ctx context.Context, steps []pipeline.AttackStep, campaignID uuid.UUID) ([]*pipeline.ExecutionResult, error) {
	results, _, err := e.ExecuteChainCapturing(ctx, steps, campaignID)
	return results, err
}

// ExecuteChainCapturing is ExecuteChain that also returns the chain's final
// variable store — the captured tokens/ids (plus the seeded {{nonce}} values).
// The adaptive loop uses it to reuse a chain's session and replay the ids it
// discovered across other endpoints (BOLA/IDOR) without re-authenticating.
func (e *Executor) ExecuteChainCapturing(ctx context.Context, steps []pipeline.AttackStep, campaignID uuid.UUID) ([]*pipeline.ExecutionResult, map[string]string, error) {
	vars := make(map[string]string)
	// Seed per-chain nonces so steps that create a resource (e.g. registering a
	// throwaway account before an authenticated attack) can build unique values
	// and avoid "already exists" failures on a persistent target. {{nonce}} is
	// alphanumeric (for emails/usernames); {{nonce_num}} is a 10-digit numeric
	// string for fields that must be digits only, like a phone number.
	now := time.Now().UnixNano()
	vars["nonce"] = strconv.FormatInt(now, 36)
	vars["nonce_num"] = fmt.Sprintf("%010d", now%1e10)
	results := make([]*pipeline.ExecutionResult, 0, len(steps))
	for _, step := range steps {
		if step.Command == "" {
			continue
		}
		s := step
		s.Command = substituteVars(step.Command, vars)
		s.CleanupCommand = substituteVars(step.CleanupCommand, vars)
		res, err := e.executeWithVars(ctx, s, campaignID, vars)
		if res != nil {
			results = append(results, res)
		}
		if err != nil {
			// Abort only on a fatal control-flow error (assist EOF / cleanup
			// registration), never on a per-step execution failure (res != nil).
			if res == nil {
				return results, vars, err
			}
		}
	}
	return results, vars, nil
}

func (e *Executor) executeWithVars(ctx context.Context, step pipeline.AttackStep, campaignID uuid.UUID, vars map[string]string) (*pipeline.ExecutionResult, error) {
	start := time.Now()

	// 1. Scope validation — HARD STOP if fails.
	//
	// The broad ValidateCommand scan extracts every domain/IP-looking substring
	// from the command line and checks each against scope. That is the right
	// safety net for a shelled-out tool, where a target could appear anywhere in
	// the argv — but it produces false positives on a builtin verb like httpreq,
	// whose command carries JSON bodies and --capture selectors (e.g.
	// $.posts.0.author.vehicleid) that look like domains but are never network
	// targets. A builtin does its own precise scope check on the exact host it
	// contacts (runHTTPReq validates opts.url), so that is authoritative here;
	// running the broad scan on top would wrongly block legitimate API-attack
	// bodies and captures.
	if !isBuiltinVerb(firstToken(step.Command)) {
		if err := scope.ValidateCommand(step.Command, *e.scopeDef); err != nil {
			return &pipeline.ExecutionResult{
				StepID:          step.ID,
				CampaignID:      campaignID,
				CommandExecuted: step.Command,
				Output:          fmt.Sprintf("BLOCKED: %s", err),
				Success:         false,
				ExecutedAt:      start,
				DurationMs:      0,
			}, fmt.Errorf("scope violation: %w", err)
		}
	}

	// 1b. Allowlist gate (#43) — when enabled, the LLM is only allowed
	// to invoke binaries that match a registered tool adapter. Blocks
	// shell wrappers (bash, sh -c, python -c) and arbitrary interpreters
	// that a prompt injection could try to reach as a bridge to RCE.
	// Both the primary command and the cleanup command must pass.
	if err := e.checkExecutableAllowed(step.Command); err != nil {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          fmt.Sprintf("BLOCKED: %s", err),
			Success:         false,
			ExecutedAt:      start,
			DurationMs:      0,
		}, err
	}
	if step.CleanupCommand != "" {
		if err := e.checkExecutableAllowed(step.CleanupCommand); err != nil {
			return nil, fmt.Errorf("cleanup command rejected: %w", err)
		}
	}

	// 2. Register cleanup BEFORE execution — only after the cleanup
	// command has passed the same executable policy as the primary step.
	if step.CleanupCommand != "" && e.cleanup != nil {
		if err := e.cleanup.Register(ctx, campaignID, step.CleanupCommand, step.Name); err != nil {
			return nil, fmt.Errorf("failed to register cleanup: %w", err)
		}
	}

	// 3. Dry run mode — return command without executing
	if e.dryRun {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          "[DRY RUN] Command would be: " + step.Command,
			Success:         true,
			ExecutedAt:      start,
			DurationMs:      0,
		}, nil
	}

	// 4. Execute with timeout
	timeout := 120 * time.Second
	execCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// Quote-aware parsing; rejects shell metacharacters (|, >, <, &, ;, $(),
	// backticks, newlines) unless they appear inside quotes. An attack step
	// that needs a real shell must wrap its command in `sh -c "..."` explicitly.
	parts, err := parseCommand(step.Command)
	if err != nil {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          fmt.Sprintf("BLOCKED: %s", err),
			Success:         false,
			ExecutedAt:      start,
			DurationMs:      0,
		}, fmt.Errorf("unsafe command: %w", err)
	}

	// Safe-mode: reject destructive tokens before execution. This is
	// a belt-and-braces check — programs that disallow automated
	// scanning would NEVER want rm / DROP / etc. to fire accidentally.
	if e.safeMode {
		for _, part := range parts {
			// Split each token into words so quoted args like 'drop table users'
			// are checked word-by-word, not as a single opaque string.
			for _, word := range strings.Fields(part) {
				if _, bad := destructiveTokens[strings.ToLower(word)]; bad {
					msg := fmt.Sprintf("safe-mode blocked destructive token %q (drop --safe-mode to override)", word)
					return &pipeline.ExecutionResult{
						StepID: step.ID, CampaignID: campaignID, CommandExecuted: step.Command,
						Output: "BLOCKED: " + msg, Success: false, ExecutedAt: start,
					}, fmt.Errorf("%s", msg)
				}
			}
		}
	}

	// Built-in verbs (httpreq) are handled in-process, after scope + safe-mode
	// (already applied above) and after assist confirmation. Confirm first so a
	// human still gates every request in assist mode.
	if isBuiltinVerb(parts[0]) {
		if e.confirm != nil {
			proceed, cerr := e.confirm(step)
			if cerr != nil {
				return nil, fmt.Errorf("assist confirm: %w", cerr)
			}
			if !proceed {
				return &pipeline.ExecutionResult{
					StepID:          step.ID,
					CampaignID:      campaignID,
					CommandExecuted: step.Command,
					Output:          "[SKIPPED by assist mode]",
					Success:         false,
					ExecutedAt:      start,
					DurationMs:      int(time.Since(start).Milliseconds()),
				}, nil
			}
		}
		if strings.ToLower(parts[0]) == builtinJWT {
			return e.runJWT(execCtx, parts, step, campaignID, start, vars)
		}
		return e.runHTTPReq(execCtx, parts, step, campaignID, start, vars)
	}

	// Assist mode (4.6.4) — pause for human y/N before firing the
	// command. Skips are recorded as a non-fatal step result so the
	// rest of the swarm can keep planning around the gap. An error
	// from confirm aborts the whole campaign (e.g. EOF on stdin →
	// scripted run with no terminal).
	if e.confirm != nil {
		proceed, cerr := e.confirm(step)
		if cerr != nil {
			return nil, fmt.Errorf("assist confirm: %w", cerr)
		}
		if !proceed {
			return &pipeline.ExecutionResult{
				StepID:          step.ID,
				CampaignID:      campaignID,
				CommandExecuted: step.Command,
				Output:          "[SKIPPED by assist mode]",
				Success:         false,
				ExecutedAt:      start,
				DurationMs:      int(time.Since(start).Milliseconds()),
			}, nil
		}
	}

	cmd := exec.CommandContext(execCtx, parts[0], parts[1:]...)
	output, cmdErr := cmd.CombinedOutput()
	err = cmdErr

	result := &pipeline.ExecutionResult{
		StepID:          step.ID,
		CampaignID:      campaignID,
		CommandExecuted: step.Command,
		Output:          string(output),
		ExecutedAt:      start,
		DurationMs:      int(time.Since(start).Milliseconds()),
	}

	if err != nil {
		result.Success = false
		result.Output += "\nError: " + err.Error()
		return result, nil
	}

	// 5. Check output against expected pattern
	if step.ExpectedOutputPattern != "" {
		matched, _ := regexp.MatchString(step.ExpectedOutputPattern, string(output))
		result.Success = matched
	} else {
		result.Success = true
	}

	// 6. Build evidence
	result.Evidence = []pipeline.Evidence{
		{
			Type:        "command_output",
			Content:     string(output),
			Timestamp:   time.Now(),
			Description: fmt.Sprintf("Output from: %s", step.Name),
		},
	}

	return result, nil
}

// DryRun returns the command string with all substitutions applied, without executing.
func (e *Executor) DryRun(step pipeline.AttackStep) string {
	return fmt.Sprintf("[DRY RUN] %s\n  Cleanup: %s\n  Expected: %s",
		step.Command, step.CleanupCommand, step.ExpectedOutputPattern)
}

// checkExecutableAllowed enforces the per-binary allowlist. Returns nil
// when (a) no allowlist is configured (preserves test behaviour and
// callers that don't opt in) or (b) the first token's basename is in the
// allowed set. Uses shellsafe.Parse so quoted-arg edge cases and
// metachar rejection are consistent with the main parse path.
func (e *Executor) checkExecutableAllowed(cmd string) error {
	if len(e.allowedExecutables) == 0 {
		return nil
	}
	parts, err := parseCommand(cmd)
	if err != nil {
		// Parse failure here just means the main parse will also fail
		// with the same error; returning nil keeps a single source of
		// blame for shell-metachar rejections.
		return nil
	}
	exe := strings.ToLower(filepath.Base(strings.TrimSpace(parts[0])))
	if exe == "" {
		return fmt.Errorf("empty executable")
	}
	// Built-in verbs (httpreq) are handled in-process, not by launching a
	// binary, so they are not subject to the binary allowlist — they still
	// pass scope validation and safe-mode in Execute.
	if isBuiltinVerb(exe) {
		return nil
	}
	if _, ok := e.allowedExecutables[exe]; !ok {
		return fmt.Errorf("executable %q is not in the allowed tool surface (use a registered tool adapter or extend the allowlist)", exe)
	}
	return nil
}