package client

import (
	"encoding/base64"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"crypto-miner-agent/deploy"
)

type commandResult struct {
	action  string
	success bool
	message string
}

func captureCommandResult(t *testing.T, c *AgentClient) (done <-chan struct{}, result *commandResult) {
	t.Helper()
	ch := make(chan struct{})
	var mu sync.Mutex
	out := &commandResult{}
	c.commandResultHook = func(action string, success bool, message string) {
		mu.Lock()
		out.action = action
		out.success = success
		out.message = message
		mu.Unlock()
		close(ch)
	}
	t.Cleanup(func() { c.commandResultHook = nil })
	return ch, out
}

// traversalPaths lists every variant that must be rejected by upload/download.
var traversalPaths = []struct {
	name string
	path string
}{
	{name: "unix_relative", path: "../../etc/passwd"},
	{name: "windows_relative", path: `..\..\Windows\System32\config\sam`},
	{name: "embedded_traversal", path: "uploads/../../outside.txt"},
	{name: "absolute_with_traversal", path: "/var/log/../../etc/shadow"},
	{name: "dotdot_only", path: ".."},
	{name: "mixed_separators", path: `foo\..\bar\..\..\secret`},
	{name: "tilde_traversal", path: "~/../../etc/passwd"},
	{name: "desktop_prefix_traversal", path: "@desktop/../../outside.txt"},
	{name: "desktop_colon_traversal", path: "desktop:../../payload.bin"},
	{name: "midpath_dotdot", path: "safe/inner/../../../etc/shadow"},
}

func TestUploadCommandRejectsPathTraversal(t *testing.T) {
	data := base64.StdEncoding.EncodeToString([]byte("payload"))

	cases := traversalPaths

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := deploy.ResolveRemotePath(tc.path)
			if err == nil {
				t.Fatalf("ResolveRemotePath(%q) should reject traversal", tc.path)
			}
			if !strings.Contains(err.Error(), "path traversal") {
				t.Fatalf("ResolveRemotePath(%q) error = %q, want path traversal rejection", tc.path, err.Error())
			}

			c := newTestClient(t)
			done, got := captureCommandResult(t, c)
			c.handleCommand("upload", 0, "", tc.path, data, "")
			<-done

			if got.action != "upload" {
				t.Fatalf("action = %q, want upload", got.action)
			}
			if got.success {
				t.Fatalf("upload with %q should fail (success=true, message=%q)", tc.path, got.message)
			}
			if !strings.Contains(got.message, "path traversal") {
				t.Fatalf("message = %q, want path traversal error from ResolveRemotePath", got.message)
			}
		})
	}
}


func TestDownloadCommandRejectsPathTraversal(t *testing.T) {
	cases := traversalPaths

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := deploy.ResolveRemotePath(tc.path)
			if err == nil {
				t.Fatalf("ResolveRemotePath(%q) should reject traversal", tc.path)
			}
			if !strings.Contains(err.Error(), "path traversal") {
				t.Fatalf("ResolveRemotePath(%q) error = %q, want path traversal rejection", tc.path, err.Error())
			}

			c := newTestClient(t)
			done, got := captureCommandResult(t, c)
			c.handleCommand("download", 0, "", tc.path, "", "")
			<-done

			if got.action != "download" {
				t.Fatalf("action = %q, want download", got.action)
			}
			if got.success {
				t.Fatalf("download with %q should fail (success=true, message=%q)", tc.path, got.message)
			}
			if !strings.Contains(got.message, "path traversal") {
				t.Fatalf("message = %q, want path traversal error from ResolveRemotePath", got.message)
			}
		})
	}
}
func TestUploadCommandRejectsInvalidBase64(t *testing.T) {
	c := newTestClient(t)
	done, got := captureCommandResult(t, c)
	c.handleCommand("upload", 0, "", "notes.txt", "not-valid-base64!!!", "")
	<-done
	if got.success {
		t.Fatalf("expected base64 failure, got %q", got.message)
	}
	if !strings.Contains(got.message, "invalid base64") {
		t.Fatalf("message=%q", got.message)
	}
}

func TestDownloadCommandReadsSafeFile(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "payload.bin")
	content := []byte("download-me")
	if err := os.WriteFile(src, content, 0644); err != nil {
		t.Fatal(err)
	}
	data := base64.StdEncoding.EncodeToString(content)

	c := newTestClient(t)
	done, got := captureCommandResult(t, c)
	c.handleCommand("download", 0, "", src, "", "")
	<-done
	if !got.success {
		t.Fatalf("download failed: %s", got.message)
	}
	if got.message != data {
		t.Fatalf("encoded mismatch: got len=%d want len=%d", len(got.message), len(data))
	}
}

func TestUploadCommandAcceptsSafePath(t *testing.T) {
	dir := t.TempDir()
	dest := dir + "/notes.txt"
	data := base64.StdEncoding.EncodeToString([]byte("ok"))

	c := newTestClient(t)
	done, got := captureCommandResult(t, c)
	c.handleCommand("upload", 0, "", dest, data, "")
	<-done

	if !got.success {
		t.Fatalf("safe upload failed: %s", got.message)
	}
	if got.action != "upload" {
		t.Fatalf("action = %q, want upload", got.action)
	}
}