1package tools
2
3import (
4 "context"
5 "encoding/base64"
6 "encoding/binary"
7 "encoding/json"
8 "fmt"
9 "os"
10 "os/exec"
11 "reflect"
12 "runtime"
13 "strings"
14 "testing"
15 "time"
16 "unicode/utf16"
17
18 "github.com/patriceckhart/zot/packages/provider"
19)
20
21func TestResolvePowerShell(t *testing.T) {
22 for _, tc := range []struct{ name, goos, available, want string }{
23 {"preferred", "windows", "pwsh.exe", "pwsh.exe"},
24 {"fallback", "windows", "powershell.exe", "powershell.exe"},
25 {"missing", "windows", "", ""},
26 {"unsupported", "linux", "pwsh.exe", ""},
27 } {
28 t.Run(tc.name, func(t *testing.T) {
29 var calls []string
30 got, err := resolvePowerShell(tc.goos, func(name string) (string, error) {
31 calls = append(calls, name)
32 if name == tc.available {
33 return name, nil
34 }
35 return "", fmt.Errorf("not found")
36 })
37 if got != tc.want || (err != nil) != (tc.want == "") {
38 t.Fatalf("got %q, %v", got, err)
39 }
40 if tc.goos != "windows" && len(calls) != 0 {
41 t.Fatal("searched on unsupported OS")
42 }
43 if tc.name == "preferred" && len(calls) != 1 {
44 t.Fatal("did not prefer pwsh")
45 }
46 })
47 }
48}
49
50func TestPowerShellArgs(t *testing.T) {
51 command := "Write-Output 'Grüße 世界'\nWrite-Output \"quoted\""
52 args := powerShellArgs(command)
53 want := []string{"-NoLogo", "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-OutputFormat", "Text", "-EncodedCommand"}
54 if !reflect.DeepEqual(args[:len(args)-1], want) {
55 t.Fatalf("args: %v", args)
56 }
57 data, err := base64.StdEncoding.DecodeString(args[len(args)-1])
58 if err != nil {
59 t.Fatal(err)
60 }
61 units := make([]uint16, len(data)/2)
62 for i := range units {
63 units[i] = binary.LittleEndian.Uint16(data[i*2:])
64 }
65 script := string(utf16.Decode(units))
66 if !strings.HasSuffix(script, "\n"+command) || !strings.Contains(script, "UTF8Encoding") {
67 t.Fatalf("script: %q", script)
68 }
69}
70
71func TestPowerShellPermissions(t *testing.T) {
72 sb := NewSandbox(t.TempDir())
73 tool := &PowerShellTool{Sandbox: sb}
74 if err := tool.checkPermission(""); err != nil {
75 t.Fatal(err)
76 }
77 sb.Lock()
78 if _, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"Get-Location"}`), nil); err == nil || !strings.Contains(err.Error(), "jailed") {
79 t.Fatalf("got %v", err)
80 }
81 sb.Unlock()
82 sb.Permissions = &PermissionSet{}
83 for _, mode := range []string{"", "none", "allowlist", "invalid", "ask"} {
84 sb.Permissions.Bash.Mode = mode
85 sb.Permissions.Bash.Allow = []string{"Get-Location"}
86 if err := tool.checkPermission("Get-Location"); (err == nil) != (mode == "ask") {
87 t.Fatalf("mode %q: %v", mode, err)
88 }
89 }
90}
91
92func TestPowerShellInvalidArgs(t *testing.T) {
93 tool := &PowerShellTool{}
94 for _, raw := range []string{`{`, `{}`, `{"command":" "}`, `{"command":123}`} {
95 if _, err := tool.Execute(context.Background(), json.RawMessage(raw), nil); err == nil {
96 t.Fatalf("accepted %s", raw)
97 }
98 }
99 if runtime.GOOS != "windows" {
100 _, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"Get-Location"}`), nil)
101 if err == nil || !strings.Contains(err.Error(), "only available on Windows") {
102 t.Fatalf("got %v", err)
103 }
104 }
105}
106
107func TestPowerShellExecution(t *testing.T) {
108 if runtime.GOOS != "windows" {
109 t.Skip("native Windows execution")
110 }
111 if _, err := resolvePowerShell(runtime.GOOS, exec.LookPath); err != nil {
112 t.Skip(err)
113 }
114 cwd := t.TempDir()
115 tool := &PowerShellTool{CWD: cwd}
116 for _, tc := range []struct {
117 name, command, want string
118 failure bool
119 timeout int
120 }{
121 {"unicode and quoting", "Write-Output 'Grüße 世界'; Write-Output '\"quoted\"'", "Grüße 世界", false, 0},
122 {"cwd", "(Get-Location).Path", "", false, 0},
123 {"stderr", "[Console]::Error.WriteLine('stderr output')", "stderr output", false, 0},
124 {"exit", "exit 7", "[exit 7]", true, 0},
125 {"error", "throw 'failure'", "failure", true, 0},
126 {"timeout", "Start-Sleep -Seconds 60", "", true, 1},
127 } {
128 t.Run(tc.name, func(t *testing.T) {
129 ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
130 defer cancel()
131 var streamed strings.Builder
132 res, err := tool.Execute(ctx, mustJSON(t, map[string]any{"command": tc.command, "timeout": tc.timeout}), func(s string) { streamed.WriteString(s) })
133 if err != nil {
134 t.Fatal(err)
135 }
136 text := res.Content[0].(provider.TextBlock).Text
137 if res.IsError != tc.failure || !strings.Contains(text, tc.want) || !strings.HasPrefix(text, "PS> ") {
138 t.Fatalf("result: %+v, %s", res, text)
139 }
140 if tc.name == "cwd" {
141 // Windows temp paths can use short names or different casing
142 // than Get-Location. Compare directory identity, not spelling.
143 gotPath := strings.TrimSpace(streamed.String())
144 gotInfo, err := os.Stat(gotPath)
145 if err != nil {
146 t.Fatalf("stat reported cwd %q (expected %q): %v", gotPath, cwd, err)
147 }
148 wantInfo, err := os.Stat(cwd)
149 if err != nil {
150 t.Fatal(err)
151 }
152 if !os.SameFile(gotInfo, wantInfo) {
153 t.Fatalf("reported cwd %q is not the expected directory %q", gotPath, cwd)
154 }
155 }
156 if tc.name == "unicode and quoting" && !strings.Contains(streamed.String(), "Grüße 世界") {
157 t.Fatalf("stream: %q", streamed.String())
158 }
159 })
160 }
161 ctx, cancel := context.WithCancel(context.Background())
162 cancel()
163 res, err := tool.Execute(ctx, json.RawMessage(`{"command":"Get-Location"}`), nil)
164 if err == nil && !res.IsError {
165 t.Fatal("cancellation succeeded")
166 }
167}