mirror of
https://github.com/BishopFox/sliver.git
synced 2026-11-09 17:09:39 -10:00
Compare commits
52
Commits
ai/tui
...
fix/linux-ver
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6d51a23cc9 | ||
|
|
3d6eec47be | ||
|
|
464fb7008c | ||
|
|
ffa5f89f5f | ||
|
|
f19cbdc957 | ||
|
|
83bfc8d8e6 | ||
|
|
365bb3596e | ||
|
|
9ee96c2ce6 | ||
|
|
e776e266d0 | ||
|
|
d080c8a646 | ||
|
|
cb6402245e | ||
|
|
e4357277ec | ||
|
|
5dc2ab1490 | ||
|
|
9f034a9956 | ||
|
|
f700535997 | ||
|
|
31bca034d8 | ||
|
|
fac4ad5fb7 | ||
|
|
c128d33514 | ||
|
|
4a37262bcf | ||
|
|
a2295112bd | ||
|
|
834dff0f6e | ||
|
|
f4a0431fdd | ||
|
|
d072b328c0 | ||
|
|
435bbcbd7e | ||
|
|
84dc28e584 | ||
|
|
639592cb02 | ||
|
|
8c7a1a6b40 | ||
|
|
06274c4061 | ||
|
|
255477e572 | ||
|
|
0acc6911d8 | ||
|
|
4ca026906b | ||
|
|
97a249e28e | ||
|
|
9f331687d3 | ||
|
|
14d68d6fcd | ||
|
|
05c5334fea | ||
|
|
d277a14c57 | ||
|
|
a7b5dbe6ee | ||
|
|
c188dc1e69 | ||
|
|
9fc9d1762e | ||
|
|
4ef86448d0 | ||
|
|
02117a6072 | ||
|
|
c272c8abac | ||
|
|
cdf2412b39 | ||
|
|
40f52791d3 | ||
|
|
a5d8bd16b9 | ||
|
|
bf113f2151 | ||
|
|
4cd010afb9 | ||
|
|
04b0bdd199 | ||
|
|
9682db93b3 | ||
|
|
1d7ffb4ac2 | ||
|
|
c5e7cd6354 | ||
|
|
b20ba7284b |
No files matched your search
@@ -1,5 +1,7 @@
|
||||
/crush
|
||||
|
||||
/examples
|
||||
|
||||
/.vscode
|
||||
/.vscode/*
|
||||
/.devcontainer
|
||||
|
||||
@@ -1,22 +1,33 @@
|
||||
#
|
||||
# Makefile for Sliver
|
||||
#
|
||||
|
||||
ifneq ($(OS),Windows_NT)
|
||||
|
||||
GO ?= go
|
||||
ARTIFACT_SUFFIX ?=
|
||||
ENV =
|
||||
TAGS ?= -tags go_sqlite
|
||||
CGO_ENABLED = 0
|
||||
|
||||
ifneq (,$(findstring cgo_sqlite,$(TAGS)))
|
||||
CGO_ENABLED = 1
|
||||
endif
|
||||
|
||||
#
|
||||
# Prerequisites
|
||||
#
|
||||
#
|
||||
# Makefile for Sliver
|
||||
#
|
||||
|
||||
GO ?= go
|
||||
ARTIFACT_SUFFIX ?=
|
||||
ENV =
|
||||
TAGS ?= -tags go_sqlite
|
||||
CGO_ENABLED = 0
|
||||
|
||||
ifneq (,$(findstring cgo_sqlite,$(TAGS)))
|
||||
CGO_ENABLED = 1
|
||||
endif
|
||||
|
||||
MIN_SUPPORTED_GO_MAJOR_VERSION = 1
|
||||
MIN_SUPPORTED_GO_MINOR_VERSION = 25
|
||||
GO_VERSION_VALIDATION_ERR_MSG = Golang version is not supported, please update to at least $(MIN_SUPPORTED_GO_MAJOR_VERSION).$(MIN_SUPPORTED_GO_MINOR_VERSION)
|
||||
|
||||
SLIVER_PUBLIC_KEY ?= RWTZPg959v3b7tLG7VzKHRB1/QT+d3c71Uzetfa44qAoX5rH7mGoQTTR
|
||||
ARMORY_PUBLIC_KEY ?= RWSBpxpRWDrD7Fe+VvRE3c2VEDC2NK80rlNCj+BX0gz44Xw07r6KQD9L
|
||||
ARMORY_REPO_URL ?= https://api.github.com/repos/sliverarmory/armory/releases
|
||||
CLIENT_ASSETS_PKG = github.com/bishopfox/sliver/client/assets
|
||||
SLIVER_UPDATE_PKG = github.com/bishopfox/sliver/client/command/update
|
||||
PB_COMPILERS = protoc protoc-gen-go protoc-gen-go-grpc
|
||||
|
||||
ifneq ($(OS),Windows_NT)
|
||||
|
||||
#
|
||||
# Prerequisites
|
||||
#
|
||||
# https://stackoverflow.com/questions/5618615/check-if-a-program-exists-from-a-makefile
|
||||
EXECUTABLES = uname sed git date cut $(GO)
|
||||
K := $(foreach exec,$(EXECUTABLES),\
|
||||
@@ -25,21 +36,12 @@ K := $(foreach exec,$(EXECUTABLES),\
|
||||
#
|
||||
# Build Information
|
||||
#
|
||||
GO_MAJOR_VERSION = $(shell $(GO) version | cut -c 14- | cut -d' ' -f1 | cut -d'.' -f1)
|
||||
GO_MINOR_VERSION = $(shell $(GO) version | cut -c 14- | cut -d' ' -f1 | cut -d'.' -f2)
|
||||
MIN_SUPPORTED_GO_MAJOR_VERSION = 1
|
||||
MIN_SUPPORTED_GO_MINOR_VERSION = 25
|
||||
GO_VERSION_VALIDATION_ERR_MSG = Golang version is not supported, please update to at least $(MIN_SUPPORTED_GO_MAJOR_VERSION).$(MIN_SUPPORTED_GO_MINOR_VERSION)
|
||||
|
||||
SLIVER_PUBLIC_KEY ?= RWTZPg959v3b7tLG7VzKHRB1/QT+d3c71Uzetfa44qAoX5rH7mGoQTTR
|
||||
ARMORY_PUBLIC_KEY ?= RWSBpxpRWDrD7Fe+VvRE3c2VEDC2NK80rlNCj+BX0gz44Xw07r6KQD9L
|
||||
ARMORY_REPO_URL ?= https://api.github.com/repos/sliverarmory/armory/releases
|
||||
CLIENT_ASSETS_PKG = github.com/bishopfox/sliver/client/assets
|
||||
SLIVER_UPDATE_PKG = github.com/bishopfox/sliver/client/command/update
|
||||
|
||||
LDFLAGS = -ldflags "-s -w \
|
||||
-X $(SLIVER_UPDATE_PKG).SliverPublicKey=$(SLIVER_PUBLIC_KEY) \
|
||||
-X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) \
|
||||
GO_MAJOR_VERSION = $(shell $(GO) version | cut -c 14- | cut -d' ' -f1 | cut -d'.' -f1)
|
||||
GO_MINOR_VERSION = $(shell $(GO) version | cut -c 14- | cut -d' ' -f1 | cut -d'.' -f2)
|
||||
|
||||
LDFLAGS = -ldflags "-s -w \
|
||||
-X $(SLIVER_UPDATE_PKG).SliverPublicKey=$(SLIVER_PUBLIC_KEY) \
|
||||
-X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) \
|
||||
-X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
# Debug builds shouldn't be stripped (-s -w flags)
|
||||
@@ -49,15 +51,14 @@ LDFLAGS_DEBUG = -ldflags "-X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMOR
|
||||
SED_INPLACE := sed -i
|
||||
STATIC_TARGET := linux
|
||||
|
||||
UNAME_S := $(shell uname -s)
|
||||
UNAME_P := $(shell uname -p)
|
||||
|
||||
# Programs required for generating protobuf/grpc files
|
||||
PB_COMPILERS = protoc protoc-gen-go protoc-gen-go-grpc
|
||||
ifeq ($(MAKECMDGOALS), pb)
|
||||
K := $(foreach exec,$(PB_COMPILERS),\
|
||||
$(if $(shell which $(exec)),some string,$(error "Missing protobuf util $(exec) in PATH")))
|
||||
endif
|
||||
UNAME_S := $(shell uname -s)
|
||||
UNAME_P := $(shell uname -p)
|
||||
|
||||
# Programs required for generating protobuf/grpc files
|
||||
ifeq ($(MAKECMDGOALS), pb)
|
||||
K := $(foreach exec,$(PB_COMPILERS),\
|
||||
$(if $(shell which $(exec)),some string,$(error "Missing protobuf util $(exec) in PATH")))
|
||||
endif
|
||||
|
||||
# *** Darwin ***
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
@@ -186,43 +187,23 @@ clean:
|
||||
touch ./.downloaded_assets
|
||||
|
||||
|
||||
#
|
||||
# >>> WINDOWS <<<
|
||||
#
|
||||
else
|
||||
|
||||
GO ?= go
|
||||
ARTIFACT_SUFFIX ?=
|
||||
ENV =
|
||||
TAGS ?= -tags go_sqlite
|
||||
CGO_ENABLED = 0
|
||||
|
||||
ifneq (,$(findstring cgo_sqlite,$(TAGS)))
|
||||
CGO_ENABLED = 1
|
||||
endif
|
||||
|
||||
SHELL := cmd.exe
|
||||
.SHELLFLAGS := /C
|
||||
|
||||
GO_VERSION := $(patsubst go%,%,$(strip $(shell $(GO) env GOVERSION)))
|
||||
GO_MAJOR_VERSION := $(word 1,$(subst ., ,$(GO_VERSION)))
|
||||
GO_MINOR_VERSION := $(word 2,$(subst ., ,$(GO_VERSION)))
|
||||
MIN_SUPPORTED_GO_MAJOR_VERSION = 1
|
||||
MIN_SUPPORTED_GO_MINOR_VERSION = 25
|
||||
GO_VERSION_VALIDATION_ERR_MSG = Golang version is not supported, please update to at least $(MIN_SUPPORTED_GO_MAJOR_VERSION).$(MIN_SUPPORTED_GO_MINOR_VERSION)
|
||||
|
||||
SLIVER_PUBLIC_KEY ?= RWTZPg959v3b7tLG7VzKHRB1/QT+d3c71Uzetfa44qAoX5rH7mGoQTTR
|
||||
ARMORY_PUBLIC_KEY ?= RWSBpxpRWDrD7Fe+VvRE3c2VEDC2NK80rlNCj+BX0gz44Xw07r6KQD9L
|
||||
ARMORY_REPO_URL ?= https://api.github.com/repos/sliverarmory/armory/releases
|
||||
CLIENT_ASSETS_PKG = github.com/bishopfox/sliver/client/assets
|
||||
SLIVER_UPDATE_PKG = github.com/bishopfox/sliver/client/command/update
|
||||
|
||||
LDFLAGS = -ldflags "-s -w -X $(SLIVER_UPDATE_PKG).SliverPublicKey=$(SLIVER_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
LDFLAGS_DEBUG = -ldflags "-X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
PB_COMPILERS = protoc protoc-gen-go protoc-gen-go-grpc
|
||||
COMMA := ,
|
||||
#
|
||||
# >>> WINDOWS <<<
|
||||
#
|
||||
else
|
||||
|
||||
SHELL := cmd.exe
|
||||
.SHELLFLAGS := /C
|
||||
|
||||
GO_VERSION := $(patsubst go%,%,$(strip $(shell $(GO) env GOVERSION)))
|
||||
GO_MAJOR_VERSION := $(word 1,$(subst ., ,$(GO_VERSION)))
|
||||
GO_MINOR_VERSION := $(word 2,$(subst ., ,$(GO_VERSION)))
|
||||
|
||||
LDFLAGS = -ldflags "-s -w -X $(SLIVER_UPDATE_PKG).SliverPublicKey=$(SLIVER_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
LDFLAGS_DEBUG = -ldflags "-X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
COMMA := ,
|
||||
|
||||
ifeq ($(MAKECMDGOALS), linux)
|
||||
LDFLAGS = -ldflags "-s -w -extldflags '-static' -X $(CLIENT_ASSETS_PKG).DefaultArmoryPublicKey=$(ARMORY_PUBLIC_KEY) -X $(CLIENT_ASSETS_PKG).DefaultArmoryRepoURL=$(ARMORY_REPO_URL)"
|
||||
|
||||
+141
-44
@@ -1,17 +1,40 @@
|
||||
package ai
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2026 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
"github.com/bishopfox/sliver/client/termio"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/charmbracelet/colorprofile"
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/term"
|
||||
)
|
||||
|
||||
var openAIProgramTTY = tea.OpenTTY
|
||||
|
||||
// AICmd launches the AI conversation TUI.
|
||||
func AICmd(_ *cobra.Command, con *console.SliverClient, _ []string) {
|
||||
if con == nil || con.Rpc == nil {
|
||||
@@ -23,24 +46,72 @@ func AICmd(_ *cobra.Command, con *console.SliverClient, _ []string) {
|
||||
|
||||
listenerID, listener := con.CreateEventListener()
|
||||
defer con.RemoveEventListener(listenerID)
|
||||
restoreEventNotifications := con.SuppressEventNotifications()
|
||||
defer restoreEventNotifications()
|
||||
|
||||
model := newAIModel(con, buildAIContext(con), listener)
|
||||
model.showExperimentalWarningModal()
|
||||
|
||||
width, height := 100, 30
|
||||
if w, h, err := term.GetSize(0); err == nil && w > 0 && h > 0 {
|
||||
width, height = w, h
|
||||
}
|
||||
|
||||
program := tea.NewProgram(
|
||||
model,
|
||||
opts := []tea.ProgramOption{
|
||||
tea.WithWindowSize(width, height),
|
||||
tea.WithColorProfile(colorprofile.TrueColor),
|
||||
)
|
||||
}
|
||||
ttyOpts, cleanup := configureAIProgramTTY()
|
||||
if cleanup != nil {
|
||||
defer cleanup()
|
||||
}
|
||||
opts = append(opts, ttyOpts...)
|
||||
|
||||
program := tea.NewProgram(model, opts...)
|
||||
if _, err := program.Run(); err != nil {
|
||||
con.PrintErrorf("AI TUI error: %s\n", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Console logging can swap stdout/stderr to pipes so session output can be tee'd
|
||||
// to disk. Bubble Tea needs a real TTY output to enable terminal features like
|
||||
// mouse reporting, so bind the AI TUI back to the interactive terminal when
|
||||
// stdout is no longer the controlling TTY.
|
||||
func configureAIProgramTTY() ([]tea.ProgramOption, func()) {
|
||||
stdinTTY := isAITTY(os.Stdin)
|
||||
stdoutTTY := isAITTY(os.Stdout)
|
||||
|
||||
switch {
|
||||
case !stdinTTY || stdoutTTY:
|
||||
return nil, nil
|
||||
|
||||
case isAITTY(termio.InteractiveInput()) && isAITTY(termio.InteractiveOutput()):
|
||||
return []tea.ProgramOption{
|
||||
tea.WithInput(termio.InteractiveInput()),
|
||||
tea.WithOutput(termio.InteractiveOutput()),
|
||||
}, nil
|
||||
|
||||
default:
|
||||
inTTY, outTTY, err := openAIProgramTTY()
|
||||
if err != nil {
|
||||
return nil, nil
|
||||
}
|
||||
return []tea.ProgramOption{
|
||||
tea.WithInput(inTTY),
|
||||
tea.WithOutput(outTTY),
|
||||
}, func() {
|
||||
_ = inTTY.Close()
|
||||
if outTTY != inTTY {
|
||||
_ = outTTY.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isAITTY(file *os.File) bool {
|
||||
return file != nil && term.IsTerminal(int(file.Fd()))
|
||||
}
|
||||
|
||||
type aiContext struct {
|
||||
target aiTargetSummary
|
||||
connection aiConnectionSummary
|
||||
@@ -48,13 +119,15 @@ type aiContext struct {
|
||||
}
|
||||
|
||||
type aiTargetSummary struct {
|
||||
Label string
|
||||
Host string
|
||||
OS string
|
||||
Arch string
|
||||
C2 string
|
||||
Mode string
|
||||
Details []string
|
||||
SessionID string
|
||||
BeaconID string
|
||||
Label string
|
||||
Host string
|
||||
OS string
|
||||
Arch string
|
||||
C2 string
|
||||
Mode string
|
||||
Details []string
|
||||
}
|
||||
|
||||
type aiConnectionSummary struct {
|
||||
@@ -64,16 +137,64 @@ type aiConnectionSummary struct {
|
||||
State string
|
||||
}
|
||||
|
||||
func defaultAITargetSummary() aiTargetSummary {
|
||||
return aiTargetSummary{
|
||||
Label: "No active target",
|
||||
Host: "Select a session or beacon with `use` or ctrl+s",
|
||||
Mode: "offline preview",
|
||||
C2: "n/a",
|
||||
OS: "unknown",
|
||||
Arch: "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
func aiSessionTargetSummary(session *clientpb.Session) aiTargetSummary {
|
||||
if session == nil {
|
||||
return defaultAITargetSummary()
|
||||
}
|
||||
|
||||
return aiTargetSummary{
|
||||
SessionID: session.ID,
|
||||
Label: fmt.Sprintf("Session %s", fallback(session.Name, session.ID)),
|
||||
Host: fallback(session.Hostname, "<unknown host>"),
|
||||
OS: fallback(session.OS, "unknown"),
|
||||
Arch: fallback(session.Arch, "unknown"),
|
||||
C2: fallback(session.ActiveC2, "unknown"),
|
||||
Mode: "interactive session",
|
||||
Details: []string{
|
||||
fmt.Sprintf("User: %s", fallback(session.Username, "<unknown>")),
|
||||
fmt.Sprintf("PID: %d", session.PID),
|
||||
fmt.Sprintf("Remote: %s", fallback(session.RemoteAddress, "<unknown>")),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func aiBeaconTargetSummary(beacon *clientpb.Beacon) aiTargetSummary {
|
||||
if beacon == nil {
|
||||
return defaultAITargetSummary()
|
||||
}
|
||||
|
||||
return aiTargetSummary{
|
||||
BeaconID: beacon.ID,
|
||||
Label: fmt.Sprintf("Beacon %s", fallback(beacon.Name, beacon.ID)),
|
||||
Host: fallback(beacon.Hostname, "<unknown host>"),
|
||||
OS: fallback(beacon.OS, "unknown"),
|
||||
Arch: fallback(beacon.Arch, "unknown"),
|
||||
C2: fallback(beacon.ActiveC2, "unknown"),
|
||||
Mode: "asynchronous beacon",
|
||||
Details: []string{
|
||||
fmt.Sprintf("User: %s", fallback(beacon.Username, "<unknown>")),
|
||||
fmt.Sprintf("PID: %d", beacon.PID),
|
||||
fmt.Sprintf("Remote: %s", fallback(beacon.RemoteAddress, "<unknown>")),
|
||||
fmt.Sprintf("Interval: %s", time.Duration(beacon.Interval).String()),
|
||||
fmt.Sprintf("Next checkin: %s", formatUnix(beacon.NextCheckin)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func buildAIContext(con *console.SliverClient) aiContext {
|
||||
ctx := aiContext{
|
||||
target: aiTargetSummary{
|
||||
Label: "No active target",
|
||||
Host: "Select a session or beacon with `use`",
|
||||
Mode: "offline preview",
|
||||
C2: "n/a",
|
||||
OS: "unknown",
|
||||
Arch: "unknown",
|
||||
},
|
||||
target: defaultAITargetSummary(),
|
||||
connection: aiConnectionSummary{
|
||||
Profile: "<disconnected>",
|
||||
Server: "<unknown>",
|
||||
@@ -96,33 +217,9 @@ func buildAIContext(con *console.SliverClient) aiContext {
|
||||
session, beacon := con.ActiveTarget.Get()
|
||||
switch {
|
||||
case session != nil:
|
||||
ctx.target = aiTargetSummary{
|
||||
Label: fmt.Sprintf("Session %s", fallback(session.Name, session.ID)),
|
||||
Host: fallback(session.Hostname, "<unknown host>"),
|
||||
OS: fallback(session.OS, "unknown"),
|
||||
Arch: fallback(session.Arch, "unknown"),
|
||||
C2: fallback(session.ActiveC2, "unknown"),
|
||||
Mode: "interactive session",
|
||||
Details: []string{
|
||||
fmt.Sprintf("User: %s", fallback(session.Username, "<unknown>")),
|
||||
fmt.Sprintf("PID: %d", session.PID),
|
||||
fmt.Sprintf("Remote: %s", fallback(session.RemoteAddress, "<unknown>")),
|
||||
},
|
||||
}
|
||||
ctx.target = aiSessionTargetSummary(session)
|
||||
case beacon != nil:
|
||||
ctx.target = aiTargetSummary{
|
||||
Label: fmt.Sprintf("Beacon %s", fallback(beacon.Name, beacon.ID)),
|
||||
Host: fallback(beacon.Hostname, "<unknown host>"),
|
||||
OS: fallback(beacon.OS, "unknown"),
|
||||
Arch: fallback(beacon.Arch, "unknown"),
|
||||
C2: fallback(beacon.ActiveC2, "unknown"),
|
||||
Mode: "asynchronous beacon",
|
||||
Details: []string{
|
||||
fmt.Sprintf("User: %s", fallback(beacon.Username, "<unknown>")),
|
||||
fmt.Sprintf("Interval: %s", time.Duration(beacon.Interval).String()),
|
||||
fmt.Sprintf("Next checkin: %s", formatUnix(beacon.NextCheckin)),
|
||||
},
|
||||
}
|
||||
ctx.target = aiBeaconTargetSummary(beacon)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,5 +1,23 @@
|
||||
package ai
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2026 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"github.com/bishopfox/sliver/client/command/help"
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
|
||||
+136
-34
@@ -1,5 +1,23 @@
|
||||
package ai
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2026 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
@@ -49,7 +67,7 @@ func TestLoadAIStateCmdLoadsConversationOverRPC(t *testing.T) {
|
||||
rpcClient, cleanup := newAITestRPCClient(t, server)
|
||||
defer cleanup()
|
||||
|
||||
cmd := loadAIStateCmd(&console.SliverClient{Rpc: rpcClient}, "")
|
||||
cmd := loadAIStateCmd(&console.SliverClient{Rpc: rpcClient}, aiTargetSummary{}, "")
|
||||
msg := cmd()
|
||||
|
||||
loaded, ok := msg.(aiLoadedMsg)
|
||||
@@ -87,24 +105,27 @@ func TestLoadAIStateCmdCreatesConversationWhenNoneExist(t *testing.T) {
|
||||
providersResp: &clientpb.AIProviderConfigs{
|
||||
Providers: []*clientpb.AIProviderConfig{{Name: "openai", Configured: true}},
|
||||
Config: &clientpb.AIConfigSummary{
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
Valid: true,
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
SystemPrompt: "Stay concise.",
|
||||
Valid: true,
|
||||
},
|
||||
},
|
||||
conversationsResp: &clientpb.AIConversations{},
|
||||
saveConversationResp: &clientpb.AIConversation{
|
||||
ID: "conv-created",
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
Title: "New conversation",
|
||||
ID: "conv-created",
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
Title: "New conversation",
|
||||
SystemPrompt: "Stay concise.",
|
||||
},
|
||||
conversationByID: map[string]*clientpb.AIConversation{
|
||||
"conv-created": {
|
||||
ID: "conv-created",
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
Title: "New conversation",
|
||||
ID: "conv-created",
|
||||
Provider: "openai",
|
||||
Model: "gpt-test",
|
||||
Title: "New conversation",
|
||||
SystemPrompt: "Stay concise.",
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -112,7 +133,7 @@ func TestLoadAIStateCmdCreatesConversationWhenNoneExist(t *testing.T) {
|
||||
rpcClient, cleanup := newAITestRPCClient(t, server)
|
||||
defer cleanup()
|
||||
|
||||
msg := loadAIStateCmd(&console.SliverClient{Rpc: rpcClient}, "")()
|
||||
msg := loadAIStateCmd(&console.SliverClient{Rpc: rpcClient}, aiTargetSummary{}, "")()
|
||||
|
||||
loaded, ok := msg.(aiLoadedMsg)
|
||||
if !ok {
|
||||
@@ -134,6 +155,9 @@ func TestLoadAIStateCmdCreatesConversationWhenNoneExist(t *testing.T) {
|
||||
if request.GetProvider() != "openai" || request.GetModel() != "gpt-test" || request.GetTitle() != "New conversation" {
|
||||
t.Fatalf("unexpected SaveAIConversation request: %+v", request)
|
||||
}
|
||||
if request.GetSystemPrompt() != "Stay concise." {
|
||||
t.Fatalf("unexpected SaveAIConversation system prompt: %q", request.GetSystemPrompt())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmitPromptCmdCreatesConversationAndSavesUserMessage(t *testing.T) {
|
||||
@@ -157,9 +181,11 @@ func TestSubmitPromptCmdCreatesConversationAndSavesUserMessage(t *testing.T) {
|
||||
|
||||
msg := submitPromptCmd(
|
||||
&console.SliverClient{Rpc: rpcClient},
|
||||
aiTargetSummary{SessionID: "session-1"},
|
||||
nil,
|
||||
"openai",
|
||||
"gpt-test",
|
||||
"Stay concise.",
|
||||
"Explain the workflow.\n\nWith details.",
|
||||
)()
|
||||
|
||||
@@ -179,6 +205,12 @@ func TestSubmitPromptCmdCreatesConversationAndSavesUserMessage(t *testing.T) {
|
||||
if server.saveConversationReqs[0].GetTitle() != "Explain the workflow." {
|
||||
t.Fatalf("unexpected conversation title: %q", server.saveConversationReqs[0].GetTitle())
|
||||
}
|
||||
if server.saveConversationReqs[0].GetTargetSessionID() != "session-1" {
|
||||
t.Fatalf("expected target session to be forwarded, got %+v", server.saveConversationReqs[0])
|
||||
}
|
||||
if server.saveConversationReqs[0].GetSystemPrompt() != "Stay concise." {
|
||||
t.Fatalf("unexpected conversation system prompt: %q", server.saveConversationReqs[0].GetSystemPrompt())
|
||||
}
|
||||
if len(server.saveMessageReqs) != 1 {
|
||||
t.Fatalf("expected SaveAIConversationMessage to be called once, got %d", len(server.saveMessageReqs))
|
||||
}
|
||||
@@ -206,9 +238,11 @@ func TestSubmitPromptCmdUsesExistingConversationSettings(t *testing.T) {
|
||||
|
||||
msg := submitPromptCmd(
|
||||
&console.SliverClient{Rpc: rpcClient},
|
||||
aiTargetSummary{},
|
||||
&clientpb.AIConversation{ID: "conv-1", Provider: "anthropic", Model: "claude-test", Title: "Thread"},
|
||||
"openai",
|
||||
"gpt-test",
|
||||
"Stay concise.",
|
||||
"hello",
|
||||
)()
|
||||
|
||||
@@ -262,10 +296,16 @@ func TestDeleteConversationCmdSendsDeleteRequest(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestWaitForAIConversationEventCmdFiltersForAIEvents(t *testing.T) {
|
||||
conversation := &clientpb.AIConversation{ID: "conv-1", OperatorName: "alice"}
|
||||
data, err := proto.Marshal(conversation)
|
||||
eventPayload := &clientpb.AIConversationEvent{
|
||||
EventType: clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_CONVERSATION_UPDATED,
|
||||
Conversation: &clientpb.AIConversation{
|
||||
ID: "conv-1",
|
||||
OperatorName: "alice",
|
||||
},
|
||||
}
|
||||
data, err := proto.Marshal(eventPayload)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal conversation: %v", err)
|
||||
t.Fatalf("marshal conversation event: %v", err)
|
||||
}
|
||||
|
||||
listener := make(chan *clientpb.Event, 2)
|
||||
@@ -278,8 +318,27 @@ func TestWaitForAIConversationEventCmdFiltersForAIEvents(t *testing.T) {
|
||||
if !ok {
|
||||
t.Fatalf("expected aiConversationEventMsg, got %T", msg)
|
||||
}
|
||||
if aiEvent.conversation == nil || aiEvent.conversation.GetID() != "conv-1" {
|
||||
t.Fatalf("unexpected AI conversation event payload: %+v", aiEvent.conversation)
|
||||
if aiEvent.event == nil || aiEvent.event.GetConversation() == nil || aiEvent.event.GetConversation().GetID() != "conv-1" {
|
||||
t.Fatalf("unexpected AI conversation event payload: %+v", aiEvent.event)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitForAIConversationEventCmdReturnsToastMessages(t *testing.T) {
|
||||
listener := make(chan *clientpb.Event, 1)
|
||||
listener <- &clientpb.Event{
|
||||
EventType: consts.ClientToastEvent,
|
||||
Data: []byte("Session alpha connected"),
|
||||
Err: "info",
|
||||
}
|
||||
|
||||
msg := waitForAIConversationEventCmd(listener)()
|
||||
|
||||
toast, ok := msg.(aiToastMsg)
|
||||
if !ok {
|
||||
t.Fatalf("expected aiToastMsg, got %T", msg)
|
||||
}
|
||||
if toast.level != "info" || toast.message != "Session alpha connected" {
|
||||
t.Fatalf("unexpected toast payload: %+v", toast)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -293,7 +352,7 @@ func TestWaitForAIConversationEventCmdReturnsClosedMessage(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIModelReloadsSharedConversationEvents(t *testing.T) {
|
||||
func TestAIModelAppliesSharedConversationEventsWithoutReload(t *testing.T) {
|
||||
listener := make(chan *clientpb.Event)
|
||||
close(listener)
|
||||
|
||||
@@ -301,25 +360,30 @@ func TestAIModelReloadsSharedConversationEvents(t *testing.T) {
|
||||
connection: aiConnectionSummary{Operator: "alice"},
|
||||
}, listener)
|
||||
model.loading = false
|
||||
model.conversations = []*clientpb.AIConversation{{ID: "conv-1", Title: "First"}}
|
||||
|
||||
updated, cmd := model.Update(aiConversationEventMsg{conversation: &clientpb.AIConversation{
|
||||
ID: "conv-2",
|
||||
OperatorName: "bob",
|
||||
updated, cmd := model.Update(aiConversationEventMsg{event: &clientpb.AIConversationEvent{
|
||||
EventType: clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_CONVERSATION_UPDATED,
|
||||
Conversation: &clientpb.AIConversation{
|
||||
ID: "conv-2",
|
||||
OperatorName: "bob",
|
||||
Title: "Shared thread",
|
||||
},
|
||||
}})
|
||||
|
||||
nextModel := updated.(*aiModel)
|
||||
if !nextModel.loading {
|
||||
t.Fatal("expected shared conversation event to trigger a reload")
|
||||
if nextModel.loading {
|
||||
t.Fatal("did not expect shared conversation event to trigger a reload")
|
||||
}
|
||||
if nextModel.status != "Conversation updated on the server. Refreshing..." {
|
||||
t.Fatalf("unexpected reload status: %q", nextModel.status)
|
||||
if conversationIndexByID(nextModel.conversations, "conv-2") < 0 {
|
||||
t.Fatalf("expected shared conversation to be merged locally, got %+v", nextModel.conversations)
|
||||
}
|
||||
if cmd == nil {
|
||||
t.Fatal("expected shared conversation event to schedule a refresh")
|
||||
t.Fatal("expected shared conversation event to continue listening for events")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIModelReloadsDeleteEventsWhileAwaitingResponse(t *testing.T) {
|
||||
func TestAIModelAppliesDeleteEventsWhileAwaitingResponse(t *testing.T) {
|
||||
listener := make(chan *clientpb.Event)
|
||||
close(listener)
|
||||
|
||||
@@ -332,22 +396,36 @@ func TestAIModelReloadsDeleteEventsWhileAwaitingResponse(t *testing.T) {
|
||||
ID: "conv-2",
|
||||
OperatorName: "alice",
|
||||
UpdatedAt: 100,
|
||||
TurnState: clientpb.AIConversationTurnState_AI_TURN_STATE_IN_PROGRESS,
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "still waiting"},
|
||||
},
|
||||
}
|
||||
model.conversations = []*clientpb.AIConversation{
|
||||
{ID: "conv-2", Title: "Current"},
|
||||
{ID: "conv-1", Title: "Fallback"},
|
||||
}
|
||||
|
||||
updated, cmd := model.Update(aiConversationEventMsg{conversation: &clientpb.AIConversation{
|
||||
ID: "conv-2",
|
||||
OperatorName: "bob",
|
||||
updated, cmd := model.Update(aiConversationEventMsg{event: &clientpb.AIConversationEvent{
|
||||
EventType: clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_CONVERSATION_DELETED,
|
||||
ConversationID: "conv-2",
|
||||
Conversation: &clientpb.AIConversation{
|
||||
ID: "conv-2",
|
||||
},
|
||||
}})
|
||||
|
||||
nextModel := updated.(*aiModel)
|
||||
if !nextModel.loading {
|
||||
t.Fatal("expected delete tombstone to reload the conversation")
|
||||
if nextModel.loading {
|
||||
t.Fatal("did not expect delete tombstone to trigger a reload")
|
||||
}
|
||||
if nextModel.currentConversation != nil {
|
||||
t.Fatalf("expected deleted conversation to be cleared, got %+v", nextModel.currentConversation)
|
||||
}
|
||||
if conversationIndexByID(nextModel.conversations, "conv-2") >= 0 {
|
||||
t.Fatalf("expected deleted conversation to be removed, got %+v", nextModel.conversations)
|
||||
}
|
||||
if cmd == nil {
|
||||
t.Fatal("expected delete tombstone to schedule a refresh")
|
||||
t.Fatal("expected delete tombstone to continue listening for events")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -361,9 +439,13 @@ type aiRPCServer struct {
|
||||
conversationByID map[string]*clientpb.AIConversation
|
||||
saveConversationResp *clientpb.AIConversation
|
||||
saveMessageResp *clientpb.AIConversationMessage
|
||||
sessionsResp *clientpb.Sessions
|
||||
beaconsResp *clientpb.Beacons
|
||||
|
||||
getProvidersCalls int
|
||||
getConversationsCalls int
|
||||
getSessionsCalls int
|
||||
getBeaconsCalls int
|
||||
getConversationReqs []*clientpb.AIConversationReq
|
||||
deleteConversationReqs []*clientpb.AIConversationReq
|
||||
saveConversationReqs []*clientpb.AIConversation
|
||||
@@ -394,6 +476,26 @@ func (s *aiRPCServer) GetAIConversations(context.Context, *commonpb.Empty) (*cli
|
||||
return proto.Clone(s.conversationsResp).(*clientpb.AIConversations), nil
|
||||
}
|
||||
|
||||
func (s *aiRPCServer) GetSessions(context.Context, *commonpb.Empty) (*clientpb.Sessions, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.getSessionsCalls++
|
||||
if s.sessionsResp == nil {
|
||||
return &clientpb.Sessions{}, nil
|
||||
}
|
||||
return proto.Clone(s.sessionsResp).(*clientpb.Sessions), nil
|
||||
}
|
||||
|
||||
func (s *aiRPCServer) GetBeacons(context.Context, *commonpb.Empty) (*clientpb.Beacons, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.getBeaconsCalls++
|
||||
if s.beaconsResp == nil {
|
||||
return &clientpb.Beacons{}, nil
|
||||
}
|
||||
return proto.Clone(s.beaconsResp).(*clientpb.Beacons), nil
|
||||
}
|
||||
|
||||
func (s *aiRPCServer) GetAIConversation(_ context.Context, req *clientpb.AIConversationReq) (*clientpb.AIConversation, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
+3036
-255
File diff suppressed because it is too large.
Load diff
+1569
-14
File diff suppressed because it is too large.
Load diff
@@ -400,7 +400,13 @@ func runAliasCommand(cmd *cobra.Command, con *console.SliverClient, args []strin
|
||||
|
||||
var outFilePath *os.File
|
||||
if save, _ := cmd.Flags().GetBool("save"); save {
|
||||
outFile := filepath.Base(fmt.Sprintf("%s_%s*.log", filepath.Base(cmd.Name()), filepath.Base(session.GetHostname())))
|
||||
hostname := ""
|
||||
if session != nil {
|
||||
hostname = session.GetHostname()
|
||||
} else if beacon != nil {
|
||||
hostname = beacon.GetHostname()
|
||||
}
|
||||
outFile := filepath.Base(fmt.Sprintf("%s_%s*.log", filepath.Base(cmd.Name()), filepath.Base(hostname)))
|
||||
outFilePath, err = os.CreateTemp("", outFile)
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
|
||||
@@ -1,5 +1,23 @@
|
||||
package armory
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2026 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"github.com/rsteube/carapace"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
@@ -1,5 +1,23 @@
|
||||
package armory
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2026 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
@@ -66,6 +66,7 @@ func Commands(con *console.SliverClient) []*cobra.Command {
|
||||
flags.Bind("", false, consoleCmd, func(f *pflag.FlagSet) {
|
||||
f.String("transport", string(defaultConfig.Transport), "mcp transport (http or sse)")
|
||||
f.String("url", "", "mcp server URL (defaults to local server settings)")
|
||||
f.String("auth-token", "", "MCP auth token (defaults to the local mcp.yaml token)")
|
||||
})
|
||||
flags.BindFlagCompletions(consoleCmd, func(comp *carapace.ActionMap) {
|
||||
(*comp)["transport"] = carapace.ActionValues("http", "sse")
|
||||
|
||||
@@ -6,12 +6,14 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/bishopfox/sliver/client/version"
|
||||
mcpclient "github.com/mark3labs/mcp-go/client"
|
||||
mcptransport "github.com/mark3labs/mcp-go/client/transport"
|
||||
mcpapi "github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
@@ -36,7 +38,23 @@ func McpConsoleCmd(cmd *cobra.Command, con *console.SliverClient, args []string)
|
||||
}
|
||||
}
|
||||
|
||||
client, err := newConsoleClient(transport, targetURL)
|
||||
authToken, _ := cmd.Flags().GetString("auth-token")
|
||||
if authToken == "" {
|
||||
status := clientmcp.GetStatus()
|
||||
if status.AuthToken != "" {
|
||||
authToken = status.AuthToken
|
||||
} else {
|
||||
authInfo, loadErr := clientmcp.LoadAuthInfo()
|
||||
if loadErr == nil {
|
||||
authToken = authInfo.Token
|
||||
} else if !errors.Is(loadErr, os.ErrNotExist) {
|
||||
con.PrintErrorf("%s\n", loadErr)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
client, err := newConsoleClient(transport, targetURL, authToken)
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
@@ -68,12 +86,19 @@ func McpConsoleCmd(cmd *cobra.Command, con *console.SliverClient, args []string)
|
||||
runConsoleREPL(ctx, client, cmd.InOrStdin(), cmd.OutOrStdout())
|
||||
}
|
||||
|
||||
func newConsoleClient(transport clientmcp.Transport, targetURL string) (*mcpclient.Client, error) {
|
||||
func newConsoleClient(transport clientmcp.Transport, targetURL string, authToken string) (*mcpclient.Client, error) {
|
||||
if strings.TrimSpace(authToken) == "" {
|
||||
return nil, fmt.Errorf("mcp auth token is required (use --auth-token or start the local MCP server first)")
|
||||
}
|
||||
headers := map[string]string{
|
||||
clientmcp.AuthHeaderName(): authToken,
|
||||
}
|
||||
|
||||
switch transport {
|
||||
case clientmcp.TransportHTTP:
|
||||
return mcpclient.NewStreamableHttpClient(targetURL)
|
||||
return mcpclient.NewStreamableHttpClient(targetURL, mcptransport.WithHTTPHeaders(headers))
|
||||
case clientmcp.TransportSSE:
|
||||
return mcpclient.NewSSEMCPClient(targetURL)
|
||||
return mcpclient.NewSSEMCPClient(targetURL, mcpclient.WithHeaders(headers))
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported transport %q", transport)
|
||||
}
|
||||
|
||||
@@ -26,6 +26,12 @@ func McpCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
} else {
|
||||
con.PrintErrorf("Endpoint: %s\n", err)
|
||||
}
|
||||
if status.AuthHeader != "" {
|
||||
con.Printf("Auth Header: %s\n", status.AuthHeader)
|
||||
}
|
||||
if status.AuthToken != "" {
|
||||
con.Printf("Auth Token: %s\n", status.AuthToken)
|
||||
}
|
||||
|
||||
if status.Running && !status.StartedAt.IsZero() {
|
||||
uptime := time.Since(status.StartedAt).Truncate(time.Second)
|
||||
|
||||
@@ -43,4 +43,11 @@ func McpStartCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
if err == nil {
|
||||
con.PrintInfof("Endpoint: %s\n", endpoint)
|
||||
}
|
||||
status := slivermcp.GetStatus()
|
||||
if status.AuthHeader != "" {
|
||||
con.PrintInfof("Auth Header: %s\n", status.AuthHeader)
|
||||
}
|
||||
if status.AuthToken != "" {
|
||||
con.PrintInfof("Auth Token: %s\n", status.AuthToken)
|
||||
}
|
||||
}
|
||||
+98
-12
@@ -29,6 +29,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/client/assets"
|
||||
@@ -99,6 +100,7 @@ type SliverClient struct {
|
||||
connWg *sync.WaitGroup
|
||||
connectionHooksApplied bool
|
||||
logCommandHookApplied bool
|
||||
suppressedEventNotifs atomic.Int32
|
||||
|
||||
// These writers are always safe to use (never error); the underlying stream
|
||||
// can be swapped on server switch to avoid breaking io.MultiWriter/io.Copy.
|
||||
@@ -352,8 +354,20 @@ func (con *SliverClient) startEventLoop(ctx context.Context) {
|
||||
switch event.EventType {
|
||||
|
||||
case consts.CanaryEvent:
|
||||
con.PrintEventErrorf("%s %s has been burned (DNS Canary)", StyleBold.Render("WARNING:"), event.Session.Name)
|
||||
sessions := con.GetSessionsByName(event.Session.Name)
|
||||
if con.eventNotificationsSuppressed() {
|
||||
affected := make([]string, 0, len(sessions))
|
||||
for _, session := range sessions {
|
||||
affected = append(affected, strings.Split(session.ID, "-")[0])
|
||||
}
|
||||
message := fmt.Sprintf("%s %s has been burned (DNS Canary)", StyleBold.Render("WARNING:"), event.Session.Name)
|
||||
if len(affected) > 0 {
|
||||
message += " Affected sessions: " + strings.Join(affected, ", ")
|
||||
}
|
||||
con.emitToastEvent("error", message)
|
||||
break
|
||||
}
|
||||
con.PrintEventErrorf("%s %s has been burned (DNS Canary)", StyleBold.Render("WARNING:"), event.Session.Name)
|
||||
for _, session := range sessions {
|
||||
shortID := strings.Split(session.ID, "-")[0]
|
||||
con.PrintErrorf("\t🔥 Session %s is affected", shortID)
|
||||
@@ -361,8 +375,20 @@ func (con *SliverClient) startEventLoop(ctx context.Context) {
|
||||
|
||||
case consts.WatchtowerEvent:
|
||||
msg := string(event.Data)
|
||||
con.PrintEventErrorf("%s %s has been burned (seen on %s)", StyleBold.Render("WARNING:"), event.Session.Name, msg)
|
||||
sessions := con.GetSessionsByName(event.Session.Name)
|
||||
if con.eventNotificationsSuppressed() {
|
||||
affected := make([]string, 0, len(sessions))
|
||||
for _, session := range sessions {
|
||||
affected = append(affected, strings.Split(session.ID, "-")[0])
|
||||
}
|
||||
message := fmt.Sprintf("%s %s has been burned (seen on %s)", StyleBold.Render("WARNING:"), event.Session.Name, msg)
|
||||
if len(affected) > 0 {
|
||||
message += " Affected sessions: " + strings.Join(affected, ", ")
|
||||
}
|
||||
con.emitToastEvent("error", message)
|
||||
break
|
||||
}
|
||||
con.PrintEventErrorf("%s %s has been burned (seen on %s)", StyleBold.Render("WARNING:"), event.Session.Name, msg)
|
||||
for _, session := range sessions {
|
||||
shortID := strings.Split(session.ID, "-")[0]
|
||||
con.PrintErrorf("\t🔥 Session %s is affected", shortID)
|
||||
@@ -370,42 +396,42 @@ func (con *SliverClient) startEventLoop(ctx context.Context) {
|
||||
|
||||
case consts.JoinedEvent:
|
||||
if con.Settings.UserConnect {
|
||||
con.PrintInfof("%s has joined the game", event.Client.Operator.Name)
|
||||
con.emitConsoleNotification("info", false, "%s has joined the game", event.Client.Operator.Name)
|
||||
}
|
||||
case consts.LeftEvent:
|
||||
if con.Settings.UserConnect {
|
||||
con.PrintInfof("%s left the game", event.Client.Operator.Name)
|
||||
con.emitConsoleNotification("info", false, "%s left the game", event.Client.Operator.Name)
|
||||
}
|
||||
|
||||
case consts.JobStoppedEvent:
|
||||
job := event.Job
|
||||
con.PrintErrorf("Job #%d stopped (%s/%s)", job.ID, job.Protocol, job.Name)
|
||||
con.emitConsoleNotification("error", false, "Job #%d stopped (%s/%s)", job.ID, job.Protocol, job.Name)
|
||||
|
||||
case consts.SessionOpenedEvent:
|
||||
session := event.Session
|
||||
currentTime := time.Now().Format(time.RFC1123)
|
||||
shortID := strings.Split(session.ID, "-")[0]
|
||||
con.PrintEventInfof("Session %s %s - %s (%s) - %s/%s - %v",
|
||||
con.emitConsoleNotification("info", true, "Session %s %s - %s (%s) - %s/%s - %v",
|
||||
shortID, session.Name, session.RemoteAddress, session.Hostname, session.OS, session.Arch, currentTime)
|
||||
|
||||
case consts.SessionUpdateEvent:
|
||||
session := event.Session
|
||||
currentTime := time.Now().Format(time.RFC1123)
|
||||
shortID := strings.Split(session.ID, "-")[0]
|
||||
con.PrintInfof("Session %s has been updated - %v", shortID, currentTime)
|
||||
con.emitConsoleNotification("info", false, "Session %s has been updated - %v", shortID, currentTime)
|
||||
|
||||
case consts.SessionClosedEvent:
|
||||
session := event.Session
|
||||
currentTime := time.Now().Format(time.RFC1123)
|
||||
shortID := strings.Split(session.ID, "-")[0]
|
||||
con.PrintEventErrorf("Lost session %s %s - %s (%s) - %s/%s - %v",
|
||||
con.emitConsoleNotification("error", true, "Lost session %s %s - %s (%s) - %s/%s - %v",
|
||||
shortID, session.Name, session.RemoteAddress, session.Hostname, session.OS, session.Arch, currentTime)
|
||||
activeSession := con.ActiveTarget.GetSession()
|
||||
core.GetTunnels().CloseForSession(session.ID)
|
||||
core.CloseCursedProcesses(session.ID)
|
||||
if activeSession != nil && activeSession.ID == session.ID {
|
||||
con.ActiveTarget.Set(nil, nil)
|
||||
con.PrintErrorf("Active session disconnected")
|
||||
con.emitConsoleNotification("error", false, "Active session disconnected")
|
||||
}
|
||||
|
||||
case consts.BeaconRegisteredEvent:
|
||||
@@ -413,7 +439,7 @@ func (con *SliverClient) startEventLoop(ctx context.Context) {
|
||||
proto.Unmarshal(event.Data, beacon)
|
||||
currentTime := time.Now().Format(time.RFC1123)
|
||||
shortID := strings.Split(beacon.ID, "-")[0]
|
||||
con.PrintEventInfof("Beacon %s %s - %s (%s) - %s/%s - %v",
|
||||
con.emitConsoleNotification("info", true, "Beacon %s %s - %s (%s) - %s/%s - %v",
|
||||
shortID, beacon.Name, beacon.RemoteAddress, beacon.Hostname, beacon.OS, beacon.Arch, currentTime)
|
||||
|
||||
case consts.BeaconTaskResultEvent:
|
||||
@@ -433,6 +459,66 @@ func (con *SliverClient) CreateEventListener() (string, <-chan *clientpb.Event)
|
||||
return listenerID.String(), listener
|
||||
}
|
||||
|
||||
// SuppressEventNotifications redirects console event messages to local listeners instead of stdout.
|
||||
func (con *SliverClient) SuppressEventNotifications() func() {
|
||||
if con == nil {
|
||||
return func() {}
|
||||
}
|
||||
con.suppressedEventNotifs.Add(1)
|
||||
return func() {
|
||||
if con == nil {
|
||||
return
|
||||
}
|
||||
if con.suppressedEventNotifs.Add(-1) < 0 {
|
||||
con.suppressedEventNotifs.Store(0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (con *SliverClient) eventNotificationsSuppressed() bool {
|
||||
return con != nil && con.suppressedEventNotifs.Load() > 0
|
||||
}
|
||||
|
||||
func (con *SliverClient) emitToastEvent(level string, message string) {
|
||||
if con == nil || con.EventListeners == nil || strings.TrimSpace(message) == "" {
|
||||
return
|
||||
}
|
||||
con.triggerEventListeners(&clientpb.Event{
|
||||
EventType: consts.ClientToastEvent,
|
||||
Data: []byte(message),
|
||||
Err: strings.TrimSpace(level),
|
||||
})
|
||||
}
|
||||
|
||||
func (con *SliverClient) emitConsoleNotification(level string, emphasized bool, format string, args ...any) {
|
||||
message := fmt.Sprintf(format, args...)
|
||||
if con.eventNotificationsSuppressed() {
|
||||
con.emitToastEvent(level, message)
|
||||
return
|
||||
}
|
||||
|
||||
switch strings.ToLower(strings.TrimSpace(level)) {
|
||||
case "error", "warn", "warning":
|
||||
if emphasized {
|
||||
con.PrintEventErrorf("%s", message)
|
||||
} else {
|
||||
con.PrintErrorf("%s", message)
|
||||
}
|
||||
case "success":
|
||||
if emphasized {
|
||||
con.PrintEventSuccessf("%s", message)
|
||||
} else {
|
||||
con.PrintSuccessf("%s", message)
|
||||
}
|
||||
default:
|
||||
if emphasized {
|
||||
con.PrintEventInfof("%s", message)
|
||||
} else {
|
||||
con.PrintInfof("%s", message)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// RemoveEventListener - removes an event listener given its id.
|
||||
func (con *SliverClient) RemoveEventListener(listenerID string) {
|
||||
value, ok := con.EventListeners.LoadAndDelete(listenerID)
|
||||
@@ -477,10 +563,10 @@ func (con *SliverClient) triggerReactions(event *clientpb.Event) {
|
||||
|
||||
for _, reaction := range reactions {
|
||||
for _, line := range reaction.Commands {
|
||||
con.PrintInfof("%s '%s'", StyleBold.Render("Execute reaction:"), line)
|
||||
con.emitConsoleNotification("info", false, "%s '%s'", StyleBold.Render("Execute reaction:"), line)
|
||||
err := con.App.ActiveMenu().RunCommandLine(context.Background(), line)
|
||||
if err != nil {
|
||||
con.PrintErrorf("Reaction command error: %s\n", err)
|
||||
con.emitConsoleNotification("error", false, "Reaction command error: %s", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,8 +3,11 @@ package console
|
||||
import (
|
||||
"bytes"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
consts "github.com/bishopfox/sliver/client/constants"
|
||||
)
|
||||
|
||||
func TestDrainStdoutPipeBufferStripsSyncFrames(t *testing.T) {
|
||||
@@ -102,3 +105,33 @@ func TestSyncOutputWaitsForPipeDrain(t *testing.T) {
|
||||
t.Fatalf("unexpected final output %q", got.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSuppressEventNotificationsRedirectsMessagesToListeners(t *testing.T) {
|
||||
con := &SliverClient{
|
||||
EventListeners: &sync.Map{},
|
||||
printf: func(string, ...any) (int, error) { return 0, nil },
|
||||
}
|
||||
|
||||
listenerID, listener := con.CreateEventListener()
|
||||
defer con.RemoveEventListener(listenerID)
|
||||
|
||||
restore := con.SuppressEventNotifications()
|
||||
defer restore()
|
||||
|
||||
con.emitConsoleNotification("info", true, "Session %s connected", "alpha")
|
||||
|
||||
select {
|
||||
case event := <-listener:
|
||||
if event.GetEventType() != consts.ClientToastEvent {
|
||||
t.Fatalf("expected toast event, got %q", event.GetEventType())
|
||||
}
|
||||
if got := string(event.GetData()); got != "Session alpha connected" {
|
||||
t.Fatalf("unexpected toast payload %q", got)
|
||||
}
|
||||
if got := event.GetErr(); got != "info" {
|
||||
t.Fatalf("unexpected toast level %q", got)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for redirected event notification")
|
||||
}
|
||||
}
|
||||
@@ -119,6 +119,9 @@ const (
|
||||
|
||||
// AIConversationEvent - Fires whenever an AI conversation or message changes.
|
||||
AIConversationEvent = "ai-conversation"
|
||||
|
||||
// ClientToastEvent - Local-only event used to surface console notifications inside TUIs.
|
||||
ClientToastEvent = "client-toast"
|
||||
)
|
||||
|
||||
// Commands.
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/assets"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const (
|
||||
authConfigFileName = "mcp.yaml"
|
||||
authHeaderName = "Authorization"
|
||||
minAuthTokenLength = 8
|
||||
authTokenBytes = 16
|
||||
sseMessagePath = "/message"
|
||||
)
|
||||
|
||||
// AuthInfo describes the persisted MCP HTTP/SSE authentication settings.
|
||||
type AuthInfo struct {
|
||||
Header string
|
||||
Token string
|
||||
ConfigPath string
|
||||
}
|
||||
|
||||
// AuthHeaderName returns the required header name for MCP HTTP/SSE requests.
|
||||
func AuthHeaderName() string {
|
||||
return authHeaderName
|
||||
}
|
||||
|
||||
// AuthConfigPath returns the MCP auth config path in the client app directory.
|
||||
func AuthConfigPath() string {
|
||||
rootDir, _ := filepath.Abs(assets.GetRootAppDir())
|
||||
return filepath.Join(rootDir, authConfigFileName)
|
||||
}
|
||||
|
||||
type authConfig struct {
|
||||
Token string `yaml:"token"`
|
||||
}
|
||||
|
||||
// LoadAuthInfo loads the persisted MCP auth configuration without creating it.
|
||||
func LoadAuthInfo() (AuthInfo, error) {
|
||||
cfg, exists, err := loadAuthConfig()
|
||||
if err != nil {
|
||||
return AuthInfo{}, err
|
||||
}
|
||||
if !exists {
|
||||
return AuthInfo{}, os.ErrNotExist
|
||||
}
|
||||
cfg.Token = strings.TrimSpace(cfg.Token)
|
||||
if err := validateAuthToken(cfg.Token); err != nil {
|
||||
return AuthInfo{}, fmt.Errorf("invalid mcp token in %s: %w", AuthConfigPath(), err)
|
||||
}
|
||||
return AuthInfo{
|
||||
Header: authHeaderName,
|
||||
Token: cfg.Token,
|
||||
ConfigPath: AuthConfigPath(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ResolveAuthInfo loads the persisted auth token or creates and saves one on first use.
|
||||
func ResolveAuthInfo() (AuthInfo, bool, error) {
|
||||
cfg, exists, err := loadAuthConfig()
|
||||
if err != nil {
|
||||
return AuthInfo{}, false, err
|
||||
}
|
||||
if exists {
|
||||
cfg.Token = strings.TrimSpace(cfg.Token)
|
||||
if err := validateAuthToken(cfg.Token); err != nil {
|
||||
return AuthInfo{}, false, fmt.Errorf("invalid mcp token in %s: %w", AuthConfigPath(), err)
|
||||
}
|
||||
return AuthInfo{
|
||||
Header: authHeaderName,
|
||||
Token: cfg.Token,
|
||||
ConfigPath: AuthConfigPath(),
|
||||
}, false, nil
|
||||
}
|
||||
|
||||
token, err := generateAuthToken()
|
||||
if err != nil {
|
||||
return AuthInfo{}, false, err
|
||||
}
|
||||
cfg = &authConfig{Token: token}
|
||||
if err := saveAuthConfig(cfg); err != nil {
|
||||
return AuthInfo{}, false, err
|
||||
}
|
||||
return AuthInfo{
|
||||
Header: authHeaderName,
|
||||
Token: token,
|
||||
ConfigPath: AuthConfigPath(),
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
func loadAuthConfig() (*authConfig, bool, error) {
|
||||
data, err := os.ReadFile(AuthConfigPath())
|
||||
if os.IsNotExist(err) {
|
||||
return nil, false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
|
||||
cfg := &authConfig{}
|
||||
if err := yaml.Unmarshal(data, cfg); err != nil {
|
||||
return nil, true, err
|
||||
}
|
||||
return cfg, true, nil
|
||||
}
|
||||
|
||||
func saveAuthConfig(cfg *authConfig) error {
|
||||
data, err := yaml.Marshal(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(AuthConfigPath(), data, 0o600)
|
||||
}
|
||||
|
||||
func generateAuthToken() (string, error) {
|
||||
tokenBytes := make([]byte, authTokenBytes)
|
||||
if _, err := rand.Read(tokenBytes); err != nil {
|
||||
return "", fmt.Errorf("generate mcp auth token: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(tokenBytes), nil
|
||||
}
|
||||
|
||||
func validateAuthToken(token string) error {
|
||||
if len(strings.TrimSpace(token)) < minAuthTokenLength {
|
||||
return fmt.Errorf("token must be at least %d characters", minAuthTokenLength)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func authMiddleware(token string, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !authTokenMatches(r.Header.Get(authHeaderName), token) {
|
||||
http.Error(w, fmt.Sprintf("missing or invalid %s header", authHeaderName), http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
func authTokenMatches(actual string, expected string) bool {
|
||||
if len(expected) == 0 || len(actual) != len(expected) {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare([]byte(actual), []byte(expected)) == 1
|
||||
}
|
||||
|
||||
type authenticatedHTTPTransport struct {
|
||||
server *http.Server
|
||||
}
|
||||
|
||||
func newAuthenticatedTransport(cfg Config, handler http.Handler, token string) serverTransport {
|
||||
mux := http.NewServeMux()
|
||||
authenticatedHandler := authMiddleware(token, handler)
|
||||
|
||||
switch cfg.Transport {
|
||||
case TransportHTTP:
|
||||
mux.Handle("/mcp", authenticatedHandler)
|
||||
case TransportSSE:
|
||||
mux.Handle("/sse", authenticatedHandler)
|
||||
mux.Handle(sseMessagePath, authenticatedHandler)
|
||||
}
|
||||
|
||||
return &authenticatedHTTPTransport{
|
||||
server: &http.Server{
|
||||
Handler: mux,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *authenticatedHTTPTransport) Start(addr string) error {
|
||||
if t.server.Addr == "" {
|
||||
t.server.Addr = addr
|
||||
} else if t.server.Addr != addr {
|
||||
return fmt.Errorf("conflicting listen address: %q vs %q", t.server.Addr, addr)
|
||||
}
|
||||
return t.server.ListenAndServe()
|
||||
}
|
||||
|
||||
func (t *authenticatedHTTPTransport) Shutdown(ctx context.Context) error {
|
||||
return t.server.Shutdown(ctx)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestResolveAuthInfoCreatesAndPersistsToken(t *testing.T) {
|
||||
t.Setenv("SLIVER_CLIENT_ROOT_DIR", t.TempDir())
|
||||
|
||||
authInfo, created, err := ResolveAuthInfo()
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveAuthInfo returned error: %v", err)
|
||||
}
|
||||
if !created {
|
||||
t.Fatalf("expected first ResolveAuthInfo call to create a token")
|
||||
}
|
||||
if authInfo.Header != AuthHeaderName() {
|
||||
t.Fatalf("expected header %q, got %q", AuthHeaderName(), authInfo.Header)
|
||||
}
|
||||
if len(authInfo.Token) != authTokenBytes*2 {
|
||||
t.Fatalf("expected %d hex characters, got %d", authTokenBytes*2, len(authInfo.Token))
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(AuthConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read persisted auth config: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), "token: "+authInfo.Token) {
|
||||
t.Fatalf("persisted config did not contain generated token: %s", string(data))
|
||||
}
|
||||
|
||||
authInfoAgain, createdAgain, err := ResolveAuthInfo()
|
||||
if err != nil {
|
||||
t.Fatalf("second ResolveAuthInfo returned error: %v", err)
|
||||
}
|
||||
if createdAgain {
|
||||
t.Fatalf("expected second ResolveAuthInfo call to reuse the existing token")
|
||||
}
|
||||
if authInfoAgain.Token != authInfo.Token {
|
||||
t.Fatalf("expected persisted token %q, got %q", authInfo.Token, authInfoAgain.Token)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthInfoRejectsInvalidExistingToken(t *testing.T) {
|
||||
t.Setenv("SLIVER_CLIENT_ROOT_DIR", t.TempDir())
|
||||
|
||||
if err := os.WriteFile(AuthConfigPath(), []byte("token: short\n"), 0o600); err != nil {
|
||||
t.Fatalf("failed to seed auth config: %v", err)
|
||||
}
|
||||
|
||||
_, _, err := ResolveAuthInfo()
|
||||
if err == nil {
|
||||
t.Fatal("expected invalid token error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "at least 8 characters") {
|
||||
t.Fatalf("expected length validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthMiddlewareRequiresExactToken(t *testing.T) {
|
||||
const token = "0123456789abcdef"
|
||||
|
||||
handler := authMiddleware(token, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
header string
|
||||
statusCode int
|
||||
}{
|
||||
{name: "missing", statusCode: http.StatusUnauthorized},
|
||||
{name: "wrong", header: "wrong-token", statusCode: http.StatusUnauthorized},
|
||||
{name: "prefixed", header: "Bearer " + token, statusCode: http.StatusUnauthorized},
|
||||
{name: "exact", header: token, statusCode: http.StatusNoContent},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodGet, "/mcp", nil)
|
||||
if testCase.header != "" {
|
||||
req.Header.Set(AuthHeaderName(), testCase.header)
|
||||
}
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
handler.ServeHTTP(recorder, req)
|
||||
|
||||
if recorder.Code != testCase.statusCode {
|
||||
t.Fatalf("expected status %d, got %d", testCase.statusCode, recorder.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+24
-11
@@ -20,10 +20,13 @@ var (
|
||||
|
||||
// Status describes the running state of the MCP server.
|
||||
type Status struct {
|
||||
Running bool
|
||||
StartedAt time.Time
|
||||
Config Config
|
||||
LastError string
|
||||
Running bool
|
||||
StartedAt time.Time
|
||||
Config Config
|
||||
LastError string
|
||||
AuthHeader string
|
||||
AuthToken string
|
||||
AuthConfigPath string
|
||||
}
|
||||
|
||||
type serverTransport interface {
|
||||
@@ -42,6 +45,7 @@ type mcpManager struct {
|
||||
done chan struct{}
|
||||
logger *log.Logger
|
||||
logFile *os.File
|
||||
authInfo AuthInfo
|
||||
}
|
||||
|
||||
func newManager() *mcpManager {
|
||||
@@ -93,10 +97,13 @@ func (m *mcpManager) status() Status {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return Status{
|
||||
Running: m.running,
|
||||
StartedAt: m.startedAt,
|
||||
Config: m.cfg,
|
||||
LastError: m.lastErr,
|
||||
Running: m.running,
|
||||
StartedAt: m.startedAt,
|
||||
Config: m.cfg,
|
||||
LastError: m.lastErr,
|
||||
AuthHeader: m.authInfo.Header,
|
||||
AuthToken: m.authInfo.Token,
|
||||
AuthConfigPath: m.authInfo.ConfigPath,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,6 +112,10 @@ func (m *mcpManager) start(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
authInfo, _, err := ResolveAuthInfo()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
if m.running {
|
||||
@@ -120,9 +131,9 @@ func (m *mcpManager) start(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
var transport serverTransport
|
||||
switch cfg.Transport {
|
||||
case TransportHTTP:
|
||||
transport = mcpserver.NewStreamableHTTPServer(mcpServer.server)
|
||||
transport = newAuthenticatedTransport(cfg, mcpserver.NewStreamableHTTPServer(mcpServer.server), authInfo.Token)
|
||||
case TransportSSE:
|
||||
transport = mcpserver.NewSSEServer(mcpServer.server)
|
||||
transport = newAuthenticatedTransport(cfg, mcpserver.NewSSEServer(mcpServer.server), authInfo.Token)
|
||||
default:
|
||||
m.mu.Unlock()
|
||||
logFile.Close()
|
||||
@@ -138,11 +149,12 @@ func (m *mcpManager) start(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
m.done = make(chan struct{})
|
||||
m.logger = logger
|
||||
m.logFile = logFile
|
||||
m.authInfo = authInfo
|
||||
addr := cfg.ListenAddress
|
||||
done := m.done
|
||||
m.mu.Unlock()
|
||||
|
||||
logger.Printf("starting mcp server transport=%s listen=%s", cfg.Transport, cfg.ListenAddress)
|
||||
logger.Printf("starting mcp server transport=%s listen=%s auth_header=%s auth_config=%s", cfg.Transport, cfg.ListenAddress, authInfo.Header, authInfo.ConfigPath)
|
||||
go m.run(addr, transport, done)
|
||||
return nil
|
||||
}
|
||||
@@ -161,6 +173,7 @@ func (m *mcpManager) run(addr string, transport serverTransport, done chan struc
|
||||
m.running = false
|
||||
m.transport = nil
|
||||
m.server = nil
|
||||
m.authInfo = AuthInfo{}
|
||||
logFile := m.logFile
|
||||
m.logFile = nil
|
||||
m.logger = nil
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
# Sliver AI Agentic Loop
|
||||
|
||||
## Goal
|
||||
|
||||
Port the example agentic loop from `examples/codex/codex-rs` into Sliver's AI chat so the server can:
|
||||
|
||||
- run a multi-step Responses API turn
|
||||
- emit turn and item lifecycle events similar to Codex
|
||||
- persist reasoning and tool-call items for sync and restoration
|
||||
- keep those non-chat items out of the model context window
|
||||
- render the new items in the client TUI
|
||||
|
||||
## How The Codex Example Loop Works
|
||||
|
||||
The Codex example splits the loop into three layers.
|
||||
|
||||
### 1. Provider/Protocol Layer
|
||||
|
||||
The protocol layer receives raw model events such as:
|
||||
|
||||
- turn started / completed / failed
|
||||
- reasoning deltas and completed reasoning items
|
||||
- command/tool begin and end notifications
|
||||
- assistant messages
|
||||
|
||||
These provider events are fine-grained and often arrive as partial updates.
|
||||
|
||||
### 2. Event Processor Layer
|
||||
|
||||
`event_processor_with_jsonl_output.rs` converts the raw provider stream into a stable thread event model.
|
||||
|
||||
The important behaviors are:
|
||||
|
||||
- a turn gets a lifecycle: `turn.started`, then `turn.completed` or `turn.failed`
|
||||
- each logical item gets a stable id
|
||||
- begin/end pairs are folded into a single item lifecycle
|
||||
- reasoning is emitted as its own item type
|
||||
- tool calls and command executions are emitted as their own item types
|
||||
- active items can be updated in place before completion
|
||||
|
||||
This is the key abstraction that makes the TUI and persisted history simple.
|
||||
|
||||
### 3. TUI Layer
|
||||
|
||||
The Codex TUI consumes those stable events rather than reloading the whole thread.
|
||||
|
||||
That lets it:
|
||||
|
||||
- append new items incrementally
|
||||
- update an existing in-flight item in place
|
||||
- keep the transcript stable while a turn runs
|
||||
- restore the same item types from stored history later
|
||||
|
||||
## Sliver Port
|
||||
|
||||
Sliver ports the same shape, adapted to its existing RPC and database model.
|
||||
|
||||
### Conversation State
|
||||
|
||||
`AIConversation` now stores:
|
||||
|
||||
- `ActiveTurnID`
|
||||
- `TurnState`
|
||||
- `TargetSessionID`
|
||||
- `TargetBeaconID`
|
||||
|
||||
The turn fields let clients restore "waiting", "completed", and "failed" state without inferring it from the last chat message.
|
||||
|
||||
The target fields bind a conversation to the active session or beacon so server-side tools can default to the conversation target.
|
||||
|
||||
### Message State
|
||||
|
||||
`AIConversationMessage` now stores:
|
||||
|
||||
- `Kind`: chat, reasoning, or tool call
|
||||
- `Visibility`: context or UI-only
|
||||
- `IncludeInContext`: whether the block should be replayed into the provider context window
|
||||
- `State`: in progress, completed, or failed
|
||||
- `TurnID`
|
||||
- `ItemID`
|
||||
- tool call metadata: name, arguments, result, error text, tool call id
|
||||
|
||||
This keeps one ordered transcript table while still distinguishing chat context from UI-only telemetry.
|
||||
|
||||
### UI-Only Persistence
|
||||
|
||||
Reasoning, tool-call, and intermediate assistant chat items can be stored as normal conversation messages with `Visibility = UI_ONLY`.
|
||||
|
||||
That gives Sliver:
|
||||
|
||||
- cross-client sync through the existing conversation storage path
|
||||
- transcript restoration after reconnect or TUI restart
|
||||
- ordered replay of an entire turn, including intermediate items
|
||||
|
||||
At the same time, the model context builder keys off `IncludeInContext`, so UI visibility and provider visibility are now independent.
|
||||
|
||||
### Event Model
|
||||
|
||||
Sliver now publishes `AIConversationEvent` payloads instead of raw `AIConversation` snapshots.
|
||||
|
||||
Supported event types:
|
||||
|
||||
- `CONVERSATION_UPDATED`
|
||||
- `CONVERSATION_DELETED`
|
||||
- `TURN_STARTED`
|
||||
- `TURN_COMPLETED`
|
||||
- `TURN_FAILED`
|
||||
- `MESSAGE_STARTED`
|
||||
- `MESSAGE_UPDATED`
|
||||
- `MESSAGE_COMPLETED`
|
||||
|
||||
Each event can carry:
|
||||
|
||||
- a conversation summary
|
||||
- a single message/item
|
||||
- the turn id
|
||||
- optional error text
|
||||
|
||||
This mirrors the Codex pattern of stable turn and item lifecycle events.
|
||||
|
||||
## Server Turn Flow
|
||||
|
||||
For each new user message:
|
||||
|
||||
1. The user chat message is saved as a normal context-visible chat item.
|
||||
2. The server resolves runtime provider/model settings.
|
||||
3. A new turn id is generated and persisted with `TURN_STARTED`.
|
||||
4. If the runtime supports the Responses API loop, Sliver runs the agentic loop.
|
||||
5. Reasoning items are persisted as completed UI-only reasoning blocks.
|
||||
6. Tool calls are persisted as in-progress UI-only tool blocks, executed locally, then updated in place to completed or failed.
|
||||
7. Intermediate assistant chat blocks on tool-use turns are persisted as UI-only chat items with `IncludeInContext = false`.
|
||||
8. Tool outputs are fed back into the Responses API with `previous_response_id`.
|
||||
9. When the model returns final assistant text, that assistant message is persisted as a normal context-visible chat item with `IncludeInContext = true`.
|
||||
10. The conversation turn state is cleared with `TURN_COMPLETED`.
|
||||
|
||||
If anything fails:
|
||||
|
||||
- a UI-only failed system message is persisted
|
||||
- the conversation is marked `TURN_FAILED`
|
||||
- the failure remains visible in the transcript without polluting future model context
|
||||
|
||||
## Tool Surface
|
||||
|
||||
The first port exposes a safe read-only tool set:
|
||||
|
||||
- `list_sessions_and_beacons`
|
||||
- `fs_ls`
|
||||
- `fs_cat`
|
||||
- `fs_pwd`
|
||||
|
||||
Why read-only first:
|
||||
|
||||
- it proves the loop, storage, eventing, and TUI behavior
|
||||
- it avoids surprising destructive remote actions in the first server-side tool release
|
||||
|
||||
Filesystem tools accept explicit `session_id` or `beacon_id`, but when omitted they fall back to the conversation target.
|
||||
|
||||
Beacon-backed tools wait for task completion on the server before emitting the completed item update.
|
||||
|
||||
## Context Filtering
|
||||
|
||||
Only chat messages with:
|
||||
|
||||
- `Kind = CHAT`
|
||||
- `IncludeInContext = true`
|
||||
|
||||
are included in future provider requests.
|
||||
|
||||
That means these items are excluded from the model context window:
|
||||
|
||||
- reasoning blocks
|
||||
- tool call blocks
|
||||
- UI-only intermediate assistant chat blocks
|
||||
- failed system status blocks
|
||||
|
||||
This is the main rule that preserves transcript richness without degrading prompt quality.
|
||||
|
||||
## TUI Behavior
|
||||
|
||||
The AI TUI now consumes `AIConversationEvent` incrementally.
|
||||
|
||||
Instead of reloading the whole conversation for every server event, it:
|
||||
|
||||
- merges conversation summaries into the sidebar
|
||||
- upserts messages in place by message id or item id
|
||||
- updates waiting state from `TurnState`
|
||||
- removes deleted conversations immediately
|
||||
|
||||
Reasoning and tool-call blocks render as dedicated dark-grey transcript boxes.
|
||||
|
||||
### Reasoning Blocks
|
||||
|
||||
Reasoning blocks show:
|
||||
|
||||
- the `Reasoning` label
|
||||
- the persisted reasoning text
|
||||
- item state metadata
|
||||
|
||||
### Tool Call Blocks
|
||||
|
||||
Tool blocks show:
|
||||
|
||||
- `Tool: <name>`
|
||||
- arguments
|
||||
- result
|
||||
- error text when present
|
||||
- item state metadata
|
||||
|
||||
These blocks are restored exactly like chat messages because they are stored in the same ordered transcript.
|
||||
|
||||
## Restore And Sync
|
||||
|
||||
Because UI-only items live in the database:
|
||||
|
||||
- another Sliver client can attach to the same conversation and see the full turn history
|
||||
- reconnecting the TUI restores reasoning and tool-call blocks
|
||||
- turn state survives restart, so "in progress" and "failed" status can be reconstructed from persisted conversation metadata
|
||||
|
||||
## Differences From Codex
|
||||
|
||||
This port intentionally differs from the example in a few places.
|
||||
|
||||
- Sliver currently uses non-streaming Responses API requests, so reasoning items arrive as completed items rather than token deltas.
|
||||
- Tool execution is currently a bounded read-only server-side function set, not the full Codex local tool surface.
|
||||
- Sliver persists intermediate items in its conversation DB rather than only in a transient thread runtime.
|
||||
|
||||
The important invariant is still the same:
|
||||
|
||||
- turn lifecycle is explicit
|
||||
- item lifecycle is explicit
|
||||
- UI-only items are persisted and replayable
|
||||
- only context-visible chat reaches the model
|
||||
Generated
+1764
-2661
File diff suppressed because it is too large.
Load diff
@@ -1,5 +1,18 @@
|
||||
Sliver allows an operator to extend the local client console and its features by adding new commands based on third party tools. The easiest way to install an alias or extension is using the [armory](/docs?name=Armory).
|
||||
|
||||
## Server AI Access
|
||||
|
||||
Aliases and extensions used from the client console still live under `~/.sliver-client/aliases` and `~/.sliver-client/extensions`.
|
||||
|
||||
The server-side AI agent uses a separate store under `~/.sliver/ai/aliases` and `~/.sliver/ai/extensions`. To make a package available to the AI agent, copy the unpacked package directory into the matching server-side path so the manifest and payload artifacts are present together. For example:
|
||||
|
||||
- `~/.sliver/ai/aliases/Rubeus/alias.json`
|
||||
- `~/.sliver/ai/aliases/Rubeus/Rubeus.exe`
|
||||
- `~/.sliver/ai/extensions/nanodump/extension.json`
|
||||
- `~/.sliver/ai/extensions/nanodump/nanodump.x64.o`
|
||||
|
||||
There is currently no dedicated UI for this workflow. The AI agent can search and execute only the packages that have been copied into these server-side directories.
|
||||
|
||||
#### Aliases Command Parsing
|
||||
|
||||
**⚠️ IMPORTANT:** It's important to understand that all alias commands have certain Sliver shell flags (those that appear in `--help`). Your Sliver shell commands are lexically parsed by the Sliver shell first, and only unnamed positional arguments are passed to the alias code. This means you may need to escape certain arguments in order for them to be correctly parsed.
|
||||
|
||||
@@ -4,6 +4,8 @@ The armory downloads packages from `github.com` and `api.github.com` so you'll n
|
||||
|
||||
Aliases and extensions are installed on the "sliver client"-side, and thus are not shared among operators in [multiplayer mode](/docs?name=Multi-player+Mode).
|
||||
|
||||
The server-side AI agent uses a separate store under `~/.sliver/ai/aliases` and `~/.sliver/ai/extensions`. If you want the AI agent to discover and execute an alias or extension, copy the unpacked package directory (the manifest plus its artifacts) into the matching server-side path. There is currently no dedicated UI or install command for the AI store.
|
||||
|
||||
You can also use `armory install all` to install _everything_ if you really want to.
|
||||
|
||||
## The Official Armory
|
||||
@@ -44,6 +46,8 @@ You remove packages installed from the `armory` using the `aliases rm` and `exte
|
||||
|
||||
Installed alias and extension files are stored in `~/.sliver-client/` by default if you want to manually remove a package simply delete its corresponding directory and restart the client.
|
||||
|
||||
If you also copied a package into the server-side AI store, remove that copy separately from `~/.sliver/ai/aliases` or `~/.sliver/ai/extensions`.
|
||||
|
||||
## Private Armories
|
||||
|
||||
Sliver has experimental support for self-hosted private armories, but I have not gotten around to testing and writing the documentation for these, so you'll have to read through the source code to figure out how they work for now. I'll eventually update this documentation and release a reference implementation. If you do play around with this, just know the design is subject to change before becoming a real feature.
|
||||
@@ -6,7 +6,7 @@ BOF support is provided via the [COFF Loader](https://github.com/sliverarmory/CO
|
||||
|
||||
The easiest way to install a BOF extension, for example [`nanodump`](https://github.com/sliverarmory/nanodump), is using the [armory](/docs?name=Armory) package manager:
|
||||
|
||||
**IMPORTANT:** BOF Extensions are installed per-sliver client, they are not stored on the server. Thus extensions are not shared across operators, each operator must install the extension to use it.
|
||||
**IMPORTANT:** BOF extensions are still installed per-sliver client for normal operator use, so they are not shared across operators automatically. If you want the server-side AI agent to use a BOF, copy the unpacked BOF directory into `~/.sliver/ai/extensions` on the server as well, and copy any dependency such as `coff-loader` into that same server-side AI store.
|
||||
|
||||
```
|
||||
sliver > armory install nanodump
|
||||
@@ -114,6 +114,8 @@ Looking at the script we can see the BOF requires a single integer argument. The
|
||||
|
||||
Once the manifest is defined load it into your client using `extensions load`, locally loaded extensions do not need to be cryptographically signed. The paths in the manifest should be relative to the manifest file, parent directories are not allowed.
|
||||
|
||||
If you want the server-side AI agent to discover and run the same BOF, copy that unpacked directory into `~/.sliver/ai/extensions/<name>/` so the `extension.json` file and `.o` artifact are available on the server.
|
||||
|
||||
More details can be found on the [Aliases & Extensions](/docs?name=Aliases+and+Extensions) page.
|
||||
|
||||
**IMPORTANT:** For BOF extensions to be properly detected by the Sliver client, the `path` must use the `.o` file extension.
|
||||
@@ -47,6 +47,10 @@ sliver > mcp start --transport http
|
||||
|
||||
[*] Starting MCP server (http) on 127.0.0.1:8080
|
||||
[*] Endpoint: http://127.0.0.1:8080/mcp
|
||||
[*] Auth Header: Authorization
|
||||
[*] Auth Token: 6f90c3b3c6058fa59f570e281f3f8d39
|
||||
```
|
||||
|
||||
You can use `mcp stop` to stop the server, see `mcp start --help` for additional options when starting the MCP server. Authentication is not yet supported for MCP over HTTP/SSE.
|
||||
You can use `mcp stop` to stop the server, see `mcp start --help` for additional options when starting the MCP server.
|
||||
|
||||
On first HTTP/SSE use, Sliver generates a random 128-bit token and saves it to `~/.sliver-client/mcp.yaml`. Every MCP HTTP/SSE request must include that token exactly in the `Authorization` header. If `mcp.yaml` already exists, the listener will only start when the token is at least 8 characters long.
|
||||
@@ -59,17 +59,17 @@ require (
|
||||
github.com/ulikunitz/xz v0.5.15
|
||||
github.com/xlab/treeprint v1.2.0
|
||||
github.com/yiya1989/sshkrb5 v0.0.1
|
||||
golang.org/x/crypto v0.48.0
|
||||
golang.org/x/crypto v0.49.0
|
||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa
|
||||
golang.org/x/mod v0.33.0
|
||||
golang.org/x/net v0.51.0
|
||||
golang.org/x/net v0.52.0
|
||||
golang.org/x/sys v0.42.0
|
||||
golang.org/x/term v0.40.0
|
||||
golang.org/x/text v0.34.0
|
||||
golang.org/x/term v0.41.0
|
||||
golang.org/x/text v0.35.0
|
||||
golang.zx2c4.com/wireguard v0.0.0-20250521234502-f333402bd9cb
|
||||
golang.zx2c4.com/wireguard/wgctrl v0.0.0-20241231184526-a9ab2273dd10
|
||||
google.golang.org/api v0.269.0
|
||||
google.golang.org/grpc v1.79.1
|
||||
google.golang.org/grpc v1.79.3
|
||||
google.golang.org/protobuf v1.36.11
|
||||
gorm.io/driver/mysql v1.6.0
|
||||
gorm.io/driver/postgres v1.6.0
|
||||
@@ -291,7 +291,7 @@ require (
|
||||
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0
|
||||
github.com/lxn/win v0.0.0-20210218163916-a377121e959e // indirect
|
||||
github.com/mark3labs/mcp-go v0.43.2
|
||||
github.com/mark3labs/mcp-go v0.45.0
|
||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.20 // indirect
|
||||
@@ -324,12 +324,12 @@ require (
|
||||
go4.org/mem v0.0.0-20240501181205-ae6ca9944745 // indirect
|
||||
go4.org/netipx v0.0.0-20231129151722-fdeea329fbba // indirect
|
||||
golang.org/x/oauth2 v0.36.0 // indirect
|
||||
golang.org/x/sync v0.19.0 // indirect
|
||||
golang.org/x/sync v0.20.0 // indirect
|
||||
golang.org/x/time v0.14.0
|
||||
golang.org/x/tools v0.42.0 // indirect
|
||||
golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
|
||||
golang.zx2c4.com/wireguard/windows v0.5.3 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5 // indirect
|
||||
gopkg.in/jcmturner/aescts.v1 v1.0.1 // indirect
|
||||
gopkg.in/jcmturner/dnsutils.v1 v1.0.1 // indirect
|
||||
gopkg.in/jcmturner/goidentity.v3 v3.0.0 // indirect
|
||||
|
||||
@@ -553,6 +553,8 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
|
||||
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
|
||||
github.com/mark3labs/mcp-go v0.43.2 h1:21PUSlWWiSbUPQwXIJ5WKlETixpFpq+WBpbMGDSVy/I=
|
||||
github.com/mark3labs/mcp-go v0.43.2/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw=
|
||||
github.com/mark3labs/mcp-go v0.45.0 h1:s0S8qR/9fWaQ3pHxz7pm1uQ0DrswoSnRIxKIjbiQtkc=
|
||||
github.com/mark3labs/mcp-go v0.45.0/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw=
|
||||
github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
|
||||
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
|
||||
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
||||
@@ -868,6 +870,8 @@ golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v
|
||||
golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
|
||||
golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts=
|
||||
golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos=
|
||||
golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4=
|
||||
golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA=
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa h1:Zt3DZoOFFYkKhDT3v7Lm9FDMEV06GpzjG2jrqW+QTE0=
|
||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa/go.mod h1:K79w1Vqn7PoiZn+TkNpx3BUWUQksGO3JcVX6qIjytmA=
|
||||
@@ -911,6 +915,8 @@ golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=
|
||||
golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=
|
||||
golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0=
|
||||
golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw=
|
||||
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
|
||||
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
||||
@@ -929,6 +935,8 @@ golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
|
||||
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
@@ -980,6 +988,8 @@ golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
|
||||
golang.org/x/term v0.40.0 h1:36e4zGLqU4yhjlmxEaagx2KuYbJq3EwY8K943ZsHcvg=
|
||||
golang.org/x/term v0.40.0/go.mod h1:w2P8uVp06p2iyKKuvXIm7N/y0UCRt3UfJTfZ7oOpglM=
|
||||
golang.org/x/term v0.41.0 h1:QCgPso/Q3RTJx2Th4bDLqML4W6iJiaXFq2/ftQF13YU=
|
||||
golang.org/x/term v0.41.0/go.mod h1:3pfBgksrReYfZ5lvYM0kSO0LIkAl4Yl2bXOkKP7Ec2A=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
@@ -995,6 +1005,8 @@ golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
|
||||
golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
|
||||
golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8=
|
||||
golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA=
|
||||
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
|
||||
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
@@ -1044,6 +1056,8 @@ google.golang.org/genproto/googleapis/api v0.0.0-20260226221140-a57be14db171 h1:
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20260226221140-a57be14db171/go.mod h1:M5krXqk4GhBKvB596udGL3UyjL4I1+cTbK0orROM9ng=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 h1:ggcbiqK8WWh6l1dnltU4BgWGIGo+EVYxCaAPih/zQXQ=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5 h1:aJmi6DVGGIStN9Mobk/tZOOQUBbj0BPjZjjnOdoZKts=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
|
||||
google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
|
||||
google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg=
|
||||
google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY=
|
||||
@@ -1051,6 +1065,8 @@ google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8
|
||||
google.golang.org/grpc v1.29.1/go.mod h1:itym6AZVZYACWQqET3MqgPpjcuV5QH3BxFS3IjizoKk=
|
||||
google.golang.org/grpc v1.79.1 h1:zGhSi45ODB9/p3VAawt9a+O/MULLl9dpizzNNpq7flY=
|
||||
google.golang.org/grpc v1.79.1/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
|
||||
google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
|
||||
google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
|
||||
google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
|
||||
google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
|
||||
google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
|
||||
|
||||
@@ -119,17 +119,20 @@ func pivotStopListenerHandler(envelope *pb.Envelope, connection *transports.Conn
|
||||
}
|
||||
|
||||
func pivotPeerEnvelopeHandler(envelope *pb.Envelope, connection *transports.Connection) {
|
||||
sent, err := pivots.SendToPeer(envelope)
|
||||
nextPeerID, sent, err := pivots.SendToPeer(envelope)
|
||||
if !sent {
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("Send to peer failed, report peer failure upstream ...")
|
||||
// {{end}}
|
||||
if nextPeerID == 0 {
|
||||
return
|
||||
}
|
||||
errStr := ""
|
||||
if err != nil {
|
||||
errStr = err.Error()
|
||||
}
|
||||
data, _ := proto.Marshal(&pb.PivotPeerFailure{
|
||||
PeerID: pivots.MyPeerID,
|
||||
PeerID: nextPeerID,
|
||||
Type: pb.PeerFailureType_SEND_FAILURE,
|
||||
Err: errStr,
|
||||
})
|
||||
|
||||
@@ -145,14 +145,14 @@ func StopListener(id uint32) {
|
||||
}
|
||||
|
||||
// SendToPeer - Forward an envelope to a peer
|
||||
func SendToPeer(envelope *pb.Envelope) (bool, error) {
|
||||
func SendToPeer(envelope *pb.Envelope) (int64, bool, error) {
|
||||
pivotPeerEnvelope := &pb.PivotPeerEnvelope{}
|
||||
err := proto.Unmarshal(envelope.Data, pivotPeerEnvelope)
|
||||
if err != nil {
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("Failed to unmarshal peer envelope: %s", err)
|
||||
// {{end}}
|
||||
return false, err
|
||||
return 0, false, err
|
||||
}
|
||||
|
||||
// {{if .Config.Debug}}
|
||||
@@ -164,7 +164,7 @@ func SendToPeer(envelope *pb.Envelope) (bool, error) {
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("Failed to find next peer id: %s", err)
|
||||
// {{end}}
|
||||
return false, err
|
||||
return 0, false, err
|
||||
}
|
||||
|
||||
sent := false // Controls iteration of outer loop
|
||||
@@ -183,9 +183,9 @@ func SendToPeer(envelope *pb.Envelope) (bool, error) {
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("Failed to find peer with id %d", nextPeerID)
|
||||
// {{end}}
|
||||
return false, errors.New("peer not found")
|
||||
return nextPeerID, false, errors.New("peer not found")
|
||||
}
|
||||
return sent, nil
|
||||
return nextPeerID, sent, nil
|
||||
}
|
||||
|
||||
func findNextPeerID(pivotPeerEnvelope *pb.PivotPeerEnvelope) (int64, error) {
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
package version
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2021 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type linuxVersionInfo struct {
|
||||
Sysname string
|
||||
Release string
|
||||
Version string
|
||||
Machine string
|
||||
}
|
||||
|
||||
func readLinuxOSRelease() string {
|
||||
paths := []string{"/etc/os-release", "/usr/lib/os-release"}
|
||||
for _, path := range paths {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if osRelease := buildLinuxOSRelease(parseLinuxOSRelease(string(data))); osRelease != "" {
|
||||
return osRelease
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func parseLinuxOSRelease(content string) map[string]string {
|
||||
values := make(map[string]string)
|
||||
for _, line := range strings.Split(content, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
key, value, ok := strings.Cut(line, "=")
|
||||
if !ok || key == "" {
|
||||
continue
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
value = unescapeLinuxOSReleaseValue(unquoteLinuxOSReleaseValue(strings.TrimSpace(value)))
|
||||
values[key] = value
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func unquoteLinuxOSReleaseValue(s string) string {
|
||||
if len(s) < 2 {
|
||||
return s
|
||||
}
|
||||
if (s[0] == '"' || s[0] == '\'') && s[0] == s[len(s)-1] {
|
||||
return s[1 : len(s)-1]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func unescapeLinuxOSReleaseValue(s string) string {
|
||||
return strings.NewReplacer(
|
||||
`\$`, `$`,
|
||||
`\"`, `"`,
|
||||
`\'`, `'`,
|
||||
`\\`, `\`,
|
||||
"\\`", "`",
|
||||
).Replace(s)
|
||||
}
|
||||
|
||||
func buildLinuxOSRelease(values map[string]string) string {
|
||||
name := values["NAME"]
|
||||
version := values["VERSION"]
|
||||
if version == "" {
|
||||
version = values["VERSION_ID"]
|
||||
}
|
||||
if name != "" && version != "" {
|
||||
return fmt.Sprintf("%s %s", name, version)
|
||||
}
|
||||
if values["PRETTY_NAME"] != "" {
|
||||
return values["PRETTY_NAME"]
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
func normalizeLinuxArch(arch string) string {
|
||||
arch = strings.TrimSpace(arch)
|
||||
if arch == "" {
|
||||
return ""
|
||||
}
|
||||
switch strings.ToLower(arch) {
|
||||
case "amd64":
|
||||
return "x86_64"
|
||||
case "i386", "i486", "i586", "i686", "x86":
|
||||
return "x86"
|
||||
default:
|
||||
return arch
|
||||
}
|
||||
}
|
||||
|
||||
func formatLinuxDetailedVersion(osRelease string, versionInfo linuxVersionInfo) string {
|
||||
parts := make([]string, 0, 3)
|
||||
if osRelease != "" {
|
||||
parts = append(parts, osRelease)
|
||||
} else if versionInfo.Sysname != "" {
|
||||
parts = append(parts, versionInfo.Sysname)
|
||||
}
|
||||
|
||||
kernelParts := make([]string, 0, 2)
|
||||
if versionInfo.Release != "" {
|
||||
kernelParts = append(kernelParts, versionInfo.Release)
|
||||
}
|
||||
if versionInfo.Version != "" {
|
||||
kernelParts = append(kernelParts, versionInfo.Version)
|
||||
}
|
||||
if len(kernelParts) > 0 {
|
||||
parts = append(parts, fmt.Sprintf("kernel %s", strings.Join(kernelParts, " ")))
|
||||
}
|
||||
|
||||
if arch := normalizeLinuxArch(versionInfo.Machine); arch != "" {
|
||||
parts = append(parts, arch)
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package version
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseLinuxOSRelease(t *testing.T) {
|
||||
content := "# comment\n" +
|
||||
"NAME=\"Ubuntu\"\n" +
|
||||
"VERSION=\"24.04.2 LTS (Noble Numbat)\"\n" +
|
||||
"PRETTY_NAME=\"Ubuntu 24.04.2 LTS\"\n" +
|
||||
"SPECIAL=\"quoted \\\"value\\\" \\$HOME \\`uname\\`\"\n"
|
||||
|
||||
values := parseLinuxOSRelease(content)
|
||||
|
||||
if got, want := values["NAME"], "Ubuntu"; got != want {
|
||||
t.Fatalf("NAME = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := values["VERSION"], "24.04.2 LTS (Noble Numbat)"; got != want {
|
||||
t.Fatalf("VERSION = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := values["PRETTY_NAME"], "Ubuntu 24.04.2 LTS"; got != want {
|
||||
t.Fatalf("PRETTY_NAME = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := values["SPECIAL"], "quoted \"value\" $HOME `uname`"; got != want {
|
||||
t.Fatalf("SPECIAL = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildLinuxOSRelease(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
values map[string]string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "prefers name and version",
|
||||
values: map[string]string{
|
||||
"NAME": "Ubuntu",
|
||||
"VERSION": "24.04.2 LTS (Noble Numbat)",
|
||||
"PRETTY_NAME": "Ubuntu 24.04.2 LTS",
|
||||
},
|
||||
want: "Ubuntu 24.04.2 LTS (Noble Numbat)",
|
||||
},
|
||||
{
|
||||
name: "falls back to version id",
|
||||
values: map[string]string{
|
||||
"NAME": "Alpine Linux",
|
||||
"VERSION_ID": "3.21.3",
|
||||
},
|
||||
want: "Alpine Linux 3.21.3",
|
||||
},
|
||||
{
|
||||
name: "falls back to pretty name",
|
||||
values: map[string]string{
|
||||
"PRETTY_NAME": "Fedora Linux 41 (Workstation Edition)",
|
||||
},
|
||||
want: "Fedora Linux 41 (Workstation Edition)",
|
||||
},
|
||||
{
|
||||
name: "falls back to name",
|
||||
values: map[string]string{
|
||||
"NAME": "Linux",
|
||||
},
|
||||
want: "Linux",
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
if got := buildLinuxOSRelease(testCase.values); got != testCase.want {
|
||||
t.Fatalf("buildLinuxOSRelease() = %q, want %q", got, testCase.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatLinuxDetailedVersion(t *testing.T) {
|
||||
version := formatLinuxDetailedVersion("Ubuntu 24.04.2 LTS (Noble Numbat)", linuxVersionInfo{
|
||||
Release: "6.8.0-57-generic",
|
||||
Version: "#59-Ubuntu SMP PREEMPT_DYNAMIC Fri Feb 14 18:29:04 UTC 2025",
|
||||
Machine: "amd64",
|
||||
})
|
||||
|
||||
want := "Ubuntu 24.04.2 LTS (Noble Numbat) kernel 6.8.0-57-generic #59-Ubuntu SMP PREEMPT_DYNAMIC Fri Feb 14 18:29:04 UTC 2025 x86_64"
|
||||
if version != want {
|
||||
t.Fatalf("formatLinuxDetailedVersion() = %q, want %q", version, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatLinuxDetailedVersionFallback(t *testing.T) {
|
||||
version := formatLinuxDetailedVersion("", linuxVersionInfo{
|
||||
Sysname: "Linux",
|
||||
Release: "6.12.0",
|
||||
Machine: "i686",
|
||||
})
|
||||
|
||||
want := "Linux kernel 6.12.0 x86"
|
||||
if version != want {
|
||||
t.Fatalf("formatLinuxDetailedVersion() = %q, want %q", version, want)
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build !darwin && !(windows && (amd64 || 386)) && !(linux && (amd64 || 386))
|
||||
//go:build !darwin && !linux && !(windows && (amd64 || 386))
|
||||
|
||||
package version
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build amd64 || 386
|
||||
//go:build linux
|
||||
|
||||
package version
|
||||
|
||||
@@ -21,132 +21,36 @@ package version
|
||||
*/
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"bytes"
|
||||
//{{if .Config.Debug}}
|
||||
"log"
|
||||
//{{end}}
|
||||
"os"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
"golang.org/x/sys/unix"
|
||||
)
|
||||
|
||||
func getString(input [65]int8) string {
|
||||
var buf [65]byte
|
||||
for i, b := range input {
|
||||
buf[i] = byte(b)
|
||||
func linuxNullTerminatedString(input []byte) string {
|
||||
if index := bytes.IndexByte(input, 0); index >= 0 {
|
||||
input = input[:index]
|
||||
}
|
||||
ver := string(buf[:])
|
||||
if i := strings.Index(ver, "\x00"); i != -1 {
|
||||
ver = ver[:i]
|
||||
}
|
||||
return ver
|
||||
}
|
||||
|
||||
func readOSRelease() string {
|
||||
paths := []string{"/etc/os-release", "/usr/lib/os-release"}
|
||||
for _, path := range paths {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
return buildOSRelease(parseOSRelease(string(data)))
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func parseOSRelease(content string) map[string]string {
|
||||
values := make(map[string]string)
|
||||
for _, line := range strings.Split(content, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
key, value, ok := strings.Cut(line, "=")
|
||||
if !ok || key == "" {
|
||||
continue
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
value = unescape(unquote(strings.TrimSpace(value)))
|
||||
values[key] = value
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func unquote(s string) string {
|
||||
if len(s) < 2 {
|
||||
return s
|
||||
}
|
||||
if (s[0] == '"' || s[0] == '\'') && s[0] == s[len(s)-1] {
|
||||
return s[1 : len(s)-1]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func unescape(s string) string {
|
||||
return strings.NewReplacer(
|
||||
`\$`, `$`,
|
||||
`\"`, `"`,
|
||||
`\'`, `'`,
|
||||
`\\`, `\`,
|
||||
"\\`", "`",
|
||||
).Replace(s)
|
||||
}
|
||||
|
||||
func buildOSRelease(values map[string]string) string {
|
||||
name := values["NAME"]
|
||||
version := values["VERSION"]
|
||||
if version == "" {
|
||||
version = values["VERSION_ID"]
|
||||
}
|
||||
if name != "" && version != "" {
|
||||
return fmt.Sprintf("%s %s", name, version)
|
||||
}
|
||||
return values["PRETTY_NAME"]
|
||||
}
|
||||
|
||||
func normalizeArch(arch string) string {
|
||||
arch = strings.TrimSpace(arch)
|
||||
if arch == "" {
|
||||
return ""
|
||||
}
|
||||
switch strings.ToLower(arch) {
|
||||
case "amd64":
|
||||
return "x86_64"
|
||||
case "i386", "i486", "i586", "i686", "x86":
|
||||
return "x86"
|
||||
default:
|
||||
return arch
|
||||
}
|
||||
}
|
||||
|
||||
func formatDetailedVersion(osRelease, kernel, arch string) string {
|
||||
parts := make([]string, 0, 3)
|
||||
if osRelease != "" {
|
||||
parts = append(parts, osRelease)
|
||||
}
|
||||
if kernel != "" {
|
||||
parts = append(parts, fmt.Sprintf("kernel %s", kernel))
|
||||
}
|
||||
if arch != "" {
|
||||
parts = append(parts, arch)
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
return string(input)
|
||||
}
|
||||
|
||||
// GetVersion returns the os version information
|
||||
func GetVersion() string {
|
||||
osRelease := readOSRelease()
|
||||
var uname syscall.Utsname
|
||||
if err := syscall.Uname(&uname); err != nil {
|
||||
osRelease := readLinuxOSRelease()
|
||||
var uname unix.Utsname
|
||||
if err := unix.Uname(&uname); err != nil {
|
||||
//{{if .Config.Debug}}
|
||||
log.Printf("error getting OS version: %v", err)
|
||||
//{{end}}
|
||||
return osRelease
|
||||
}
|
||||
if osRelease == "" {
|
||||
osRelease = getString(uname.Sysname)
|
||||
}
|
||||
kernel := getString(uname.Release)
|
||||
arch := normalizeArch(getString(uname.Machine))
|
||||
return formatDetailedVersion(osRelease, kernel, arch)
|
||||
|
||||
return formatLinuxDetailedVersion(osRelease, linuxVersionInfo{
|
||||
Sysname: linuxNullTerminatedString(uname.Sysname[:]),
|
||||
Release: linuxNullTerminatedString(uname.Release[:]),
|
||||
Version: linuxNullTerminatedString(uname.Version[:]),
|
||||
Machine: linuxNullTerminatedString(uname.Machine[:]),
|
||||
})
|
||||
}
|
||||
+1000
-357
File diff suppressed because it is too large.
Load diff
@@ -1426,6 +1426,7 @@ message AIConfigSummary {
|
||||
string ThinkingLevel = 3;
|
||||
bool Valid = 4;
|
||||
string Error = 5;
|
||||
string SystemPrompt = 6;
|
||||
}
|
||||
|
||||
message AIConversationReq {
|
||||
@@ -1437,6 +1438,49 @@ message AIConversations {
|
||||
repeated AIConversation Conversations = 1;
|
||||
}
|
||||
|
||||
enum AIConversationTurnState {
|
||||
AI_TURN_STATE_IDLE = 0;
|
||||
AI_TURN_STATE_IN_PROGRESS = 1;
|
||||
AI_TURN_STATE_FAILED = 2;
|
||||
}
|
||||
|
||||
enum AIConversationMessageKind {
|
||||
AI_MESSAGE_KIND_CHAT = 0;
|
||||
AI_MESSAGE_KIND_REASONING = 1;
|
||||
AI_MESSAGE_KIND_TOOL_CALL = 2;
|
||||
}
|
||||
|
||||
enum AIConversationMessageVisibility {
|
||||
AI_MESSAGE_VISIBILITY_CONTEXT = 0;
|
||||
AI_MESSAGE_VISIBILITY_UI_ONLY = 1;
|
||||
}
|
||||
|
||||
enum AIConversationMessageState {
|
||||
AI_MESSAGE_STATE_COMPLETED = 0;
|
||||
AI_MESSAGE_STATE_IN_PROGRESS = 1;
|
||||
AI_MESSAGE_STATE_FAILED = 2;
|
||||
}
|
||||
|
||||
enum AIConversationEventType {
|
||||
AI_CONVERSATION_EVENT_TYPE_UNKNOWN = 0;
|
||||
AI_CONVERSATION_EVENT_TYPE_CONVERSATION_UPDATED = 1;
|
||||
AI_CONVERSATION_EVENT_TYPE_CONVERSATION_DELETED = 2;
|
||||
AI_CONVERSATION_EVENT_TYPE_TURN_STARTED = 3;
|
||||
AI_CONVERSATION_EVENT_TYPE_TURN_COMPLETED = 4;
|
||||
AI_CONVERSATION_EVENT_TYPE_TURN_FAILED = 5;
|
||||
AI_CONVERSATION_EVENT_TYPE_MESSAGE_STARTED = 6;
|
||||
AI_CONVERSATION_EVENT_TYPE_MESSAGE_UPDATED = 7;
|
||||
AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED = 8;
|
||||
}
|
||||
|
||||
message AIContextWindowUsage {
|
||||
int64 InputTokens = 1;
|
||||
int64 OutputTokens = 2;
|
||||
int64 TotalTokens = 3;
|
||||
int64 ContextWindowTokens = 4;
|
||||
bool ContextWindowTokensEstimated = 5;
|
||||
}
|
||||
|
||||
message AIConversation {
|
||||
string ID = 1;
|
||||
int64 CreatedAt = 2;
|
||||
@@ -1448,6 +1492,12 @@ message AIConversation {
|
||||
string Summary = 8;
|
||||
string SystemPrompt = 9;
|
||||
repeated AIConversationMessage Messages = 10;
|
||||
string ActiveTurnID = 11;
|
||||
AIConversationTurnState TurnState = 12;
|
||||
string TargetSessionID = 13;
|
||||
string TargetBeaconID = 14;
|
||||
string ThinkingLevel = 15;
|
||||
AIContextWindowUsage ContextWindowUsage = 16;
|
||||
}
|
||||
|
||||
message AIConversationMessages {
|
||||
@@ -1468,6 +1518,26 @@ message AIConversationMessage {
|
||||
string Content = 10;
|
||||
string ProviderMessageID = 11;
|
||||
string FinishReason = 12;
|
||||
AIConversationMessageKind Kind = 13;
|
||||
AIConversationMessageVisibility Visibility = 14;
|
||||
AIConversationMessageState State = 15;
|
||||
string TurnID = 16;
|
||||
string ItemID = 17;
|
||||
string ToolCallID = 18;
|
||||
string ToolName = 19;
|
||||
string ToolArguments = 20;
|
||||
string ToolResult = 21;
|
||||
string ErrorText = 22;
|
||||
optional bool IncludeInContext = 23;
|
||||
}
|
||||
|
||||
message AIConversationEvent {
|
||||
AIConversationEventType EventType = 1;
|
||||
string ConversationID = 2;
|
||||
string TurnID = 3;
|
||||
string ErrorText = 4;
|
||||
AIConversation Conversation = 5;
|
||||
AIConversationMessage Message = 6;
|
||||
}
|
||||
|
||||
// watchtower
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.34.0
|
||||
// protoc v7.34.1
|
||||
// source: commonpb/common.proto
|
||||
|
||||
package commonpb
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.34.0
|
||||
// protoc v7.34.1
|
||||
// source: dnspb/dns.proto
|
||||
|
||||
package dnspb
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.34.0
|
||||
// protoc v7.34.1
|
||||
// source: rpcpb/services.proto
|
||||
|
||||
package rpcpb
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
|
||||
// versions:
|
||||
// - protoc-gen-go-grpc v1.6.1
|
||||
// - protoc v7.34.0
|
||||
// - protoc v7.34.1
|
||||
// source: rpcpb/services.proto
|
||||
|
||||
package rpcpb
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.34.0
|
||||
// protoc v7.34.1
|
||||
// source: sliverpb/sliver.proto
|
||||
|
||||
package sliverpb
|
||||
|
||||
@@ -0,0 +1,592 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
openai "github.com/openai/openai-go/v2"
|
||||
openairesponses "github.com/openai/openai-go/v2/responses"
|
||||
"github.com/openai/openai-go/v2/shared"
|
||||
)
|
||||
|
||||
const maxAgenticLoopIterations = 100
|
||||
|
||||
// AgenticToolDefinition describes a custom function tool exposed to the model.
|
||||
type AgenticToolDefinition struct {
|
||||
Name string
|
||||
Description string
|
||||
Parameters map[string]any
|
||||
}
|
||||
|
||||
// AgenticToolExecutor executes model-requested tools for a conversation turn.
|
||||
type AgenticToolExecutor interface {
|
||||
ToolDefinitions() []AgenticToolDefinition
|
||||
CallTool(context.Context, string, string) (string, error)
|
||||
}
|
||||
|
||||
// AgenticEventSink persists and publishes non-context-window UI items.
|
||||
type AgenticEventSink interface {
|
||||
ChatMessage(context.Context, AgenticChatMessage) error
|
||||
ReasoningItem(context.Context, AgenticReasoningItem) error
|
||||
ToolCallStarted(context.Context, AgenticToolCall) error
|
||||
ToolCallCompleted(context.Context, AgenticToolCallResult) error
|
||||
}
|
||||
|
||||
// AgenticChatMessage is a provider-emitted assistant or system chat block.
|
||||
type AgenticChatMessage struct {
|
||||
ItemID string
|
||||
Role string
|
||||
Content string
|
||||
Status string
|
||||
UIOnly bool
|
||||
IncludeInContext bool
|
||||
}
|
||||
|
||||
// AgenticReasoningItem is a completed reasoning item emitted by the provider.
|
||||
type AgenticReasoningItem struct {
|
||||
ItemID string
|
||||
Summary string
|
||||
Content string
|
||||
EncryptedContent string
|
||||
Status string
|
||||
}
|
||||
|
||||
// AgenticToolCall is a provider-emitted function tool call.
|
||||
type AgenticToolCall struct {
|
||||
ItemID string
|
||||
CallID string
|
||||
Name string
|
||||
Arguments string
|
||||
Status string
|
||||
}
|
||||
|
||||
// AgenticToolCallResult is the resolved tool call state after local execution.
|
||||
type AgenticToolCallResult struct {
|
||||
AgenticToolCall
|
||||
Output string
|
||||
Error error
|
||||
}
|
||||
|
||||
// SupportsAgenticConversation reports whether the configured runtime can run the
|
||||
// Responses API tool loop used by the Sliver AI chat.
|
||||
func SupportsAgenticConversation(runtime *RuntimeConfig) bool {
|
||||
if runtime == nil || !runtime.UseResponsesAPI {
|
||||
return false
|
||||
}
|
||||
|
||||
switch NormalizeProviderName(runtime.Provider) {
|
||||
case ProviderOpenAI, ProviderOpenAICompat, ProviderOpenRouter:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// CompleteConversationAgentic runs a Codex-style Responses API loop that emits
|
||||
// reasoning and tool-call items between the user prompt and the final assistant
|
||||
// reply. The returned completion is the assistant message that should be stored
|
||||
// in the conversation context window.
|
||||
func CompleteConversationAgentic(
|
||||
ctx context.Context,
|
||||
runtime *RuntimeConfig,
|
||||
conversation *clientpb.AIConversation,
|
||||
tools AgenticToolExecutor,
|
||||
sink AgenticEventSink,
|
||||
) (*Completion, error) {
|
||||
if runtime == nil {
|
||||
return nil, fmt.Errorf("AI runtime config is required")
|
||||
}
|
||||
if conversation == nil {
|
||||
return nil, fmt.Errorf("AI conversation is required")
|
||||
}
|
||||
if !SupportsAgenticConversation(runtime) {
|
||||
return nil, fmt.Errorf("AI provider %q does not support the agentic Responses API loop", runtime.Provider)
|
||||
}
|
||||
|
||||
systemPrompt, messages, err := conversationHistory(conversation)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client := newOpenAIClient(openAIRequestOptions(runtime)...)
|
||||
currentInput := buildResponseInput(systemPrompt, messages)
|
||||
toolParams := appendOpenAIResponseTools(runtime, buildAgenticToolParams(tools))
|
||||
previousResponseID := ""
|
||||
var maxUsage *ContextWindowUsage
|
||||
|
||||
for step := 0; step < maxAgenticLoopIterations; step++ {
|
||||
params := openairesponses.ResponseNewParams{
|
||||
Input: openairesponses.ResponseNewParamsInputUnion{
|
||||
OfInputItemList: currentInput,
|
||||
},
|
||||
Model: shared.ResponsesModel(runtime.Model),
|
||||
}
|
||||
if previousResponseID != "" {
|
||||
params.PreviousResponseID = openai.String(previousResponseID)
|
||||
}
|
||||
if len(toolParams) > 0 {
|
||||
params.Tools = toolParams
|
||||
params.ParallelToolCalls = openai.Bool(false)
|
||||
}
|
||||
if runtime.MaxOutputTokens > 0 {
|
||||
params.MaxOutputTokens = openai.Int(runtime.MaxOutputTokens)
|
||||
}
|
||||
if runtime.Temperature != nil {
|
||||
params.Temperature = openai.Float(*runtime.Temperature)
|
||||
}
|
||||
if runtime.TopP != nil {
|
||||
params.TopP = openai.Float(*runtime.TopP)
|
||||
}
|
||||
if reasoning, ok := responseReasoningParam(runtime); ok {
|
||||
params.Reasoning = reasoning
|
||||
}
|
||||
|
||||
response, err := client.responses.New(ctx, params)
|
||||
if err != nil {
|
||||
return nil, formatOpenAIError(runtime.Provider, err)
|
||||
}
|
||||
|
||||
maxUsage = mergeContextWindowUsage(maxUsage, resolveContextWindowUsage(
|
||||
ctx,
|
||||
runtime,
|
||||
fallbackString(response.Model, runtime.Model),
|
||||
response.Usage.InputTokens,
|
||||
response.Usage.OutputTokens,
|
||||
response.Usage.TotalTokens,
|
||||
))
|
||||
previousResponseID = strings.TrimSpace(response.ID)
|
||||
nextInput := make(openairesponses.ResponseInputParam, 0)
|
||||
finalTexts := make([]string, 0, len(response.Output))
|
||||
chatItems := make([]AgenticChatMessage, 0)
|
||||
|
||||
for _, item := range response.Output {
|
||||
switch output := item.AsAny().(type) {
|
||||
case openairesponses.ResponseReasoningItem:
|
||||
if sink != nil {
|
||||
if err := sink.ReasoningItem(ctx, AgenticReasoningItem{
|
||||
ItemID: strings.TrimSpace(output.ID),
|
||||
Summary: responseReasoningSummary(output),
|
||||
Content: responseReasoningContent(output),
|
||||
EncryptedContent: strings.TrimSpace(output.EncryptedContent),
|
||||
Status: strings.TrimSpace(string(output.Status)),
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
case openairesponses.ResponseFunctionToolCall:
|
||||
toolCall := AgenticToolCall{
|
||||
ItemID: strings.TrimSpace(output.ID),
|
||||
CallID: strings.TrimSpace(output.CallID),
|
||||
Name: strings.TrimSpace(output.Name),
|
||||
Arguments: strings.TrimSpace(output.Arguments),
|
||||
Status: strings.TrimSpace(string(output.Status)),
|
||||
}
|
||||
if sink != nil {
|
||||
if err := sink.ToolCallStarted(ctx, toolCall); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
toolOutput, toolErr := agenticToolOutput(ctx, tools, toolCall)
|
||||
if sink != nil {
|
||||
if err := sink.ToolCallCompleted(ctx, AgenticToolCallResult{
|
||||
AgenticToolCall: toolCall,
|
||||
Output: toolOutput,
|
||||
Error: toolErr,
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
nextInput = append(nextInput, openairesponses.ResponseInputItemParamOfFunctionCallOutput(toolCall.CallID, toolOutput))
|
||||
|
||||
case openairesponses.ResponseOutputMessage:
|
||||
if text := strings.TrimSpace(responseOutputMessageText(output)); text != "" {
|
||||
chatItems = append(chatItems, AgenticChatMessage{
|
||||
ItemID: strings.TrimSpace(output.ID),
|
||||
Role: strings.TrimSpace(string(output.Role)),
|
||||
Content: text,
|
||||
Status: strings.TrimSpace(string(output.Status)),
|
||||
IncludeInContext: true,
|
||||
})
|
||||
finalTexts = append(finalTexts, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(nextInput) > 0 {
|
||||
if sink != nil {
|
||||
for _, chatItem := range chatItems {
|
||||
chatItem.UIOnly = true
|
||||
chatItem.IncludeInContext = false
|
||||
if err := sink.ChatMessage(ctx, chatItem); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
currentInput = nextInput
|
||||
continue
|
||||
}
|
||||
|
||||
finalText := joinTextBlocks(finalTexts)
|
||||
if finalText == "" {
|
||||
finalText = strings.TrimSpace(response.OutputText())
|
||||
}
|
||||
if finalText == "" {
|
||||
return nil, fmt.Errorf("%s response did not include assistant text", runtime.Provider)
|
||||
}
|
||||
|
||||
return &Completion{
|
||||
Provider: runtime.Provider,
|
||||
Model: fallbackString(response.Model, runtime.Model),
|
||||
Content: finalText,
|
||||
ProviderMessageID: strings.TrimSpace(response.ID),
|
||||
FinishReason: responseFinishReason(response),
|
||||
ContextWindowUsage: maxUsage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("AI agent loop exceeded %d provider steps", maxAgenticLoopIterations)
|
||||
}
|
||||
|
||||
func buildAgenticToolParams(tools AgenticToolExecutor) []openairesponses.ToolUnionParam {
|
||||
if tools == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
definitions := tools.ToolDefinitions()
|
||||
if len(definitions) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
params := make([]openairesponses.ToolUnionParam, 0, len(definitions))
|
||||
for _, definition := range definitions {
|
||||
name := strings.TrimSpace(definition.Name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
tool := openairesponses.ToolParamOfFunction(name, normalizeOpenAIStrictSchema(definition.Parameters), true)
|
||||
if description := strings.TrimSpace(definition.Description); description != "" && tool.OfFunction != nil {
|
||||
tool.OfFunction.Description = openai.String(description)
|
||||
}
|
||||
params = append(params, tool)
|
||||
}
|
||||
if len(params) == 0 {
|
||||
return nil
|
||||
}
|
||||
return params
|
||||
}
|
||||
|
||||
func normalizeOpenAIStrictSchema(schema map[string]any) map[string]any {
|
||||
if len(schema) == 0 {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"required": []string{},
|
||||
"additionalProperties": false,
|
||||
}
|
||||
}
|
||||
|
||||
cloned, ok := cloneJSONValue(schema).(map[string]any)
|
||||
if !ok || cloned == nil {
|
||||
return schema
|
||||
}
|
||||
|
||||
normalizeOpenAIStrictSchemaNode(cloned)
|
||||
return cloned
|
||||
}
|
||||
|
||||
func normalizeOpenAIStrictSchemaNode(node map[string]any) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if properties, ok := schemaProperties(node); ok {
|
||||
required := schemaRequiredSet(node["required"])
|
||||
names := make([]string, 0, len(properties))
|
||||
for name, rawProperty := range properties {
|
||||
names = append(names, name)
|
||||
|
||||
property, ok := rawProperty.(map[string]any)
|
||||
if !ok || property == nil {
|
||||
continue
|
||||
}
|
||||
if !required[name] {
|
||||
makeJSONSchemaNullable(property)
|
||||
}
|
||||
normalizeOpenAIStrictSchemaNode(property)
|
||||
}
|
||||
sort.Strings(names)
|
||||
node["required"] = names
|
||||
if _, exists := node["additionalProperties"]; !exists {
|
||||
node["additionalProperties"] = false
|
||||
}
|
||||
}
|
||||
|
||||
if items, ok := node["items"].(map[string]any); ok {
|
||||
normalizeOpenAIStrictSchemaNode(items)
|
||||
}
|
||||
|
||||
for _, key := range []string{"anyOf", "allOf", "oneOf"} {
|
||||
variants, ok := node[key].([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, rawVariant := range variants {
|
||||
variant, ok := rawVariant.(map[string]any)
|
||||
if !ok || variant == nil {
|
||||
continue
|
||||
}
|
||||
normalizeOpenAIStrictSchemaNode(variant)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func schemaProperties(node map[string]any) (map[string]any, bool) {
|
||||
if node == nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
rawProperties, ok := node["properties"]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
properties, ok := rawProperties.(map[string]any)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return properties, true
|
||||
}
|
||||
|
||||
func schemaRequiredSet(raw any) map[string]bool {
|
||||
required := map[string]bool{}
|
||||
switch values := raw.(type) {
|
||||
case []string:
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" {
|
||||
required[value] = true
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, rawValue := range values {
|
||||
value, ok := rawValue.(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" {
|
||||
required[value] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return required
|
||||
}
|
||||
|
||||
func makeJSONSchemaNullable(node map[string]any) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
|
||||
switch value := node["type"].(type) {
|
||||
case string:
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" && value != "null" {
|
||||
node["type"] = []string{value, "null"}
|
||||
}
|
||||
case []string:
|
||||
if !stringSliceContains(value, "null") {
|
||||
node["type"] = append(append([]string(nil), value...), "null")
|
||||
}
|
||||
case []any:
|
||||
if !anyStringSliceContains(value, "null") {
|
||||
node["type"] = append(append([]any(nil), value...), "null")
|
||||
}
|
||||
default:
|
||||
if _, ok := node["anyOf"]; ok {
|
||||
appendSchemaVariant(node, "anyOf", map[string]any{"type": "null"})
|
||||
} else if _, ok := node["properties"]; ok {
|
||||
node["type"] = []string{"object", "null"}
|
||||
} else if _, ok := node["items"]; ok {
|
||||
node["type"] = []string{"array", "null"}
|
||||
}
|
||||
}
|
||||
|
||||
if enum, ok := nullableEnum(node["enum"]); ok {
|
||||
node["enum"] = enum
|
||||
}
|
||||
}
|
||||
|
||||
func appendSchemaVariant(node map[string]any, key string, variant map[string]any) {
|
||||
if node == nil || variant == nil {
|
||||
return
|
||||
}
|
||||
|
||||
rawVariants, ok := node[key].([]any)
|
||||
if !ok {
|
||||
node[key] = []any{variant}
|
||||
return
|
||||
}
|
||||
for _, rawVariant := range rawVariants {
|
||||
existing, ok := rawVariant.(map[string]any)
|
||||
if !ok || existing == nil {
|
||||
continue
|
||||
}
|
||||
if len(existing) == len(variant) {
|
||||
matched := true
|
||||
for variantKey, variantValue := range variant {
|
||||
if existing[variantKey] != variantValue {
|
||||
matched = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if matched {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
node[key] = append(rawVariants, variant)
|
||||
}
|
||||
|
||||
func nullableEnum(raw any) (any, bool) {
|
||||
switch values := raw.(type) {
|
||||
case []any:
|
||||
for _, value := range values {
|
||||
if value == nil {
|
||||
return values, true
|
||||
}
|
||||
}
|
||||
return append(append([]any(nil), values...), nil), true
|
||||
case []string:
|
||||
enum := make([]any, 0, len(values)+1)
|
||||
for _, value := range values {
|
||||
enum = append(enum, value)
|
||||
}
|
||||
enum = append(enum, nil)
|
||||
return enum, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func stringSliceContains(values []string, needle string) bool {
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) == needle {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func anyStringSliceContains(values []any, needle string) bool {
|
||||
for _, rawValue := range values {
|
||||
value, ok := rawValue.(string)
|
||||
if ok && strings.TrimSpace(value) == needle {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func cloneJSONValue(value any) any {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
cloned := make(map[string]any, len(typed))
|
||||
for key, nested := range typed {
|
||||
cloned[key] = cloneJSONValue(nested)
|
||||
}
|
||||
return cloned
|
||||
case []any:
|
||||
cloned := make([]any, len(typed))
|
||||
for idx, nested := range typed {
|
||||
cloned[idx] = cloneJSONValue(nested)
|
||||
}
|
||||
return cloned
|
||||
case []string:
|
||||
cloned := make([]string, len(typed))
|
||||
copy(cloned, typed)
|
||||
return cloned
|
||||
default:
|
||||
return typed
|
||||
}
|
||||
}
|
||||
|
||||
func responseOutputMessageText(message openairesponses.ResponseOutputMessage) string {
|
||||
parts := make([]string, 0, len(message.Content))
|
||||
for _, content := range message.Content {
|
||||
switch part := content.AsAny().(type) {
|
||||
case openairesponses.ResponseOutputText:
|
||||
if text := strings.TrimSpace(part.Text); text != "" {
|
||||
parts = append(parts, text)
|
||||
}
|
||||
case openairesponses.ResponseOutputRefusal:
|
||||
if refusal := strings.TrimSpace(part.Refusal); refusal != "" {
|
||||
parts = append(parts, refusal)
|
||||
}
|
||||
}
|
||||
}
|
||||
return joinTextBlocks(parts)
|
||||
}
|
||||
|
||||
func responseReasoningSummary(item openairesponses.ResponseReasoningItem) string {
|
||||
summaries := make([]string, 0, len(item.Summary))
|
||||
for _, summary := range item.Summary {
|
||||
if text := strings.TrimSpace(summary.Text); text != "" {
|
||||
summaries = append(summaries, text)
|
||||
}
|
||||
}
|
||||
return joinTextBlocks(summaries)
|
||||
}
|
||||
|
||||
func responseReasoningContent(item openairesponses.ResponseReasoningItem) string {
|
||||
content := make([]string, 0, len(item.Content))
|
||||
for _, block := range item.Content {
|
||||
if text := strings.TrimSpace(block.Text); text != "" {
|
||||
content = append(content, text)
|
||||
}
|
||||
}
|
||||
return joinTextBlocks(content)
|
||||
}
|
||||
|
||||
func agenticToolOutput(ctx context.Context, tools AgenticToolExecutor, toolCall AgenticToolCall) (string, error) {
|
||||
if tools == nil {
|
||||
err := fmt.Errorf("tool executor is not configured")
|
||||
return agenticToolErrorJSON(err), err
|
||||
}
|
||||
|
||||
output, err := tools.CallTool(ctx, toolCall.Name, toolCall.Arguments)
|
||||
if err != nil {
|
||||
return agenticToolErrorJSON(err), err
|
||||
}
|
||||
|
||||
output = strings.TrimSpace(output)
|
||||
if output == "" {
|
||||
empty, marshalErr := json.Marshal(map[string]any{"ok": true})
|
||||
if marshalErr != nil {
|
||||
return `{"ok":true}`, nil
|
||||
}
|
||||
return string(empty), nil
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
func agenticToolErrorJSON(err error) string {
|
||||
payload := map[string]any{
|
||||
"ok": false,
|
||||
"error": "tool call failed",
|
||||
}
|
||||
if err != nil && strings.TrimSpace(err.Error()) != "" {
|
||||
payload["error"] = truncateForError(err.Error())
|
||||
}
|
||||
data, marshalErr := json.Marshal(payload)
|
||||
if marshalErr != nil {
|
||||
return `{"ok":false,"error":"tool call failed"}`
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
)
|
||||
|
||||
type staticAgenticToolExecutor struct {
|
||||
definitions []AgenticToolDefinition
|
||||
}
|
||||
|
||||
func (s staticAgenticToolExecutor) ToolDefinitions() []AgenticToolDefinition {
|
||||
return s.definitions
|
||||
}
|
||||
|
||||
func (staticAgenticToolExecutor) CallTool(_ context.Context, _ string, _ string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func TestBuildAgenticToolParamsNormalizesStrictOpenAISchemas(t *testing.T) {
|
||||
params := buildAgenticToolParams(staticAgenticToolExecutor{
|
||||
definitions: []AgenticToolDefinition{
|
||||
{
|
||||
Name: "fs_ls",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"session_id": map[string]any{"type": "string"},
|
||||
"beacon_id": map[string]any{"type": "string"},
|
||||
"path": map[string]any{"type": "string"},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
if len(params) != 1 || params[0].OfFunction == nil {
|
||||
t.Fatalf("expected one function tool param, got %+v", params)
|
||||
}
|
||||
|
||||
parameters := params[0].OfFunction.Parameters
|
||||
required := schemaRequiredSet(parameters["required"])
|
||||
for _, field := range []string{"session_id", "beacon_id", "path"} {
|
||||
if !required[field] {
|
||||
t.Fatalf("expected strict schema to require %q, got %#v", field, parameters["required"])
|
||||
}
|
||||
}
|
||||
|
||||
properties, ok := parameters["properties"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected strict schema properties map, got %#v", parameters["properties"])
|
||||
}
|
||||
for _, field := range []string{"session_id", "beacon_id", "path"} {
|
||||
property, ok := properties[field].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected property %q to be a schema object, got %#v", field, properties[field])
|
||||
}
|
||||
typeValues, ok := property["type"].([]string)
|
||||
if !ok {
|
||||
t.Fatalf("expected property %q type to be []string, got %#v", field, property["type"])
|
||||
}
|
||||
if !stringSliceContains(typeValues, "string") || !stringSliceContains(typeValues, "null") {
|
||||
t.Fatalf("expected property %q to allow string and null, got %#v", field, typeValues)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationAgenticRequestsReasoningSummaryForOpenAI(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
Body string
|
||||
}
|
||||
|
||||
requests := make(chan capturedRequest, 1)
|
||||
restoreClient := SetHTTPClientForTests(&http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
defer r.Body.Close()
|
||||
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
|
||||
requests <- capturedRequest{Body: string(payload)}
|
||||
return jsonResponse(http.StatusOK, `{
|
||||
"id": "resp_123",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_123",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "output_text", "text": "Done."}
|
||||
]
|
||||
}
|
||||
]
|
||||
}`), nil
|
||||
}),
|
||||
})
|
||||
defer restoreClient()
|
||||
|
||||
completion, err := CompleteConversationAgentic(context.Background(), &RuntimeConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
Model: "gpt-5.4",
|
||||
UseResponsesAPI: true,
|
||||
ThinkingLevel: "high",
|
||||
APIKey: "openai-key",
|
||||
}, &clientpb.AIConversation{
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "Say hi."},
|
||||
},
|
||||
}, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("complete agentic conversation: %v", err)
|
||||
}
|
||||
|
||||
request := <-requests
|
||||
for _, fragment := range []string{
|
||||
`"type":"web_search"`,
|
||||
`"search_context_size":"medium"`,
|
||||
`"effort":"high"`,
|
||||
`"summary":"auto"`,
|
||||
} {
|
||||
if !strings.Contains(request.Body, fragment) {
|
||||
t.Fatalf("expected agentic request body to contain %q, got %s", fragment, request.Body)
|
||||
}
|
||||
}
|
||||
if completion.Content != "Done." {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
}
|
||||
+2
-1
@@ -130,12 +130,13 @@ func SafeConfigSummary() *clientpb.AIConfigSummary {
|
||||
func SafeConfigSummaryFromConfig(cfg *configs.ServerConfig) *clientpb.AIConfigSummary {
|
||||
summary := &clientpb.AIConfigSummary{}
|
||||
if cfg != nil && cfg.AI != nil {
|
||||
summary.Model = strings.TrimSpace(cfg.AI.Model)
|
||||
summary.ThinkingLevel = strings.TrimSpace(cfg.AI.ThinkingLevel)
|
||||
summary.SystemPrompt = strings.TrimSpace(cfg.AI.SystemPrompt)
|
||||
}
|
||||
|
||||
provider, providerConfig := selectedProviderConfig(cfg)
|
||||
summary.Provider = provider
|
||||
summary.Model = configuredDefaultModel(provider, selectedAIConfig(cfg), providerConfig)
|
||||
|
||||
switch {
|
||||
case cfg == nil || cfg.AI == nil:
|
||||
|
||||
+20
-2
@@ -10,9 +10,9 @@ func TestSafeConfigSummaryFromConfigUsesExplicitConfiguredProvider(t *testing.T)
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
Model: "gpt-test",
|
||||
ThinkingLevel: "high",
|
||||
OpenAI: &configs.AIProviderConfig{APIKey: "openai-key"},
|
||||
SystemPrompt: "Stay concise.",
|
||||
OpenAI: &configs.AIProviderConfig{APIKey: "openai-key", Models: []string{"gpt-test", "gpt-test-mini"}},
|
||||
Anthropic: &configs.AIProviderConfig{APIKey: "anthropic-key"},
|
||||
},
|
||||
}
|
||||
@@ -30,6 +30,24 @@ func TestSafeConfigSummaryFromConfigUsesExplicitConfiguredProvider(t *testing.T)
|
||||
if summary.GetThinkingLevel() != "high" {
|
||||
t.Fatalf("expected thinking level %q, got %q", "high", summary.GetThinkingLevel())
|
||||
}
|
||||
if summary.GetSystemPrompt() != "Stay concise." {
|
||||
t.Fatalf("expected system prompt %q, got %q", "Stay concise.", summary.GetSystemPrompt())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSafeConfigSummaryFromConfigFallsBackToLegacySharedModel(t *testing.T) {
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
Model: "legacy-model",
|
||||
OpenAI: &configs.AIProviderConfig{APIKey: "openai-key"},
|
||||
},
|
||||
}
|
||||
|
||||
summary := SafeConfigSummaryFromConfig(cfg)
|
||||
if summary.GetModel() != "legacy-model" {
|
||||
t.Fatalf("expected legacy model %q, got %q", "legacy-model", summary.GetModel())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSafeConfigSummaryFromConfigFallsBackToConfiguredProvider(t *testing.T) {
|
||||
|
||||
+84
-374
@@ -8,14 +8,9 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||
fantasygoogle "charm.land/fantasy/providers/google"
|
||||
fantasyopenai "charm.land/fantasy/providers/openai"
|
||||
fantasyopenaicompat "charm.land/fantasy/providers/openaicompat"
|
||||
fantasyopenrouter "charm.land/fantasy/providers/openrouter"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/server/configs"
|
||||
"github.com/bishopfox/sliver/util/clientpbutil"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -24,8 +19,6 @@ const (
|
||||
defaultGoogleModel = "gemini-2.5-pro"
|
||||
defaultOpenRouterModel = "openai/gpt-5"
|
||||
defaultCompletionTimeout = 2 * time.Minute
|
||||
|
||||
anthropicThinkingBudgetLow = 1024
|
||||
)
|
||||
|
||||
var httpClient = &http.Client{Timeout: defaultCompletionTimeout}
|
||||
@@ -68,11 +61,12 @@ type RuntimeConfig struct {
|
||||
|
||||
// Completion contains the persisted assistant response details.
|
||||
type Completion struct {
|
||||
Provider string
|
||||
Model string
|
||||
Content string
|
||||
ProviderMessageID string
|
||||
FinishReason string
|
||||
Provider string
|
||||
Model string
|
||||
Content string
|
||||
ProviderMessageID string
|
||||
FinishReason string
|
||||
ContextWindowUsage *ContextWindowUsage
|
||||
}
|
||||
|
||||
type providerMessage struct {
|
||||
@@ -98,6 +92,9 @@ func ResolveRuntimeConfig(cfg *configs.ServerConfig, conversation *clientpb.AICo
|
||||
if conversation != nil {
|
||||
runtime.Provider = NormalizeProviderName(conversation.GetProvider())
|
||||
runtime.Model = strings.TrimSpace(conversation.GetModel())
|
||||
if thinkingLevel := NormalizeThinkingLevelValue(conversation.GetThinkingLevel()); thinkingLevel != "" {
|
||||
runtime.ThinkingLevel = thinkingLevel
|
||||
}
|
||||
}
|
||||
|
||||
if runtime.Provider == "" {
|
||||
@@ -113,11 +110,8 @@ func ResolveRuntimeConfig(cfg *configs.ServerConfig, conversation *clientpb.AICo
|
||||
applyProviderRuntimeConfig(runtime, providerConfig)
|
||||
}
|
||||
|
||||
if runtime.Model == "" && cfg != nil && cfg.AI != nil {
|
||||
runtime.Model = strings.TrimSpace(cfg.AI.Model)
|
||||
}
|
||||
if runtime.Model == "" {
|
||||
runtime.Model = defaultModelForProvider(runtime.Provider)
|
||||
runtime.Model = configuredDefaultModel(runtime.Provider, selectedAIConfig(cfg), providerConfig)
|
||||
}
|
||||
runtime.UseResponsesAPI = useResponsesAPI(runtime.Provider, providerConfig)
|
||||
|
||||
@@ -129,7 +123,9 @@ func ResolveRuntimeConfig(cfg *configs.ServerConfig, conversation *clientpb.AICo
|
||||
case !runtimeProviderConfigured(runtime):
|
||||
return runtime, errors.New(missingProviderConfigError(runtime.Provider))
|
||||
case runtime.Model == "":
|
||||
return runtime, fmt.Errorf("server AI provider %q is missing a model; update `ai.model` or choose a provider default", runtime.Provider)
|
||||
return runtime, fmt.Errorf("server AI provider %q is missing a model; update `ai.%s.models` or choose a provider default", runtime.Provider, providerConfigKey(runtime.Provider))
|
||||
case !completionDriverAvailable(runtime):
|
||||
return runtime, missingDriverError(runtime)
|
||||
default:
|
||||
return runtime, nil
|
||||
}
|
||||
@@ -149,229 +145,15 @@ func CompleteConversation(ctx context.Context, runtime *RuntimeConfig, conversat
|
||||
return nil, err
|
||||
}
|
||||
|
||||
provider, err := newFantasyProvider(runtime)
|
||||
driver, err := selectCompletionDriver(runtime)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
model, err := provider.LanguageModel(ctx, runtime.Model)
|
||||
if err != nil {
|
||||
return nil, formatFantasyError(runtime.Provider, err)
|
||||
}
|
||||
|
||||
response, err := model.Generate(ctx, buildFantasyCall(runtime, systemPrompt, messages))
|
||||
if err != nil {
|
||||
return nil, formatFantasyError(runtime.Provider, err)
|
||||
}
|
||||
|
||||
content := extractFantasyResponseText(response)
|
||||
if content == "" {
|
||||
return nil, fmt.Errorf("%s response did not include assistant text", runtime.Provider)
|
||||
}
|
||||
|
||||
finishReason := strings.TrimSpace(string(response.FinishReason))
|
||||
if finishReason == "" || finishReason == string(fantasy.FinishReasonUnknown) {
|
||||
finishReason = "completed"
|
||||
}
|
||||
|
||||
return &Completion{
|
||||
Provider: fallbackString(model.Provider(), runtime.Provider),
|
||||
Model: fallbackString(model.Model(), runtime.Model),
|
||||
Content: content,
|
||||
ProviderMessageID: extractProviderMessageID(response),
|
||||
FinishReason: finishReason,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func newFantasyProvider(runtime *RuntimeConfig) (fantasy.Provider, error) {
|
||||
if runtime == nil {
|
||||
return nil, fmt.Errorf("AI runtime config is required")
|
||||
}
|
||||
|
||||
headers := providerHeaders(runtime)
|
||||
|
||||
switch runtime.Provider {
|
||||
case ProviderAnthropic:
|
||||
opts := []fantasyanthropic.Option{
|
||||
fantasyanthropic.WithName(ProviderAnthropic),
|
||||
fantasyanthropic.WithHTTPClient(httpClient),
|
||||
}
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, fantasyanthropic.WithAPIKey(runtime.APIKey))
|
||||
}
|
||||
if runtime.BaseURL != "" {
|
||||
opts = append(opts, fantasyanthropic.WithBaseURL(runtime.BaseURL))
|
||||
}
|
||||
if len(headers) > 0 {
|
||||
opts = append(opts, fantasyanthropic.WithHeaders(headers))
|
||||
}
|
||||
if runtime.Project != "" && runtime.Location != "" {
|
||||
opts = append(opts, fantasyanthropic.WithVertex(runtime.Project, runtime.Location))
|
||||
}
|
||||
if runtime.SkipAuth {
|
||||
opts = append(opts, fantasyanthropic.WithSkipAuth(true))
|
||||
}
|
||||
if runtime.UseBedrock {
|
||||
opts = append(opts, fantasyanthropic.WithBedrock())
|
||||
}
|
||||
return fantasyanthropic.New(opts...)
|
||||
case ProviderGoogle:
|
||||
opts := []fantasygoogle.Option{
|
||||
fantasygoogle.WithName(ProviderGoogle),
|
||||
fantasygoogle.WithHTTPClient(httpClient),
|
||||
}
|
||||
if runtime.BaseURL != "" {
|
||||
opts = append(opts, fantasygoogle.WithBaseURL(runtime.BaseURL))
|
||||
}
|
||||
if len(headers) > 0 {
|
||||
opts = append(opts, fantasygoogle.WithHeaders(headers))
|
||||
}
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, fantasygoogle.WithGeminiAPIKey(runtime.APIKey))
|
||||
} else if runtime.Project != "" && runtime.Location != "" {
|
||||
opts = append(opts, fantasygoogle.WithVertex(runtime.Project, runtime.Location))
|
||||
}
|
||||
if runtime.SkipAuth {
|
||||
opts = append(opts, fantasygoogle.WithSkipAuth(true))
|
||||
}
|
||||
return fantasygoogle.New(opts...)
|
||||
case ProviderOpenAI:
|
||||
opts := []fantasyopenai.Option{
|
||||
fantasyopenai.WithName(ProviderOpenAI),
|
||||
fantasyopenai.WithHTTPClient(httpClient),
|
||||
}
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, fantasyopenai.WithAPIKey(runtime.APIKey))
|
||||
}
|
||||
if runtime.BaseURL != "" {
|
||||
opts = append(opts, fantasyopenai.WithBaseURL(runtime.BaseURL))
|
||||
}
|
||||
if len(headers) > 0 {
|
||||
opts = append(opts, fantasyopenai.WithHeaders(headers))
|
||||
}
|
||||
if runtime.Organization != "" {
|
||||
opts = append(opts, fantasyopenai.WithOrganization(runtime.Organization))
|
||||
}
|
||||
if runtime.Project != "" {
|
||||
opts = append(opts, fantasyopenai.WithProject(runtime.Project))
|
||||
}
|
||||
if runtime.UseResponsesAPI {
|
||||
opts = append(opts, fantasyopenai.WithUseResponsesAPI())
|
||||
}
|
||||
return fantasyopenai.New(opts...)
|
||||
case ProviderOpenAICompat:
|
||||
opts := []fantasyopenaicompat.Option{
|
||||
fantasyopenaicompat.WithName(ProviderOpenAICompat),
|
||||
fantasyopenaicompat.WithHTTPClient(httpClient),
|
||||
}
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, fantasyopenaicompat.WithAPIKey(runtime.APIKey))
|
||||
}
|
||||
if runtime.BaseURL != "" {
|
||||
opts = append(opts, fantasyopenaicompat.WithBaseURL(runtime.BaseURL))
|
||||
}
|
||||
if len(headers) > 0 {
|
||||
opts = append(opts, fantasyopenaicompat.WithHeaders(headers))
|
||||
}
|
||||
if runtime.UseResponsesAPI {
|
||||
opts = append(opts, fantasyopenaicompat.WithUseResponsesAPI())
|
||||
}
|
||||
return fantasyopenaicompat.New(opts...)
|
||||
case ProviderOpenRouter:
|
||||
opts := []fantasyopenrouter.Option{
|
||||
fantasyopenrouter.WithName(ProviderOpenRouter),
|
||||
fantasyopenrouter.WithHTTPClient(httpClient),
|
||||
}
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, fantasyopenrouter.WithAPIKey(runtime.APIKey))
|
||||
}
|
||||
if len(headers) > 0 {
|
||||
opts = append(opts, fantasyopenrouter.WithHeaders(headers))
|
||||
}
|
||||
return fantasyopenrouter.New(opts...)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported AI provider %q", runtime.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func buildFantasyCall(runtime *RuntimeConfig, systemPrompt string, messages []providerMessage) fantasy.Call {
|
||||
prompt := make(fantasy.Prompt, 0, len(messages)+1)
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
prompt = append(prompt, fantasy.Message{
|
||||
Role: fantasy.MessageRoleSystem,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: strings.TrimSpace(systemPrompt)},
|
||||
},
|
||||
})
|
||||
}
|
||||
for _, message := range messages {
|
||||
prompt = append(prompt, fantasy.Message{
|
||||
Role: toFantasyMessageRole(message.Role),
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: message.Content},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
call := fantasy.Call{
|
||||
Prompt: prompt,
|
||||
ProviderOptions: fantasyProviderOptions(runtime),
|
||||
MaxOutputTokens: optionalRuntimeInt(runtime.MaxOutputTokens),
|
||||
Temperature: copyOptionalFloat(runtime.Temperature),
|
||||
TopP: copyOptionalFloat(runtime.TopP),
|
||||
TopK: copyOptionalInt(runtime.TopK),
|
||||
PresencePenalty: copyOptionalFloat(runtime.PresencePenalty),
|
||||
FrequencyPenalty: copyOptionalFloat(runtime.FrequencyPenalty),
|
||||
}
|
||||
return call
|
||||
}
|
||||
|
||||
func fantasyProviderOptions(runtime *RuntimeConfig) fantasy.ProviderOptions {
|
||||
if runtime == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch runtime.Provider {
|
||||
case ProviderAnthropic:
|
||||
if budget := anthropicThinkingBudget(runtime.ThinkingLevel); budget > 0 {
|
||||
return fantasyanthropic.NewProviderOptions(&fantasyanthropic.ProviderOptions{
|
||||
Thinking: &fantasyanthropic.ThinkingProviderOption{BudgetTokens: budget},
|
||||
})
|
||||
}
|
||||
case ProviderGoogle:
|
||||
if thinkingConfig := googleThinkingConfig(runtime.ThinkingLevel); thinkingConfig != nil {
|
||||
return fantasy.ProviderOptions{
|
||||
fantasygoogle.Name: &fantasygoogle.ProviderOptions{
|
||||
ThinkingConfig: thinkingConfig,
|
||||
},
|
||||
}
|
||||
}
|
||||
case ProviderOpenAI:
|
||||
if effort := openAIReasoningEffort(runtime.ThinkingLevel); effort != nil {
|
||||
if runtime.UseResponsesAPI && fantasyopenai.IsResponsesModel(runtime.Model) {
|
||||
return fantasyopenai.NewResponsesProviderOptions(&fantasyopenai.ResponsesProviderOptions{
|
||||
ReasoningEffort: effort,
|
||||
})
|
||||
}
|
||||
return fantasyopenai.NewProviderOptions(&fantasyopenai.ProviderOptions{
|
||||
ReasoningEffort: effort,
|
||||
})
|
||||
}
|
||||
case ProviderOpenAICompat:
|
||||
if effort := openAIReasoningEffort(runtime.ThinkingLevel); effort != nil {
|
||||
return fantasyopenaicompat.NewProviderOptions(&fantasyopenaicompat.ProviderOptions{
|
||||
ReasoningEffort: effort,
|
||||
})
|
||||
}
|
||||
case ProviderOpenRouter:
|
||||
if reasoning := openRouterReasoningOptions(runtime.ThinkingLevel); reasoning != nil {
|
||||
return fantasyopenrouter.NewProviderOptions(&fantasyopenrouter.ProviderOptions{
|
||||
Reasoning: reasoning,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
return driver.CompleteConversation(ctx, runtime, &completionRequest{
|
||||
SystemPrompt: systemPrompt,
|
||||
Messages: messages,
|
||||
})
|
||||
}
|
||||
|
||||
func conversationHistory(conversation *clientpb.AIConversation) (string, []providerMessage, error) {
|
||||
@@ -385,6 +167,12 @@ func conversationHistory(conversation *clientpb.AIConversation) (string, []provi
|
||||
if message == nil {
|
||||
continue
|
||||
}
|
||||
if !clientpbutil.AIConversationMessageIncludesContext(message) {
|
||||
continue
|
||||
}
|
||||
if message.GetKind() != clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT {
|
||||
continue
|
||||
}
|
||||
|
||||
role := strings.ToLower(strings.TrimSpace(message.GetRole()))
|
||||
content := strings.TrimSpace(message.GetContent())
|
||||
@@ -432,82 +220,69 @@ func defaultModelForProvider(provider string) string {
|
||||
}
|
||||
}
|
||||
|
||||
func configuredDefaultModel(provider string, aiConfig *configs.AIConfig, providerConfig *configs.AIProviderConfig) string {
|
||||
if providerConfig != nil {
|
||||
for _, model := range providerConfig.Models {
|
||||
model = strings.TrimSpace(model)
|
||||
if model != "" {
|
||||
return model
|
||||
}
|
||||
}
|
||||
}
|
||||
if aiConfig != nil {
|
||||
if legacyModel := strings.TrimSpace(aiConfig.Model); legacyModel != "" {
|
||||
return legacyModel
|
||||
}
|
||||
}
|
||||
return defaultModelForProvider(provider)
|
||||
}
|
||||
|
||||
func providerConfigKey(provider string) string {
|
||||
switch NormalizeProviderName(provider) {
|
||||
case ProviderOpenAICompat:
|
||||
return "openai_compat"
|
||||
default:
|
||||
return NormalizeProviderName(provider)
|
||||
}
|
||||
}
|
||||
|
||||
func selectedAIConfig(cfg *configs.ServerConfig) *configs.AIConfig {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
return cfg.AI
|
||||
}
|
||||
|
||||
func completionDriverAvailable(runtime *RuntimeConfig) bool {
|
||||
_, err := selectCompletionDriver(runtime)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func missingDriverError(runtime *RuntimeConfig) error {
|
||||
if runtime == nil {
|
||||
return fmt.Errorf("AI runtime config is required")
|
||||
}
|
||||
_, err := selectCompletionDriver(runtime)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func NormalizeThinkingLevelValue(value string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "low", "medium", "high", "xhigh", "disabled":
|
||||
return strings.ToLower(strings.TrimSpace(value))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeThinkingLevel(cfg *configs.ServerConfig) string {
|
||||
if cfg == nil || cfg.AI == nil {
|
||||
return ""
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(cfg.AI.ThinkingLevel)) {
|
||||
case "low", "medium", "high", "disabled":
|
||||
return strings.ToLower(strings.TrimSpace(cfg.AI.ThinkingLevel))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func openAIReasoningEffort(thinkingLevel string) *fantasyopenai.ReasoningEffort {
|
||||
switch strings.ToLower(strings.TrimSpace(thinkingLevel)) {
|
||||
case "low":
|
||||
return fantasyopenai.ReasoningEffortOption(fantasyopenai.ReasoningEffortLow)
|
||||
case "medium":
|
||||
return fantasyopenai.ReasoningEffortOption(fantasyopenai.ReasoningEffortMedium)
|
||||
case "high":
|
||||
return fantasyopenai.ReasoningEffortOption(fantasyopenai.ReasoningEffortHigh)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func anthropicThinkingBudget(thinkingLevel string) int64 {
|
||||
switch strings.ToLower(strings.TrimSpace(thinkingLevel)) {
|
||||
case "low":
|
||||
return anthropicThinkingBudgetLow
|
||||
case "medium":
|
||||
return anthropicThinkingBudgetLow * 2
|
||||
case "high":
|
||||
return anthropicThinkingBudgetLow * 4
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func googleThinkingConfig(thinkingLevel string) *fantasygoogle.ThinkingConfig {
|
||||
var budget int64
|
||||
switch strings.ToLower(strings.TrimSpace(thinkingLevel)) {
|
||||
case "low":
|
||||
budget = 1024
|
||||
case "medium":
|
||||
budget = 2048
|
||||
case "high":
|
||||
budget = 4096
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
|
||||
return &fantasygoogle.ThinkingConfig{
|
||||
ThinkingBudget: fantasy.Opt(budget),
|
||||
}
|
||||
}
|
||||
|
||||
func openRouterReasoningOptions(thinkingLevel string) *fantasyopenrouter.ReasoningOptions {
|
||||
switch strings.ToLower(strings.TrimSpace(thinkingLevel)) {
|
||||
case "low":
|
||||
return &fantasyopenrouter.ReasoningOptions{
|
||||
Enabled: fantasy.Opt(true),
|
||||
Effort: fantasyopenrouter.ReasoningEffortOption(fantasyopenrouter.ReasoningEffortLow),
|
||||
}
|
||||
case "medium":
|
||||
return &fantasyopenrouter.ReasoningOptions{
|
||||
Enabled: fantasy.Opt(true),
|
||||
Effort: fantasyopenrouter.ReasoningEffortOption(fantasyopenrouter.ReasoningEffortMedium),
|
||||
}
|
||||
case "high":
|
||||
return &fantasyopenrouter.ReasoningOptions{
|
||||
Enabled: fantasy.Opt(true),
|
||||
Effort: fantasyopenrouter.ReasoningEffortOption(fantasyopenrouter.ReasoningEffortHigh),
|
||||
}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
return NormalizeThinkingLevelValue(cfg.AI.ThinkingLevel)
|
||||
}
|
||||
|
||||
func applyProviderRuntimeConfig(runtime *RuntimeConfig, providerConfig *configs.AIProviderConfig) {
|
||||
@@ -583,71 +358,6 @@ func runtimeProviderConfigured(runtime *RuntimeConfig) bool {
|
||||
}
|
||||
}
|
||||
|
||||
func extractFantasyResponseText(response *fantasy.Response) string {
|
||||
if response == nil {
|
||||
return ""
|
||||
}
|
||||
blocks := make([]string, 0, len(response.Content))
|
||||
for _, content := range response.Content {
|
||||
if content == nil || content.GetType() != fantasy.ContentTypeText {
|
||||
continue
|
||||
}
|
||||
textContent, ok := fantasy.AsContentType[fantasy.TextContent](content)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(textContent.Text) == "" {
|
||||
continue
|
||||
}
|
||||
blocks = append(blocks, textContent.Text)
|
||||
}
|
||||
return joinTextBlocks(blocks)
|
||||
}
|
||||
|
||||
func extractProviderMessageID(_ *fantasy.Response) string {
|
||||
// Fantasy's unified response model does not currently expose a stable top-level
|
||||
// provider message ID for plain text generations.
|
||||
return ""
|
||||
}
|
||||
|
||||
func formatFantasyError(provider string, err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var providerErr *fantasy.ProviderError
|
||||
if errors.As(err, &providerErr) {
|
||||
message := strings.TrimSpace(providerErr.Error())
|
||||
if providerErr.StatusCode > 0 {
|
||||
if message == "" {
|
||||
message = http.StatusText(providerErr.StatusCode)
|
||||
}
|
||||
return fmt.Errorf("%s API request failed with HTTP %d: %s", provider, providerErr.StatusCode, truncateForError(message))
|
||||
}
|
||||
if message != "" {
|
||||
return fmt.Errorf("%s API request failed: %s", provider, truncateForError(message))
|
||||
}
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func toFantasyMessageRole(role string) fantasy.MessageRole {
|
||||
switch strings.ToLower(strings.TrimSpace(role)) {
|
||||
case "assistant":
|
||||
return fantasy.MessageRoleAssistant
|
||||
default:
|
||||
return fantasy.MessageRoleUser
|
||||
}
|
||||
}
|
||||
|
||||
func optionalRuntimeInt(value int64) *int64 {
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return fantasy.Opt(value)
|
||||
}
|
||||
|
||||
func copyOptionalFloat(value *float64) *float64 {
|
||||
if value == nil {
|
||||
return nil
|
||||
|
||||
+478
-65
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/server/configs"
|
||||
"github.com/openai/openai-go/v2/shared"
|
||||
)
|
||||
|
||||
func TestCompleteConversationOpenAIUsesConfiguredCredentialsAndSettings(t *testing.T) {
|
||||
@@ -38,6 +39,13 @@ func TestCompleteConversationOpenAIUsesConfiguredCredentialsAndSettings(t *testi
|
||||
"id": "resp_123",
|
||||
"model": "gpt-5.2",
|
||||
"status": "completed",
|
||||
"usage": {
|
||||
"input_tokens": 12000,
|
||||
"input_tokens_details": {},
|
||||
"output_tokens": 6000,
|
||||
"output_tokens_details": {},
|
||||
"total_tokens": 18000
|
||||
},
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
@@ -95,7 +103,10 @@ func TestCompleteConversationOpenAIUsesConfiguredCredentialsAndSettings(t *testi
|
||||
}
|
||||
for _, fragment := range []string{
|
||||
`"model":"gpt-5.2"`,
|
||||
`"type":"web_search"`,
|
||||
`"search_context_size":"medium"`,
|
||||
`"effort":"high"`,
|
||||
`"summary":"auto"`,
|
||||
`"Keep it brief."`,
|
||||
`"Earlier answer."`,
|
||||
`"Explain the workflow."`,
|
||||
@@ -108,17 +119,25 @@ func TestCompleteConversationOpenAIUsesConfiguredCredentialsAndSettings(t *testi
|
||||
if completion.Content != "OpenAI assistant reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.ProviderMessageID != "" {
|
||||
t.Fatalf("expected fantasy-backed completion to leave provider message id empty, got %q", completion.ProviderMessageID)
|
||||
if completion.ProviderMessageID != "resp_123" {
|
||||
t.Fatalf("unexpected provider message id: %q", completion.ProviderMessageID)
|
||||
}
|
||||
if completion.ContextWindowUsage == nil {
|
||||
t.Fatal("expected context window usage to be captured")
|
||||
}
|
||||
if completion.ContextWindowUsage.TotalTokens != 18000 {
|
||||
t.Fatalf("unexpected total tokens: %d", completion.ContextWindowUsage.TotalTokens)
|
||||
}
|
||||
if completion.ContextWindowUsage.ContextWindowTokens != openAIFrontierContextWindowTokens {
|
||||
t.Fatalf("unexpected context window tokens: %d", completion.ContextWindowUsage.ContextWindowTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationAnthropicUsesConfiguredCredentialsAndSettings(t *testing.T) {
|
||||
func TestCompleteConversationOpenAIUsesDefaultBaseURLWhenUnset(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
Path string
|
||||
APIKey string
|
||||
AnthropicVersion string
|
||||
Body string
|
||||
Path string
|
||||
Authorization string
|
||||
Body string
|
||||
}
|
||||
|
||||
requests := make(chan capturedRequest, 1)
|
||||
@@ -132,25 +151,192 @@ func TestCompleteConversationAnthropicUsesConfiguredCredentialsAndSettings(t *te
|
||||
}
|
||||
|
||||
requests <- capturedRequest{
|
||||
Path: r.URL.Path,
|
||||
APIKey: r.Header.Get("x-api-key"),
|
||||
AnthropicVersion: r.Header.Get("anthropic-version"),
|
||||
Body: string(payload),
|
||||
Path: r.URL.Path,
|
||||
Authorization: r.Header.Get("Authorization"),
|
||||
Body: string(payload),
|
||||
}
|
||||
|
||||
return jsonResponse(http.StatusOK, `{
|
||||
"id": "msg_123",
|
||||
"model": "claude-sonnet-4-0",
|
||||
"stop_reason": "end_turn",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "internal reasoning"},
|
||||
{"type": "text", "text": "Anthropic assistant reply"}
|
||||
"id": "resp_default_base",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "output_text", "text": "OpenAI default-base reply"}
|
||||
]
|
||||
}
|
||||
]
|
||||
}`), nil
|
||||
}),
|
||||
})
|
||||
defer restoreClient()
|
||||
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
OpenAI: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-5.4"},
|
||||
APIKey: "openai-key",
|
||||
UseResponsesAPI: boolPtr(true),
|
||||
},
|
||||
Anthropic: &configs.AIProviderConfig{},
|
||||
},
|
||||
}
|
||||
conversation := &clientpb.AIConversation{
|
||||
Provider: ProviderOpenAI,
|
||||
Model: "gpt-5.4",
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "Say hi."},
|
||||
},
|
||||
}
|
||||
|
||||
runtime, err := ResolveRuntimeConfig(cfg, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve runtime config: %v", err)
|
||||
}
|
||||
if runtime.BaseURL != "" {
|
||||
t.Fatalf("expected openai runtime base url to remain unset, got %q", runtime.BaseURL)
|
||||
}
|
||||
|
||||
completion, err := CompleteConversation(context.Background(), runtime, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("complete conversation: %v", err)
|
||||
}
|
||||
|
||||
request := <-requests
|
||||
if request.Path != "/v1/responses" {
|
||||
t.Fatalf("unexpected openai default-base request path: got=%q want=%q", request.Path, "/v1/responses")
|
||||
}
|
||||
if request.Authorization != "Bearer openai-key" {
|
||||
t.Fatalf("unexpected authorization header: %q", request.Authorization)
|
||||
}
|
||||
if !strings.Contains(request.Body, `"Say hi."`) {
|
||||
t.Fatalf("expected openai request body to contain the user prompt, got %s", request.Body)
|
||||
}
|
||||
if completion.Content != "OpenAI default-base reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.ProviderMessageID != "resp_default_base" {
|
||||
t.Fatalf("unexpected provider message id: %q", completion.ProviderMessageID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationOpenAIChatEnablesWebSearch(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
Path string
|
||||
Authorization string
|
||||
Body string
|
||||
}
|
||||
|
||||
requests := make(chan capturedRequest, 1)
|
||||
restoreClient := SetHTTPClientForTests(&http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
defer r.Body.Close()
|
||||
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
|
||||
requests <- capturedRequest{
|
||||
Path: r.URL.Path,
|
||||
Authorization: r.Header.Get("Authorization"),
|
||||
Body: string(payload),
|
||||
}
|
||||
|
||||
return jsonResponse(http.StatusOK, `{
|
||||
"id": "chatcmpl_openai_123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5.4",
|
||||
"usage": {
|
||||
"prompt_tokens": 32000,
|
||||
"completion_tokens": 4000,
|
||||
"total_tokens": 36000
|
||||
},
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "OpenAI chat reply"
|
||||
}
|
||||
}
|
||||
]
|
||||
}`), nil
|
||||
}),
|
||||
})
|
||||
defer restoreClient()
|
||||
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
OpenAI: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-5.4"},
|
||||
APIKey: "openai-key",
|
||||
UseResponsesAPI: boolPtr(false),
|
||||
},
|
||||
Anthropic: &configs.AIProviderConfig{},
|
||||
},
|
||||
}
|
||||
conversation := &clientpb.AIConversation{
|
||||
Provider: ProviderOpenAI,
|
||||
Model: "gpt-5.4",
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "Say hi."},
|
||||
},
|
||||
}
|
||||
|
||||
runtime, err := ResolveRuntimeConfig(cfg, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve runtime config: %v", err)
|
||||
}
|
||||
if runtime.UseResponsesAPI {
|
||||
t.Fatal("expected openai runtime to use chat completions for this test")
|
||||
}
|
||||
|
||||
completion, err := CompleteConversation(context.Background(), runtime, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("complete conversation: %v", err)
|
||||
}
|
||||
|
||||
request := <-requests
|
||||
if request.Path != "/v1/chat/completions" {
|
||||
t.Fatalf("unexpected openai chat request path: got=%q want=%q", request.Path, "/v1/chat/completions")
|
||||
}
|
||||
if request.Authorization != "Bearer openai-key" {
|
||||
t.Fatalf("unexpected authorization header: %q", request.Authorization)
|
||||
}
|
||||
for _, fragment := range []string{
|
||||
`"model":"gpt-5.4"`,
|
||||
`"web_search_options":{"search_context_size":"medium"}`,
|
||||
`"Say hi."`,
|
||||
} {
|
||||
if !strings.Contains(request.Body, fragment) {
|
||||
t.Fatalf("expected openai chat request body to contain %q, got %s", fragment, request.Body)
|
||||
}
|
||||
}
|
||||
if completion.Content != "OpenAI chat reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.ProviderMessageID != "chatcmpl_openai_123" {
|
||||
t.Fatalf("unexpected provider message id: %q", completion.ProviderMessageID)
|
||||
}
|
||||
if completion.ContextWindowUsage == nil {
|
||||
t.Fatal("expected openai chat usage to be captured")
|
||||
}
|
||||
if completion.ContextWindowUsage.TotalTokens != 36000 {
|
||||
t.Fatalf("unexpected openai chat total tokens: %d", completion.ContextWindowUsage.TotalTokens)
|
||||
}
|
||||
if completion.ContextWindowUsage.ContextWindowTokens != openAI5_4ContextWindowTokens {
|
||||
t.Fatalf("unexpected openai chat context window tokens: %d", completion.ContextWindowUsage.ContextWindowTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRuntimeConfigAnthropicRequiresAnInstalledDriver(t *testing.T) {
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderAnthropic,
|
||||
@@ -162,57 +348,16 @@ func TestCompleteConversationAnthropicUsesConfiguredCredentialsAndSettings(t *te
|
||||
OpenAI: &configs.AIProviderConfig{},
|
||||
},
|
||||
}
|
||||
conversation := &clientpb.AIConversation{
|
||||
Provider: ProviderAnthropic,
|
||||
SystemPrompt: "Use short answers.",
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "Hello"},
|
||||
{Role: "assistant", Content: "Hi there"},
|
||||
{Role: "user", Content: "What changed?"},
|
||||
},
|
||||
}
|
||||
|
||||
runtime, err := ResolveRuntimeConfig(cfg, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve runtime config: %v", err)
|
||||
}
|
||||
runtime, err := ResolveRuntimeConfig(cfg, &clientpb.AIConversation{Provider: ProviderAnthropic})
|
||||
if runtime.Model != defaultAnthropicModel {
|
||||
t.Fatalf("unexpected default anthropic model: got=%q want=%q", runtime.Model, defaultAnthropicModel)
|
||||
}
|
||||
|
||||
completion, err := CompleteConversation(context.Background(), runtime, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("complete conversation: %v", err)
|
||||
if err == nil {
|
||||
t.Fatal("expected anthropic runtime resolution to require a dedicated driver")
|
||||
}
|
||||
|
||||
request := <-requests
|
||||
if request.Path != "/edge/v1/messages" {
|
||||
t.Fatalf("unexpected anthropic request path: got=%q want=%q", request.Path, "/edge/v1/messages")
|
||||
}
|
||||
if request.APIKey != "anthropic-key" {
|
||||
t.Fatalf("unexpected anthropic api key header: %q", request.APIKey)
|
||||
}
|
||||
if request.AnthropicVersion == "" {
|
||||
t.Fatal("expected anthropic version header to be set")
|
||||
}
|
||||
for _, fragment := range []string{
|
||||
`"model":"claude-sonnet-4-0"`,
|
||||
`"Use short answers."`,
|
||||
`"budget_tokens":2048`,
|
||||
`"Hello"`,
|
||||
`"Hi there"`,
|
||||
`"What changed?"`,
|
||||
} {
|
||||
if !strings.Contains(request.Body, fragment) {
|
||||
t.Fatalf("expected anthropic request body to contain %q, got %s", fragment, request.Body)
|
||||
}
|
||||
}
|
||||
|
||||
if completion.Content != "Anthropic assistant reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.FinishReason != "stop" {
|
||||
t.Fatalf("unexpected finish reason: %q", completion.FinishReason)
|
||||
if !strings.Contains(err.Error(), `provider "anthropic"`) || !strings.Contains(err.Error(), "SDK driver") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -220,8 +365,8 @@ func TestResolveRuntimeConfigOpenAICompatAllowsBaseURLWithoutAPIKey(t *testing.T
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAICompat,
|
||||
Model: "gpt-oss-120b",
|
||||
OpenAICompat: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-oss-120b"},
|
||||
BaseURL: "http://127.0.0.1:8080/v1",
|
||||
},
|
||||
},
|
||||
@@ -244,6 +389,58 @@ func TestResolveRuntimeConfigOpenAICompatAllowsBaseURLWithoutAPIKey(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRuntimeConfigUsesConversationThinkingOverride(t *testing.T) {
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAI,
|
||||
ThinkingLevel: "high",
|
||||
OpenAI: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-5.4"},
|
||||
APIKey: "openai-key",
|
||||
UseResponsesAPI: boolPtr(true),
|
||||
},
|
||||
Anthropic: &configs.AIProviderConfig{},
|
||||
},
|
||||
}
|
||||
|
||||
runtime, err := ResolveRuntimeConfig(cfg, &clientpb.AIConversation{
|
||||
Provider: ProviderOpenAI,
|
||||
ThinkingLevel: "disabled",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("resolve runtime config: %v", err)
|
||||
}
|
||||
if runtime.ThinkingLevel != "disabled" {
|
||||
t.Fatalf("expected conversation thinking override %q, got %q", "disabled", runtime.ThinkingLevel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReasoningEffortForThinkingLevelSupportsXHigh(t *testing.T) {
|
||||
effort, ok := reasoningEffortForThinkingLevel("xhigh")
|
||||
if !ok {
|
||||
t.Fatal("expected xhigh reasoning effort to be supported")
|
||||
}
|
||||
if string(effort) != "xhigh" {
|
||||
t.Fatalf("expected reasoning effort %q, got %q", "xhigh", effort)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseReasoningParamSkipsSummaryForOpenAICompat(t *testing.T) {
|
||||
reasoning, ok := responseReasoningParam(&RuntimeConfig{
|
||||
Provider: ProviderOpenAICompat,
|
||||
ThinkingLevel: "high",
|
||||
})
|
||||
if !ok {
|
||||
t.Fatal("expected response reasoning params to be available")
|
||||
}
|
||||
if reasoning.Effort != shared.ReasoningEffortHigh {
|
||||
t.Fatalf("expected reasoning effort %q, got %q", shared.ReasoningEffortHigh, reasoning.Effort)
|
||||
}
|
||||
if reasoning.Summary != "" {
|
||||
t.Fatalf("expected openai-compatible runtime to skip reasoning summary, got %q", reasoning.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationOpenAICompatUsesBaseURLWithoutAuth(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
Path string
|
||||
@@ -294,8 +491,8 @@ func TestCompleteConversationOpenAICompatUsesBaseURLWithoutAuth(t *testing.T) {
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenAICompat,
|
||||
Model: "gpt-oss-120b",
|
||||
OpenAICompat: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-oss-120b"},
|
||||
BaseURL: "http://127.0.0.1:8080/v1",
|
||||
},
|
||||
},
|
||||
@@ -330,6 +527,222 @@ func TestCompleteConversationOpenAICompatUsesBaseURLWithoutAuth(t *testing.T) {
|
||||
if completion.Content != "OpenAI-compatible assistant reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.ProviderMessageID != "chatcmpl_123" {
|
||||
t.Fatalf("unexpected provider message id: %q", completion.ProviderMessageID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationOpenRouterUsesDefaultBaseURL(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
Path string
|
||||
Authorization string
|
||||
Body string
|
||||
}
|
||||
|
||||
requests := make(chan capturedRequest, 1)
|
||||
restoreClient := SetHTTPClientForTests(&http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
defer r.Body.Close()
|
||||
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
|
||||
requests <- capturedRequest{
|
||||
Path: r.URL.Path,
|
||||
Authorization: r.Header.Get("Authorization"),
|
||||
Body: string(payload),
|
||||
}
|
||||
|
||||
return jsonResponse(http.StatusOK, `{
|
||||
"id": "chatcmpl_or_123",
|
||||
"object": "chat.completion",
|
||||
"model": "openai/gpt-5",
|
||||
"usage": {
|
||||
"prompt_tokens": 14000,
|
||||
"completion_tokens": 2000,
|
||||
"total_tokens": 16000
|
||||
},
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "OpenRouter assistant reply"
|
||||
}
|
||||
}
|
||||
]
|
||||
}`), nil
|
||||
}),
|
||||
})
|
||||
defer restoreClient()
|
||||
|
||||
cfg := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ProviderOpenRouter,
|
||||
ThinkingLevel: "medium",
|
||||
OpenRouter: &configs.AIProviderConfig{
|
||||
APIKey: "openrouter-key",
|
||||
},
|
||||
},
|
||||
}
|
||||
conversation := &clientpb.AIConversation{
|
||||
Provider: ProviderOpenRouter,
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{Role: "user", Content: "Say hi."},
|
||||
},
|
||||
}
|
||||
|
||||
runtime, err := ResolveRuntimeConfig(cfg, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve runtime config: %v", err)
|
||||
}
|
||||
if runtime.Model != defaultOpenRouterModel {
|
||||
t.Fatalf("unexpected default openrouter model: got=%q want=%q", runtime.Model, defaultOpenRouterModel)
|
||||
}
|
||||
|
||||
completion, err := CompleteConversation(context.Background(), runtime, conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("complete conversation: %v", err)
|
||||
}
|
||||
|
||||
request := <-requests
|
||||
if request.Path != "/api/v1/chat/completions" {
|
||||
t.Fatalf("unexpected openrouter request path: got=%q want=%q", request.Path, "/api/v1/chat/completions")
|
||||
}
|
||||
if request.Authorization != "Bearer openrouter-key" {
|
||||
t.Fatalf("unexpected authorization header: %q", request.Authorization)
|
||||
}
|
||||
for _, fragment := range []string{
|
||||
`"model":"openai/gpt-5"`,
|
||||
`"Say hi."`,
|
||||
} {
|
||||
if !strings.Contains(request.Body, fragment) {
|
||||
t.Fatalf("expected openrouter request body to contain %q, got %s", fragment, request.Body)
|
||||
}
|
||||
}
|
||||
if completion.Content != "OpenRouter assistant reply" {
|
||||
t.Fatalf("unexpected completion content: %q", completion.Content)
|
||||
}
|
||||
if completion.ProviderMessageID != "chatcmpl_or_123" {
|
||||
t.Fatalf("unexpected provider message id: %q", completion.ProviderMessageID)
|
||||
}
|
||||
if completion.ContextWindowUsage == nil {
|
||||
t.Fatal("expected openrouter usage to be captured")
|
||||
}
|
||||
if completion.ContextWindowUsage.TotalTokens != 16000 {
|
||||
t.Fatalf("unexpected openrouter total tokens: %d", completion.ContextWindowUsage.TotalTokens)
|
||||
}
|
||||
if completion.ContextWindowUsage.ContextWindowTokens != openAIFrontierContextWindowTokens {
|
||||
t.Fatalf("unexpected openrouter context window tokens: %d", completion.ContextWindowUsage.ContextWindowTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConversationHistoryUsesExplicitContextFlagIndependentlyOfVisibility(t *testing.T) {
|
||||
systemPrompt, messages, err := conversationHistory(&clientpb.AIConversation{
|
||||
SystemPrompt: "Stay concise.",
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: "Visible in the transcript but excluded from context.",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: boolPtr(false),
|
||||
},
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: "UI-only but still included in context.",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY,
|
||||
IncludeInContext: boolPtr(true),
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "What changed?",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: boolPtr(true),
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("conversationHistory: %v", err)
|
||||
}
|
||||
if systemPrompt != "Stay concise." {
|
||||
t.Fatalf("unexpected system prompt: %q", systemPrompt)
|
||||
}
|
||||
if len(messages) != 2 {
|
||||
t.Fatalf("unexpected message count: got=%d want=%d", len(messages), 2)
|
||||
}
|
||||
if messages[0].Content != "UI-only but still included in context." {
|
||||
t.Fatalf("unexpected retained assistant message: %+v", messages[0])
|
||||
}
|
||||
if messages[1].Role != "user" || messages[1].Content != "What changed?" {
|
||||
t.Fatalf("unexpected trailing user message: %+v", messages[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestConversationHistoryTreatsLeadingSystemMessageAsSystemPrompt(t *testing.T) {
|
||||
systemPrompt, messages, err := conversationHistory(&clientpb.AIConversation{
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{
|
||||
Role: "system",
|
||||
Content: "Stay concise.",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: boolPtr(true),
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "What changed?",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: boolPtr(true),
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("conversationHistory: %v", err)
|
||||
}
|
||||
if systemPrompt != "Stay concise." {
|
||||
t.Fatalf("unexpected system prompt: %q", systemPrompt)
|
||||
}
|
||||
if len(messages) != 1 {
|
||||
t.Fatalf("unexpected message count: got=%d want=%d", len(messages), 1)
|
||||
}
|
||||
if messages[0].Role != "user" || messages[0].Content != "What changed?" {
|
||||
t.Fatalf("unexpected trailing user message: %+v", messages[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestConversationHistoryFallsBackToVisibilityWhenContextFlagIsUnset(t *testing.T) {
|
||||
_, messages, err := conversationHistory(&clientpb.AIConversation{
|
||||
Messages: []*clientpb.AIConversationMessage{
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: "Legacy context-visible message.",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "Legacy prompt.",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("conversationHistory: %v", err)
|
||||
}
|
||||
if len(messages) != 2 {
|
||||
t.Fatalf("unexpected legacy message count: got=%d want=%d", len(messages), 2)
|
||||
}
|
||||
if messages[0].Content != "Legacy context-visible message." {
|
||||
t.Fatalf("unexpected legacy assistant message: %+v", messages[0])
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
@@ -0,0 +1,312 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIFrontierContextWindowTokens = 400000
|
||||
openAIChatContextWindowTokens = 128000
|
||||
openAI4_1ContextWindowTokens = 1047576
|
||||
openAI5_4ContextWindowTokens = 1050000
|
||||
openAIReasoningContextWindowTokens = 200000
|
||||
openAIRealtimeContextWindowTokens = 32000
|
||||
openAIRealtimeMiniContextWindowTokens = 16000
|
||||
openRouterContextWindowCacheTTL = 6 * time.Hour
|
||||
)
|
||||
|
||||
// ContextWindowUsage captures the provider-reported token usage for a turn and,
|
||||
// when available, the model context window used to render the client-side meter.
|
||||
type ContextWindowUsage struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
ContextWindowTokens int64
|
||||
ContextWindowTokensEstimated bool
|
||||
}
|
||||
|
||||
func newContextWindowUsage(inputTokens, outputTokens, totalTokens, windowTokens int64, estimated bool) *ContextWindowUsage {
|
||||
if inputTokens == 0 && outputTokens == 0 && totalTokens == 0 {
|
||||
return nil
|
||||
}
|
||||
return &ContextWindowUsage{
|
||||
InputTokens: inputTokens,
|
||||
OutputTokens: outputTokens,
|
||||
TotalTokens: totalTokens,
|
||||
ContextWindowTokens: windowTokens,
|
||||
ContextWindowTokensEstimated: estimated,
|
||||
}
|
||||
}
|
||||
|
||||
func cloneContextWindowUsage(usage *ContextWindowUsage) *ContextWindowUsage {
|
||||
if usage == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *usage
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func contextWindowUsageRank(usage *ContextWindowUsage) (int64, int64, int64) {
|
||||
if usage == nil {
|
||||
return 0, 0, 0
|
||||
}
|
||||
return usage.TotalTokens, usage.InputTokens, usage.OutputTokens
|
||||
}
|
||||
|
||||
func mergeContextWindowUsage(current *ContextWindowUsage, candidate *ContextWindowUsage) *ContextWindowUsage {
|
||||
if candidate == nil {
|
||||
return cloneContextWindowUsage(current)
|
||||
}
|
||||
if current == nil {
|
||||
return cloneContextWindowUsage(candidate)
|
||||
}
|
||||
|
||||
result := cloneContextWindowUsage(current)
|
||||
currentTotal, currentInput, currentOutput := contextWindowUsageRank(result)
|
||||
candidateTotal, candidateInput, candidateOutput := contextWindowUsageRank(candidate)
|
||||
if candidateTotal > currentTotal ||
|
||||
(candidateTotal == currentTotal && candidateInput > currentInput) ||
|
||||
(candidateTotal == currentTotal && candidateInput == currentInput && candidateOutput > currentOutput) {
|
||||
result = cloneContextWindowUsage(candidate)
|
||||
}
|
||||
|
||||
if result.ContextWindowTokens == 0 && candidate.ContextWindowTokens > 0 {
|
||||
result.ContextWindowTokens = candidate.ContextWindowTokens
|
||||
result.ContextWindowTokensEstimated = candidate.ContextWindowTokensEstimated
|
||||
}
|
||||
if result.ContextWindowTokens > 0 && result.ContextWindowTokensEstimated &&
|
||||
candidate.ContextWindowTokens > 0 && !candidate.ContextWindowTokensEstimated {
|
||||
result.ContextWindowTokens = candidate.ContextWindowTokens
|
||||
result.ContextWindowTokensEstimated = false
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func resolveContextWindowUsage(ctx context.Context, runtime *RuntimeConfig, model string, inputTokens, outputTokens, totalTokens int64) *ContextWindowUsage {
|
||||
windowTokens, estimated := resolveContextWindowTokens(ctx, runtime, model)
|
||||
return newContextWindowUsage(inputTokens, outputTokens, totalTokens, windowTokens, estimated)
|
||||
}
|
||||
|
||||
func resolveContextWindowTokens(ctx context.Context, runtime *RuntimeConfig, model string) (int64, bool) {
|
||||
model = normalizeContextWindowModelID(model)
|
||||
if model == "" {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
if tokens := openAIKnownContextWindowTokens(model); tokens > 0 {
|
||||
return tokens, true
|
||||
}
|
||||
|
||||
if openRouterCompatibleRuntime(runtime) {
|
||||
if tokens, ok := openRouterContextWindowTokens(ctx, runtime, model); ok {
|
||||
return tokens, false
|
||||
}
|
||||
}
|
||||
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func openAIKnownContextWindowTokens(model string) int64 {
|
||||
model = normalizeContextWindowModelID(model)
|
||||
switch {
|
||||
case model == "codex-mini-latest":
|
||||
return openAIReasoningContextWindowTokens
|
||||
case strings.HasPrefix(model, "o3"), strings.HasPrefix(model, "o4-"):
|
||||
return openAIReasoningContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-5.4"):
|
||||
return openAI5_4ContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-4.1"):
|
||||
return openAI4_1ContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-5.1-chat"), strings.HasPrefix(model, "gpt-5-chat"):
|
||||
return openAIChatContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-5"):
|
||||
return openAIFrontierContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-4o-mini-realtime"):
|
||||
return openAIRealtimeMiniContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-4o-realtime"):
|
||||
return openAIRealtimeContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-4o"), strings.HasPrefix(model, "chatgpt-4o"):
|
||||
return openAIChatContextWindowTokens
|
||||
case strings.Contains(model, "gpt-4-turbo"):
|
||||
return openAIChatContextWindowTokens
|
||||
case strings.HasPrefix(model, "gpt-4"):
|
||||
return 8192
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeContextWindowModelID(model string) string {
|
||||
model = strings.ToLower(strings.TrimSpace(model))
|
||||
if idx := strings.Index(model, "/"); idx >= 0 {
|
||||
model = model[idx+1:]
|
||||
}
|
||||
if idx := strings.Index(model, ":"); idx >= 0 {
|
||||
model = model[:idx]
|
||||
}
|
||||
return strings.TrimSpace(model)
|
||||
}
|
||||
|
||||
func openRouterCompatibleRuntime(runtime *RuntimeConfig) bool {
|
||||
if runtime == nil {
|
||||
return false
|
||||
}
|
||||
switch NormalizeProviderName(runtime.Provider) {
|
||||
case ProviderOpenRouter:
|
||||
return true
|
||||
case ProviderOpenAICompat:
|
||||
baseURL := strings.ToLower(strings.TrimSpace(openAIBaseURL(runtime)))
|
||||
return strings.Contains(baseURL, "openrouter.ai")
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
type openRouterModelsResponse struct {
|
||||
Data []openRouterModel `json:"data"`
|
||||
}
|
||||
|
||||
type openRouterModel struct {
|
||||
ID string `json:"id"`
|
||||
TopProvider struct {
|
||||
ContextLength int64 `json:"context_length"`
|
||||
} `json:"top_provider"`
|
||||
}
|
||||
|
||||
type cachedOpenRouterContextWindow struct {
|
||||
tokens int64
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
var openRouterContextWindowCache = struct {
|
||||
sync.RWMutex
|
||||
entries map[string]cachedOpenRouterContextWindow
|
||||
}{
|
||||
entries: map[string]cachedOpenRouterContextWindow{},
|
||||
}
|
||||
|
||||
func openRouterContextWindowTokens(ctx context.Context, runtime *RuntimeConfig, model string) (int64, bool) {
|
||||
model = normalizeContextWindowModelID(model)
|
||||
if model == "" {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
cacheKey := openRouterContextWindowCacheKey(runtime, model)
|
||||
if tokens, ok := cachedOpenRouterContextWindowTokens(cacheKey); ok {
|
||||
return tokens, true
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, openRouterModelsURL(runtime), nil)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
if runtime != nil && !runtime.SkipAuth && strings.TrimSpace(runtime.APIKey) != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(runtime.APIKey))
|
||||
}
|
||||
for key, value := range providerHeaders(runtime) {
|
||||
req.Header.Set(key, value)
|
||||
}
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
payload := &openRouterModelsResponse{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(payload); err != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
for _, candidate := range payload.Data {
|
||||
candidateID := normalizeContextWindowModelID(candidate.ID)
|
||||
if candidateID != model {
|
||||
continue
|
||||
}
|
||||
if candidate.TopProvider.ContextLength <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
storeOpenRouterContextWindowTokens(cacheKey, candidate.TopProvider.ContextLength)
|
||||
return candidate.TopProvider.ContextLength, true
|
||||
}
|
||||
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func openRouterModelsURL(runtime *RuntimeConfig) string {
|
||||
baseURL := strings.TrimSpace(openAIBaseURL(runtime))
|
||||
if baseURL == "" {
|
||||
baseURL = defaultOpenRouterBaseURL
|
||||
}
|
||||
|
||||
parsed, err := url.Parse(baseURL)
|
||||
if err != nil || parsed == nil {
|
||||
return defaultOpenRouterBaseURL + "/models"
|
||||
}
|
||||
|
||||
path := strings.TrimRight(parsed.Path, "/")
|
||||
switch {
|
||||
case strings.HasSuffix(path, "/chat/completions"):
|
||||
path = strings.TrimSuffix(path, "/chat/completions")
|
||||
case strings.HasSuffix(path, "/responses"):
|
||||
path = strings.TrimSuffix(path, "/responses")
|
||||
case path == "":
|
||||
path = "/api/v1"
|
||||
}
|
||||
if !strings.HasSuffix(path, "/models") {
|
||||
path += "/models"
|
||||
}
|
||||
parsed.Path = path
|
||||
parsed.RawQuery = ""
|
||||
parsed.Fragment = ""
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
func openRouterContextWindowCacheKey(runtime *RuntimeConfig, model string) string {
|
||||
return openRouterModelsURL(runtime) + "|" + normalizeContextWindowModelID(model)
|
||||
}
|
||||
|
||||
func cachedOpenRouterContextWindowTokens(cacheKey string) (int64, bool) {
|
||||
if cacheKey == "" {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
openRouterContextWindowCache.RLock()
|
||||
cached, ok := openRouterContextWindowCache.entries[cacheKey]
|
||||
openRouterContextWindowCache.RUnlock()
|
||||
if !ok || cached.tokens <= 0 || now.After(cached.expiresAt) {
|
||||
if ok && now.After(cached.expiresAt) {
|
||||
openRouterContextWindowCache.Lock()
|
||||
delete(openRouterContextWindowCache.entries, cacheKey)
|
||||
openRouterContextWindowCache.Unlock()
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
return cached.tokens, true
|
||||
}
|
||||
|
||||
func storeOpenRouterContextWindowTokens(cacheKey string, tokens int64) {
|
||||
if cacheKey == "" || tokens <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
openRouterContextWindowCache.Lock()
|
||||
openRouterContextWindowCache.entries[cacheKey] = cachedOpenRouterContextWindow{
|
||||
tokens: tokens,
|
||||
expiresAt: time.Now().Add(openRouterContextWindowCacheTTL),
|
||||
}
|
||||
openRouterContextWindowCache.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestResolveContextWindowUsageOpenRouterUsesModelsAPIForUnknownModels(t *testing.T) {
|
||||
openRouterContextWindowCache.Lock()
|
||||
openRouterContextWindowCache.entries = map[string]cachedOpenRouterContextWindow{}
|
||||
openRouterContextWindowCache.Unlock()
|
||||
|
||||
requests := make(chan string, 1)
|
||||
restoreClient := SetHTTPClientForTests(&http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
requests <- r.URL.Path
|
||||
return jsonResponse(http.StatusOK, `{
|
||||
"data": [
|
||||
{
|
||||
"id": "anthropic/claude-sonnet-4-0",
|
||||
"top_provider": {
|
||||
"context_length": 200000
|
||||
}
|
||||
}
|
||||
]
|
||||
}`), nil
|
||||
}),
|
||||
})
|
||||
defer restoreClient()
|
||||
|
||||
usage := resolveContextWindowUsage(context.Background(), &RuntimeConfig{
|
||||
Provider: ProviderOpenRouter,
|
||||
APIKey: "openrouter-key",
|
||||
}, "anthropic/claude-sonnet-4-0", 24000, 1000, 25000)
|
||||
if usage == nil {
|
||||
t.Fatal("expected context window usage")
|
||||
}
|
||||
if usage.ContextWindowTokens != 200000 {
|
||||
t.Fatalf("unexpected context window tokens: %d", usage.ContextWindowTokens)
|
||||
}
|
||||
if usage.ContextWindowTokensEstimated {
|
||||
t.Fatal("expected OpenRouter model lookup to be treated as exact")
|
||||
}
|
||||
|
||||
requestPath := <-requests
|
||||
if requestPath != "/api/v1/models" {
|
||||
t.Fatalf("unexpected models lookup path: %q", requestPath)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
type completionRequest struct {
|
||||
SystemPrompt string
|
||||
Messages []providerMessage
|
||||
}
|
||||
|
||||
type sdkDriver interface {
|
||||
Name() string
|
||||
Supports(*RuntimeConfig) bool
|
||||
CompleteConversation(context.Context, *RuntimeConfig, *completionRequest) (*Completion, error)
|
||||
}
|
||||
|
||||
var completionDrivers = []sdkDriver{
|
||||
newOpenAIDriver(),
|
||||
}
|
||||
|
||||
func selectCompletionDriver(runtime *RuntimeConfig) (sdkDriver, error) {
|
||||
if runtime == nil {
|
||||
return nil, fmt.Errorf("AI runtime config is required")
|
||||
}
|
||||
|
||||
for _, driver := range completionDrivers {
|
||||
if driver.Supports(runtime) {
|
||||
return driver, nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf(
|
||||
"server AI provider %q does not have an available SDK driver; use openai, openai-compat, or openrouter, or add a driver in server/ai",
|
||||
runtime.Provider,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,389 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
openai "github.com/openai/openai-go/v2"
|
||||
"github.com/openai/openai-go/v2/option"
|
||||
openairesponses "github.com/openai/openai-go/v2/responses"
|
||||
"github.com/openai/openai-go/v2/shared"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIDriverName = "openai-go"
|
||||
defaultOpenRouterBaseURL = "https://openrouter.ai/api/v1"
|
||||
openAIWebSearchContext = "medium"
|
||||
)
|
||||
|
||||
type openAIDriver struct{}
|
||||
|
||||
type openAIClient struct {
|
||||
chat openai.ChatCompletionService
|
||||
responses openairesponses.ResponseService
|
||||
}
|
||||
|
||||
func newOpenAIDriver() sdkDriver {
|
||||
return &openAIDriver{}
|
||||
}
|
||||
|
||||
func (d *openAIDriver) Name() string {
|
||||
return openAIDriverName
|
||||
}
|
||||
|
||||
func (d *openAIDriver) Supports(runtime *RuntimeConfig) bool {
|
||||
if runtime == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
switch NormalizeProviderName(runtime.Provider) {
|
||||
case ProviderOpenAI, ProviderOpenAICompat, ProviderOpenRouter:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (d *openAIDriver) CompleteConversation(ctx context.Context, runtime *RuntimeConfig, request *completionRequest) (*Completion, error) {
|
||||
if runtime == nil {
|
||||
return nil, fmt.Errorf("AI runtime config is required")
|
||||
}
|
||||
if request == nil {
|
||||
return nil, fmt.Errorf("AI completion request is required")
|
||||
}
|
||||
|
||||
client := newOpenAIClient(openAIRequestOptions(runtime)...)
|
||||
if runtime.UseResponsesAPI {
|
||||
return d.completeResponses(ctx, client, runtime, request)
|
||||
}
|
||||
return d.completeChat(ctx, client, runtime, request)
|
||||
}
|
||||
|
||||
func newOpenAIClient(opts ...option.RequestOption) *openAIClient {
|
||||
return &openAIClient{
|
||||
chat: openai.NewChatCompletionService(opts...),
|
||||
responses: openairesponses.NewResponseService(opts...),
|
||||
}
|
||||
}
|
||||
|
||||
func openAIRequestOptions(runtime *RuntimeConfig) []option.RequestOption {
|
||||
opts := []option.RequestOption{
|
||||
option.WithHTTPClient(httpClient),
|
||||
option.WithRequestTimeout(defaultCompletionTimeout),
|
||||
option.WithMaxRetries(0),
|
||||
}
|
||||
|
||||
if runtime != nil && NormalizeProviderName(runtime.Provider) == ProviderOpenAI {
|
||||
opts = append(opts, option.WithEnvironmentProduction())
|
||||
}
|
||||
if baseURL := openAIBaseURL(runtime); baseURL != "" {
|
||||
opts = append(opts, option.WithBaseURL(baseURL))
|
||||
}
|
||||
if runtime != nil {
|
||||
if runtime.APIKey != "" {
|
||||
opts = append(opts, option.WithAPIKey(runtime.APIKey))
|
||||
}
|
||||
if runtime.Organization != "" {
|
||||
opts = append(opts, option.WithOrganization(runtime.Organization))
|
||||
}
|
||||
if runtime.Project != "" {
|
||||
opts = append(opts, option.WithProject(runtime.Project))
|
||||
}
|
||||
for key, value := range providerHeaders(runtime) {
|
||||
opts = append(opts, option.WithHeader(key, value))
|
||||
}
|
||||
if runtime.TopK != nil && NormalizeProviderName(runtime.Provider) != ProviderOpenAI {
|
||||
opts = append(opts, option.WithJSONSet("top_k", *runtime.TopK))
|
||||
}
|
||||
}
|
||||
|
||||
return opts
|
||||
}
|
||||
|
||||
func openAIBaseURL(runtime *RuntimeConfig) string {
|
||||
if runtime == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
if runtime.BaseURL != "" {
|
||||
return runtime.BaseURL
|
||||
}
|
||||
if NormalizeProviderName(runtime.Provider) == ProviderOpenRouter {
|
||||
return defaultOpenRouterBaseURL
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func openAIWebSearchEnabled(runtime *RuntimeConfig) bool {
|
||||
return runtime != nil && NormalizeProviderName(runtime.Provider) == ProviderOpenAI
|
||||
}
|
||||
|
||||
func appendOpenAIResponseTools(runtime *RuntimeConfig, tools []openairesponses.ToolUnionParam) []openairesponses.ToolUnionParam {
|
||||
if !openAIWebSearchEnabled(runtime) {
|
||||
return tools
|
||||
}
|
||||
|
||||
for _, tool := range tools {
|
||||
if tool.OfWebSearch != nil || tool.OfWebSearchPreview != nil {
|
||||
return tools
|
||||
}
|
||||
}
|
||||
|
||||
webSearch := openairesponses.ToolParamOfWebSearch(openairesponses.WebSearchToolTypeWebSearch)
|
||||
if webSearch.OfWebSearch != nil {
|
||||
webSearch.OfWebSearch.SearchContextSize = openairesponses.WebSearchToolSearchContextSizeMedium
|
||||
}
|
||||
|
||||
return append(append([]openairesponses.ToolUnionParam{}, tools...), webSearch)
|
||||
}
|
||||
|
||||
func applyOpenAIChatWebSearch(runtime *RuntimeConfig, params *openai.ChatCompletionNewParams) {
|
||||
if params == nil || !openAIWebSearchEnabled(runtime) {
|
||||
return
|
||||
}
|
||||
|
||||
params.WebSearchOptions = openai.ChatCompletionNewParamsWebSearchOptions{
|
||||
SearchContextSize: openAIWebSearchContext,
|
||||
}
|
||||
}
|
||||
|
||||
func (d *openAIDriver) completeResponses(ctx context.Context, client *openAIClient, runtime *RuntimeConfig, request *completionRequest) (*Completion, error) {
|
||||
params := openairesponses.ResponseNewParams{
|
||||
Input: openairesponses.ResponseNewParamsInputUnion{
|
||||
OfInputItemList: buildResponseInput(request.SystemPrompt, request.Messages),
|
||||
},
|
||||
Model: shared.ResponsesModel(runtime.Model),
|
||||
}
|
||||
if tools := appendOpenAIResponseTools(runtime, nil); len(tools) > 0 {
|
||||
params.Tools = tools
|
||||
}
|
||||
if runtime.MaxOutputTokens > 0 {
|
||||
params.MaxOutputTokens = openai.Int(runtime.MaxOutputTokens)
|
||||
}
|
||||
if runtime.Temperature != nil {
|
||||
params.Temperature = openai.Float(*runtime.Temperature)
|
||||
}
|
||||
if runtime.TopP != nil {
|
||||
params.TopP = openai.Float(*runtime.TopP)
|
||||
}
|
||||
if reasoning, ok := responseReasoningParam(runtime); ok {
|
||||
params.Reasoning = reasoning
|
||||
}
|
||||
|
||||
response, err := client.responses.New(ctx, params)
|
||||
if err != nil {
|
||||
return nil, formatOpenAIError(runtime.Provider, err)
|
||||
}
|
||||
|
||||
content := strings.TrimSpace(response.OutputText())
|
||||
if content == "" {
|
||||
return nil, fmt.Errorf("%s response did not include assistant text", runtime.Provider)
|
||||
}
|
||||
|
||||
model := fallbackString(response.Model, runtime.Model)
|
||||
return &Completion{
|
||||
Provider: runtime.Provider,
|
||||
Model: model,
|
||||
Content: content,
|
||||
ProviderMessageID: strings.TrimSpace(response.ID),
|
||||
FinishReason: responseFinishReason(response),
|
||||
ContextWindowUsage: resolveContextWindowUsage(
|
||||
ctx,
|
||||
runtime,
|
||||
model,
|
||||
response.Usage.InputTokens,
|
||||
response.Usage.OutputTokens,
|
||||
response.Usage.TotalTokens,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (d *openAIDriver) completeChat(ctx context.Context, client *openAIClient, runtime *RuntimeConfig, request *completionRequest) (*Completion, error) {
|
||||
params := openai.ChatCompletionNewParams{
|
||||
Messages: buildChatMessages(request.SystemPrompt, request.Messages),
|
||||
Model: shared.ChatModel(runtime.Model),
|
||||
}
|
||||
applyOpenAIChatWebSearch(runtime, ¶ms)
|
||||
if runtime.MaxOutputTokens > 0 {
|
||||
if NormalizeProviderName(runtime.Provider) == ProviderOpenAI {
|
||||
params.MaxCompletionTokens = openai.Int(runtime.MaxOutputTokens)
|
||||
} else {
|
||||
params.MaxTokens = openai.Int(runtime.MaxOutputTokens)
|
||||
}
|
||||
}
|
||||
if runtime.Temperature != nil {
|
||||
params.Temperature = openai.Float(*runtime.Temperature)
|
||||
}
|
||||
if runtime.TopP != nil {
|
||||
params.TopP = openai.Float(*runtime.TopP)
|
||||
}
|
||||
if runtime.PresencePenalty != nil {
|
||||
params.PresencePenalty = openai.Float(*runtime.PresencePenalty)
|
||||
}
|
||||
if runtime.FrequencyPenalty != nil {
|
||||
params.FrequencyPenalty = openai.Float(*runtime.FrequencyPenalty)
|
||||
}
|
||||
if effort, ok := reasoningEffortForThinkingLevel(runtime.ThinkingLevel); ok {
|
||||
params.ReasoningEffort = effort
|
||||
}
|
||||
|
||||
response, err := client.chat.New(ctx, params)
|
||||
if err != nil {
|
||||
return nil, formatOpenAIError(runtime.Provider, err)
|
||||
}
|
||||
|
||||
content, finishReason := extractChatCompletionResult(response)
|
||||
if content == "" {
|
||||
return nil, fmt.Errorf("%s response did not include assistant text", runtime.Provider)
|
||||
}
|
||||
|
||||
model := fallbackString(response.Model, runtime.Model)
|
||||
return &Completion{
|
||||
Provider: runtime.Provider,
|
||||
Model: model,
|
||||
Content: content,
|
||||
ProviderMessageID: strings.TrimSpace(response.ID),
|
||||
FinishReason: finishReason,
|
||||
ContextWindowUsage: resolveContextWindowUsage(
|
||||
ctx,
|
||||
runtime,
|
||||
model,
|
||||
response.Usage.PromptTokens,
|
||||
response.Usage.CompletionTokens,
|
||||
response.Usage.TotalTokens,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func buildResponseInput(systemPrompt string, messages []providerMessage) openairesponses.ResponseInputParam {
|
||||
input := make(openairesponses.ResponseInputParam, 0, len(messages)+1)
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
input = append(input, openairesponses.ResponseInputItemParamOfMessage(strings.TrimSpace(systemPrompt), openairesponses.EasyInputMessageRoleSystem))
|
||||
}
|
||||
for _, message := range messages {
|
||||
input = append(input, openairesponses.ResponseInputItemParamOfMessage(message.Content, responseMessageRole(message.Role)))
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
func buildChatMessages(systemPrompt string, messages []providerMessage) []openai.ChatCompletionMessageParamUnion {
|
||||
chatMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages)+1)
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
chatMessages = append(chatMessages, openai.SystemMessage(strings.TrimSpace(systemPrompt)))
|
||||
}
|
||||
for _, message := range messages {
|
||||
switch strings.ToLower(strings.TrimSpace(message.Role)) {
|
||||
case "assistant":
|
||||
chatMessages = append(chatMessages, openai.AssistantMessage(message.Content))
|
||||
default:
|
||||
chatMessages = append(chatMessages, openai.UserMessage(message.Content))
|
||||
}
|
||||
}
|
||||
return chatMessages
|
||||
}
|
||||
|
||||
func responseMessageRole(role string) openairesponses.EasyInputMessageRole {
|
||||
switch strings.ToLower(strings.TrimSpace(role)) {
|
||||
case "assistant":
|
||||
return openairesponses.EasyInputMessageRoleAssistant
|
||||
default:
|
||||
return openairesponses.EasyInputMessageRoleUser
|
||||
}
|
||||
}
|
||||
|
||||
func extractChatCompletionResult(response *openai.ChatCompletion) (string, string) {
|
||||
if response == nil || len(response.Choices) == 0 {
|
||||
return "", ""
|
||||
}
|
||||
|
||||
for _, choice := range response.Choices {
|
||||
content := strings.TrimSpace(choice.Message.Content)
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
finishReason := strings.TrimSpace(choice.FinishReason)
|
||||
if finishReason == "" {
|
||||
finishReason = "completed"
|
||||
}
|
||||
return content, finishReason
|
||||
}
|
||||
|
||||
return "", ""
|
||||
}
|
||||
|
||||
func responseFinishReason(response *openairesponses.Response) string {
|
||||
if response == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
if response.Status == openairesponses.ResponseStatusIncomplete && strings.TrimSpace(response.IncompleteDetails.Reason) != "" {
|
||||
return strings.TrimSpace(response.IncompleteDetails.Reason)
|
||||
}
|
||||
|
||||
finishReason := strings.TrimSpace(string(response.Status))
|
||||
if finishReason == "" {
|
||||
return "completed"
|
||||
}
|
||||
return finishReason
|
||||
}
|
||||
|
||||
func reasoningEffortForThinkingLevel(thinkingLevel string) (shared.ReasoningEffort, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(thinkingLevel)) {
|
||||
case "low":
|
||||
return shared.ReasoningEffortLow, true
|
||||
case "medium":
|
||||
return shared.ReasoningEffortMedium, true
|
||||
case "high":
|
||||
return shared.ReasoningEffortHigh, true
|
||||
case "xhigh":
|
||||
return shared.ReasoningEffort("xhigh"), true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func responseReasoningParam(runtime *RuntimeConfig) (shared.ReasoningParam, bool) {
|
||||
effort, ok := reasoningEffortForThinkingLevel(runtime.ThinkingLevel)
|
||||
if !ok {
|
||||
return shared.ReasoningParam{}, false
|
||||
}
|
||||
|
||||
reasoning := shared.ReasoningParam{Effort: effort}
|
||||
switch NormalizeProviderName(runtime.Provider) {
|
||||
case ProviderOpenAI, ProviderOpenRouter:
|
||||
// Native OpenAI-style Responses backends can return a reasoning summary
|
||||
// for UI-only transcript blocks. Use "auto" because summary modes are
|
||||
// model-specific; hard-coding "concise" can suppress summaries on models
|
||||
// that only expose a different summarizer.
|
||||
reasoning.Summary = shared.ReasoningSummaryAuto
|
||||
}
|
||||
|
||||
return reasoning, true
|
||||
}
|
||||
|
||||
func formatOpenAIError(provider string, err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var apiErr *openai.Error
|
||||
if errors.As(err, &apiErr) {
|
||||
message := strings.TrimSpace(apiErr.Message)
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(apiErr.RawJSON())
|
||||
}
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(err.Error())
|
||||
}
|
||||
if apiErr.StatusCode > 0 {
|
||||
return fmt.Errorf("%s API request failed with HTTP %d: %s", provider, apiErr.StatusCode, truncateForError(message))
|
||||
}
|
||||
if message != "" {
|
||||
return fmt.Errorf("%s API request failed: %s", provider, truncateForError(message))
|
||||
}
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package assets
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
const (
|
||||
// AIDirName stores server-side AI helper data under the Sliver root.
|
||||
AIDirName = "ai"
|
||||
|
||||
// AIAliasesDirName stores aliases made available to the server-side AI agent.
|
||||
AIAliasesDirName = "aliases"
|
||||
|
||||
// AIExtensionsDirName stores extensions made available to the server-side AI agent.
|
||||
AIExtensionsDirName = "extensions"
|
||||
)
|
||||
|
||||
// GetAIDir returns the server-side AI data directory: ~/.sliver/ai
|
||||
func GetAIDir() string {
|
||||
return ensureAIDir(filepath.Join(GetRootAppDir(), AIDirName))
|
||||
}
|
||||
|
||||
// GetAIAliasesDir returns the server-side AI alias directory: ~/.sliver/ai/aliases
|
||||
func GetAIAliasesDir() string {
|
||||
return ensureAIDir(filepath.Join(GetAIDir(), AIAliasesDirName))
|
||||
}
|
||||
|
||||
// GetAIExtensionsDir returns the server-side AI extension directory: ~/.sliver/ai/extensions
|
||||
func GetAIExtensionsDir() string {
|
||||
return ensureAIDir(filepath.Join(GetAIDir(), AIExtensionsDirName))
|
||||
}
|
||||
|
||||
func ensureAIDir(dir string) string {
|
||||
if _, err := os.Stat(dir); os.IsNotExist(err) {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
setupLog.Fatalf("Cannot create AI data dir %s", err)
|
||||
}
|
||||
}
|
||||
return dir
|
||||
}
|
||||
+305
-30
@@ -18,10 +18,62 @@ package configs
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import "strings"
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/server/assets"
|
||||
"github.com/bishopfox/sliver/server/log"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const (
|
||||
aiConfigFileName = "ai.yaml"
|
||||
aiProviderAnthropic = "anthropic"
|
||||
aiProviderGoogle = "google"
|
||||
aiProviderOpenAI = "openai"
|
||||
aiProviderOpenAICompat = "openai-compat"
|
||||
aiProviderOpenRouter = "openrouter"
|
||||
defaultAISystemPrompt = `You are Sliver's AI copilot for authorized security testing, detection engineering, lab work, and incident-response support in environments the operator is explicitly permitted to assess.
|
||||
|
||||
Your job is to help the operator make careful, high-signal decisions inside the Sliver workflow.
|
||||
|
||||
Operating rules:
|
||||
- Assume all activity must stay within the operator's authorized scope. If a request is ambiguous, unusually risky, destructive, or inconsistent with the available context, pause, state the concern, and ask for confirmation or offer a safer alternative.
|
||||
- Use the current conversation, target metadata, and Sliver context to tailor answers. If OS, privilege, transport, or session/beacon state matters, say so before recommending commands.
|
||||
- Never fabricate command output, host state, credentials, files, loot, or tool results. When information is missing, say exactly what is unknown.
|
||||
- Prefer the smallest next step that increases certainty. Start with low-noise, reversible discovery before collection, execution, or configuration changes.
|
||||
- Highlight operational tradeoffs: stealth, telemetry, privilege requirements, target stability, cleanup burden, and chances of breaking access.
|
||||
- Prefer actions that preserve operator access, host stability, and forensic defensibility. Avoid destructive, noisy, or irreversible steps unless the operator explicitly asks for them and the purpose is clear.
|
||||
- Treat credentials, secrets, tokens, and loot as sensitive. Do not suggest unnecessary exposure, duplication, or broad collection.
|
||||
- When troubleshooting, separate confirmed facts from hypotheses and propose the next diagnostic step rather than many speculative changes at once.
|
||||
- When multiple approaches are possible, present the safest practical option first, then note faster or more aggressive alternatives only if they materially help.
|
||||
- When suggesting Sliver commands, make them concrete, minimal, and easy to audit. Include assumptions and prerequisites when they matter.
|
||||
- When asked to draft code, scripts, or one-liners, optimize for readability, explicitness, and minimal side effects.
|
||||
|
||||
Response style:
|
||||
- Be concise, structured, and operator-focused.
|
||||
- Prefer short checklists or step plans over long essays.
|
||||
- Distinguish clearly between facts, assumptions, and recommendations.
|
||||
- If the request appears outside authorized security work, refuse and redirect to safer high-level guidance.`
|
||||
)
|
||||
|
||||
var (
|
||||
aiConfigLog = log.NamedLogger("config", "ai")
|
||||
supportedAIProviders = []string{
|
||||
aiProviderAnthropic,
|
||||
aiProviderGoogle,
|
||||
aiProviderOpenAI,
|
||||
aiProviderOpenAICompat,
|
||||
aiProviderOpenRouter,
|
||||
}
|
||||
)
|
||||
|
||||
// AIProviderConfig - Shared AI provider configuration.
|
||||
type AIProviderConfig struct {
|
||||
Models []string `json:"models,omitempty" yaml:"models,omitempty"`
|
||||
APIKey string `json:"api_key" yaml:"api_key"`
|
||||
BaseURL string `json:"base_url" yaml:"base_url"`
|
||||
Headers map[string]string `json:"headers,omitempty" yaml:"headers,omitempty"`
|
||||
@@ -37,8 +89,9 @@ type AIProviderConfig struct {
|
||||
// AIConfig - Server-side AI provider configuration.
|
||||
type AIConfig struct {
|
||||
Provider string `json:"provider" yaml:"provider"`
|
||||
Model string `json:"model" yaml:"model"`
|
||||
Model string `json:"model,omitempty" yaml:"model,omitempty"` // Deprecated legacy shared model field; migrated into provider-specific models on save.
|
||||
ThinkingLevel string `json:"thinking_level" yaml:"thinking_level"`
|
||||
SystemPrompt string `json:"system_prompt" yaml:"system_prompt"`
|
||||
MaxOutputTokens int64 `json:"max_output_tokens,omitempty" yaml:"max_output_tokens,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty" yaml:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty" yaml:"top_p,omitempty"`
|
||||
@@ -52,6 +105,18 @@ type AIConfig struct {
|
||||
OpenRouter *AIProviderConfig `json:"openrouter" yaml:"openrouter"`
|
||||
}
|
||||
|
||||
type aiConfigEnvelope struct {
|
||||
AI *AIConfig `json:"ai" yaml:"ai"`
|
||||
}
|
||||
|
||||
// GetAIConfigPath - File path to ai.yaml.
|
||||
func GetAIConfigPath() string {
|
||||
appDir := assets.GetRootAppDir()
|
||||
aiConfigPath := filepath.Join(appDir, "configs", aiConfigFileName)
|
||||
aiConfigLog.Debugf("Loading config from %s", aiConfigPath)
|
||||
return aiConfigPath
|
||||
}
|
||||
|
||||
func defaultAIProviderConfig() *AIProviderConfig {
|
||||
return &AIProviderConfig{
|
||||
Headers: map[string]string{},
|
||||
@@ -63,6 +128,7 @@ func defaultAIConfig() *AIConfig {
|
||||
Provider: "",
|
||||
Model: "",
|
||||
ThinkingLevel: "",
|
||||
SystemPrompt: defaultAISystemPrompt,
|
||||
MaxOutputTokens: 0,
|
||||
Temperature: nil,
|
||||
TopP: nil,
|
||||
@@ -77,47 +143,135 @@ func defaultAIConfig() *AIConfig {
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeAIConfig(config *ServerConfig) {
|
||||
if config.AI == nil {
|
||||
config.AI = defaultAIConfig()
|
||||
// Save - Save AI config file to disk.
|
||||
func (c *AIConfig) Save() error {
|
||||
config := normalizeAIConfig(c)
|
||||
configPath := GetAIConfigPath()
|
||||
configDir := filepath.Dir(configPath)
|
||||
if _, err := os.Stat(configDir); os.IsNotExist(err) {
|
||||
aiConfigLog.Debugf("Creating config dir %s", configDir)
|
||||
err := os.MkdirAll(configDir, 0700)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
config.AI.Provider = strings.ToLower(strings.TrimSpace(config.AI.Provider))
|
||||
config.AI.Model = strings.TrimSpace(config.AI.Model)
|
||||
config.AI.ThinkingLevel = strings.ToLower(strings.TrimSpace(config.AI.ThinkingLevel))
|
||||
if config.AI.MaxOutputTokens < 0 {
|
||||
config.AI.MaxOutputTokens = 0
|
||||
data, err := yaml.Marshal(config)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
config.AI.Temperature = normalizeOptionalFloat(config.AI.Temperature)
|
||||
config.AI.TopP = normalizeOptionalFloat(config.AI.TopP)
|
||||
config.AI.TopK = normalizeOptionalInt(config.AI.TopK)
|
||||
config.AI.PresencePenalty = normalizeOptionalFloat(config.AI.PresencePenalty)
|
||||
config.AI.FrequencyPenalty = normalizeOptionalFloat(config.AI.FrequencyPenalty)
|
||||
if config.AI.Anthropic == nil {
|
||||
config.AI.Anthropic = defaultAIProviderConfig()
|
||||
aiConfigLog.Infof("Saving config to %s", configPath)
|
||||
err = os.WriteFile(configPath, data, 0600)
|
||||
if err != nil {
|
||||
aiConfigLog.Errorf("Failed to write config %s", err)
|
||||
}
|
||||
normalizeAIProviderConfig(config.AI.Anthropic)
|
||||
if config.AI.Google == nil {
|
||||
config.AI.Google = defaultAIProviderConfig()
|
||||
return err
|
||||
}
|
||||
|
||||
// GetAIConfig - Get AI config value from ai.yaml.
|
||||
func GetAIConfig() *AIConfig {
|
||||
return getAIConfig(nil)
|
||||
}
|
||||
|
||||
func getAIConfig(migratedConfig *AIConfig) *AIConfig {
|
||||
configPath := GetAIConfigPath()
|
||||
config := defaultAIConfig()
|
||||
if _, err := os.Stat(configPath); !os.IsNotExist(err) {
|
||||
data, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
aiConfigLog.Errorf("Failed to read config file %s", err)
|
||||
return config
|
||||
}
|
||||
err = yaml.Unmarshal(data, config)
|
||||
if err != nil {
|
||||
aiConfigLog.Errorf("Failed to parse config file %s", err)
|
||||
return config
|
||||
}
|
||||
} else if migratedConfig != nil {
|
||||
config = migratedConfig
|
||||
aiConfigLog.Infof("Migrating embedded AI config to %s", configPath)
|
||||
} else {
|
||||
aiConfigLog.Warnf("Config file does not exist, using defaults")
|
||||
}
|
||||
normalizeAIProviderConfig(config.AI.Google)
|
||||
if config.AI.OpenAI == nil {
|
||||
config.AI.OpenAI = defaultAIProviderConfig()
|
||||
|
||||
config = normalizeAIConfig(config)
|
||||
|
||||
err := config.Save() // This updates the config with any missing fields
|
||||
if err != nil {
|
||||
aiConfigLog.Errorf("Failed to save default config %s", err)
|
||||
return config
|
||||
}
|
||||
normalizeAIProviderConfig(config.AI.OpenAI)
|
||||
if config.AI.OpenAICompat == nil {
|
||||
config.AI.OpenAICompat = defaultAIProviderConfig()
|
||||
return config
|
||||
}
|
||||
|
||||
func aiConfigFromYAML(data []byte) *AIConfig {
|
||||
envelope := &aiConfigEnvelope{}
|
||||
if err := yaml.Unmarshal(data, envelope); err != nil {
|
||||
aiConfigLog.Errorf("Failed to parse embedded AI config %s", err)
|
||||
return nil
|
||||
}
|
||||
normalizeAIProviderConfig(config.AI.OpenAICompat)
|
||||
if config.AI.OpenRouter == nil {
|
||||
config.AI.OpenRouter = defaultAIProviderConfig()
|
||||
if envelope.AI == nil {
|
||||
return nil
|
||||
}
|
||||
normalizeAIProviderConfig(config.AI.OpenRouter)
|
||||
return normalizeAIConfig(envelope.AI)
|
||||
}
|
||||
|
||||
func aiConfigFromJSON(data []byte) *AIConfig {
|
||||
envelope := &aiConfigEnvelope{}
|
||||
if err := json.Unmarshal(data, envelope); err != nil {
|
||||
aiConfigLog.Errorf("Failed to parse embedded AI config %s", err)
|
||||
return nil
|
||||
}
|
||||
if envelope.AI == nil {
|
||||
return nil
|
||||
}
|
||||
return normalizeAIConfig(envelope.AI)
|
||||
}
|
||||
|
||||
func normalizeAIConfig(config *AIConfig) *AIConfig {
|
||||
if config == nil {
|
||||
config = defaultAIConfig()
|
||||
}
|
||||
config.Provider = normalizeAIProviderName(config.Provider)
|
||||
config.Model = strings.TrimSpace(config.Model)
|
||||
config.ThinkingLevel = strings.ToLower(strings.TrimSpace(config.ThinkingLevel))
|
||||
config.SystemPrompt = strings.TrimSpace(config.SystemPrompt)
|
||||
if config.MaxOutputTokens < 0 {
|
||||
config.MaxOutputTokens = 0
|
||||
}
|
||||
config.Temperature = normalizeOptionalFloat(config.Temperature)
|
||||
config.TopP = normalizeOptionalFloat(config.TopP)
|
||||
config.TopK = normalizeOptionalInt(config.TopK)
|
||||
config.PresencePenalty = normalizeOptionalFloat(config.PresencePenalty)
|
||||
config.FrequencyPenalty = normalizeOptionalFloat(config.FrequencyPenalty)
|
||||
if config.Anthropic == nil {
|
||||
config.Anthropic = defaultAIProviderConfig()
|
||||
}
|
||||
normalizeAIProviderConfig(config.Anthropic)
|
||||
if config.Google == nil {
|
||||
config.Google = defaultAIProviderConfig()
|
||||
}
|
||||
normalizeAIProviderConfig(config.Google)
|
||||
if config.OpenAI == nil {
|
||||
config.OpenAI = defaultAIProviderConfig()
|
||||
}
|
||||
normalizeAIProviderConfig(config.OpenAI)
|
||||
if config.OpenAICompat == nil {
|
||||
config.OpenAICompat = defaultAIProviderConfig()
|
||||
}
|
||||
normalizeAIProviderConfig(config.OpenAICompat)
|
||||
if config.OpenRouter == nil {
|
||||
config.OpenRouter = defaultAIProviderConfig()
|
||||
}
|
||||
normalizeAIProviderConfig(config.OpenRouter)
|
||||
migrateLegacyAIModel(config)
|
||||
return config
|
||||
}
|
||||
|
||||
func normalizeAIProviderConfig(provider *AIProviderConfig) {
|
||||
if provider == nil {
|
||||
return
|
||||
}
|
||||
provider.Models = normalizeStringSlice(provider.Models)
|
||||
provider.APIKey = strings.TrimSpace(provider.APIKey)
|
||||
provider.BaseURL = strings.TrimSpace(provider.BaseURL)
|
||||
provider.UserAgent = strings.TrimSpace(provider.UserAgent)
|
||||
@@ -174,3 +328,124 @@ func normalizeStringMap(values map[string]string) map[string]string {
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func migrateLegacyAIModel(config *AIConfig) {
|
||||
if config == nil {
|
||||
return
|
||||
}
|
||||
legacyModel := strings.TrimSpace(config.Model)
|
||||
if legacyModel == "" {
|
||||
return
|
||||
}
|
||||
|
||||
provider := legacyAIModelProvider(config)
|
||||
providerConfig := aiProviderConfigByName(config, provider)
|
||||
if providerConfig != nil {
|
||||
providerConfig.Models = normalizeStringSlice(append([]string{legacyModel}, providerConfig.Models...))
|
||||
}
|
||||
config.Model = ""
|
||||
}
|
||||
|
||||
func legacyAIModelProvider(config *AIConfig) string {
|
||||
if config == nil {
|
||||
return aiProviderOpenAI
|
||||
}
|
||||
if provider := normalizeAIProviderName(config.Provider); isSupportedAIProvider(provider) {
|
||||
return provider
|
||||
}
|
||||
for _, provider := range supportedAIProviders {
|
||||
if aiProviderConfigured(provider, aiProviderConfigByName(config, provider)) {
|
||||
return provider
|
||||
}
|
||||
}
|
||||
return aiProviderOpenAI
|
||||
}
|
||||
|
||||
func normalizeAIProviderName(name string) string {
|
||||
normalized := strings.ToLower(strings.TrimSpace(name))
|
||||
switch normalized {
|
||||
case "gemini":
|
||||
return aiProviderGoogle
|
||||
case "openai_compat", "openaicompat", "openai-compatible":
|
||||
return aiProviderOpenAICompat
|
||||
default:
|
||||
return normalized
|
||||
}
|
||||
}
|
||||
|
||||
func isSupportedAIProvider(provider string) bool {
|
||||
switch normalizeAIProviderName(provider) {
|
||||
case aiProviderAnthropic, aiProviderGoogle, aiProviderOpenAI, aiProviderOpenAICompat, aiProviderOpenRouter:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func aiProviderConfigByName(config *AIConfig, provider string) *AIProviderConfig {
|
||||
if config == nil {
|
||||
return nil
|
||||
}
|
||||
switch normalizeAIProviderName(provider) {
|
||||
case aiProviderAnthropic:
|
||||
return config.Anthropic
|
||||
case aiProviderGoogle:
|
||||
return config.Google
|
||||
case aiProviderOpenAI:
|
||||
return config.OpenAI
|
||||
case aiProviderOpenAICompat:
|
||||
return config.OpenAICompat
|
||||
case aiProviderOpenRouter:
|
||||
return config.OpenRouter
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func aiProviderConfigured(provider string, providerConfig *AIProviderConfig) bool {
|
||||
if providerConfig == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
switch normalizeAIProviderName(provider) {
|
||||
case aiProviderAnthropic:
|
||||
return strings.TrimSpace(providerConfig.APIKey) != "" ||
|
||||
providerConfig.UseBedrock ||
|
||||
(strings.TrimSpace(providerConfig.Project) != "" && strings.TrimSpace(providerConfig.Location) != "")
|
||||
case aiProviderGoogle:
|
||||
return strings.TrimSpace(providerConfig.APIKey) != "" ||
|
||||
(strings.TrimSpace(providerConfig.Project) != "" && strings.TrimSpace(providerConfig.Location) != "")
|
||||
case aiProviderOpenAI:
|
||||
return strings.TrimSpace(providerConfig.APIKey) != ""
|
||||
case aiProviderOpenAICompat:
|
||||
return strings.TrimSpace(providerConfig.BaseURL) != ""
|
||||
case aiProviderOpenRouter:
|
||||
return strings.TrimSpace(providerConfig.APIKey) != ""
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeStringSlice(values []string) []string {
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
normalized := make([]string, 0, len(values))
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[value]; exists {
|
||||
continue
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
normalized = append(normalized, value)
|
||||
}
|
||||
if len(normalized) == 0 {
|
||||
return nil
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
package configs
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAIConfigParsesYAML(t *testing.T) {
|
||||
t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
|
||||
|
||||
configPath := GetAIConfigPath()
|
||||
if err := os.MkdirAll(filepath.Dir(configPath), 0700); err != nil {
|
||||
t.Fatalf("failed to create config dir: %v", err)
|
||||
}
|
||||
|
||||
data := []byte(`provider: "openai"
|
||||
thinking_level: "high"
|
||||
system_prompt: "Stay concise."
|
||||
max_output_tokens: 2048
|
||||
temperature: 0.25
|
||||
anthropic:
|
||||
api_key: "anthropic-key"
|
||||
base_url: "https://api.anthropic.example"
|
||||
use_bedrock: true
|
||||
openai:
|
||||
models:
|
||||
- "gpt-test"
|
||||
- "gpt-mini-test"
|
||||
api_key: "openai-key"
|
||||
base_url: "https://api.openai.example"
|
||||
organization: "org-test"
|
||||
project: "proj-test"
|
||||
use_responses_api: true
|
||||
google:
|
||||
api_key: "google-key"
|
||||
project: "vertex-project"
|
||||
location: "us-central1"
|
||||
skip_auth: true
|
||||
openai_compat:
|
||||
base_url: "http://127.0.0.1:8080/v1"
|
||||
openrouter:
|
||||
api_key: "openrouter-key"
|
||||
user_agent: "sliver-test/1.0"
|
||||
`)
|
||||
if err := os.WriteFile(configPath, data, 0600); err != nil {
|
||||
t.Fatalf("failed to write yaml config: %v", err)
|
||||
}
|
||||
|
||||
config := GetAIConfig()
|
||||
if config.Provider != "openai" {
|
||||
t.Fatalf("expected provider %q, got %q", "openai", config.Provider)
|
||||
}
|
||||
if config.Model != "" {
|
||||
t.Fatalf("expected legacy model field to remain empty, got %q", config.Model)
|
||||
}
|
||||
if config.ThinkingLevel != "high" {
|
||||
t.Fatalf("expected thinking level %q, got %q", "high", config.ThinkingLevel)
|
||||
}
|
||||
if config.SystemPrompt != "Stay concise." {
|
||||
t.Fatalf("expected system prompt %q, got %q", "Stay concise.", config.SystemPrompt)
|
||||
}
|
||||
if config.MaxOutputTokens != 2048 {
|
||||
t.Fatalf("expected max_output_tokens %d, got %d", 2048, config.MaxOutputTokens)
|
||||
}
|
||||
if config.Temperature == nil || *config.Temperature != 0.25 {
|
||||
t.Fatalf("expected temperature %v, got %#v", 0.25, config.Temperature)
|
||||
}
|
||||
if config.Anthropic == nil || config.Anthropic.APIKey != "anthropic-key" || !config.Anthropic.UseBedrock {
|
||||
t.Fatalf("expected anthropic config to load, got %#v", config.Anthropic)
|
||||
}
|
||||
if config.OpenAI == nil || config.OpenAI.APIKey != "openai-key" || config.OpenAI.Organization != "org-test" || config.OpenAI.Project != "proj-test" {
|
||||
t.Fatalf("expected openai config to load, got %#v", config.OpenAI)
|
||||
}
|
||||
if config.OpenAI == nil || len(config.OpenAI.Models) != 2 || config.OpenAI.Models[0] != "gpt-test" || config.OpenAI.Models[1] != "gpt-mini-test" {
|
||||
t.Fatalf("expected openai models to load, got %#v", config.OpenAI)
|
||||
}
|
||||
if config.OpenAI == nil || config.OpenAI.UseResponsesAPI == nil || !*config.OpenAI.UseResponsesAPI {
|
||||
t.Fatalf("expected openai responses api flag to load")
|
||||
}
|
||||
if config.Google == nil || config.Google.APIKey != "google-key" || config.Google.Project != "vertex-project" || config.Google.Location != "us-central1" || !config.Google.SkipAuth {
|
||||
t.Fatalf("expected google config to load, got %#v", config.Google)
|
||||
}
|
||||
if config.OpenAICompat == nil || config.OpenAICompat.BaseURL != "http://127.0.0.1:8080/v1" {
|
||||
t.Fatalf("expected openai_compat config to load, got %#v", config.OpenAICompat)
|
||||
}
|
||||
if config.OpenRouter == nil || config.OpenRouter.APIKey != "openrouter-key" || config.OpenRouter.UserAgent != "sliver-test/1.0" {
|
||||
t.Fatalf("expected openrouter config to load, got %#v", config.OpenRouter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIConfigWritesDefault(t *testing.T) {
|
||||
t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
|
||||
|
||||
config := GetAIConfig()
|
||||
if config.Provider != "" || config.Model != "" || config.ThinkingLevel != "" {
|
||||
t.Fatalf("expected empty default ai selections, got %#v", config)
|
||||
}
|
||||
if config.SystemPrompt != defaultAISystemPrompt {
|
||||
t.Fatalf("expected default system prompt %q, got %q", defaultAISystemPrompt, config.SystemPrompt)
|
||||
}
|
||||
if config.Anthropic == nil || config.Google == nil || config.OpenAI == nil || config.OpenAICompat == nil || config.OpenRouter == nil {
|
||||
t.Fatalf("expected default provider configs to exist, got %#v", config)
|
||||
}
|
||||
if _, err := os.Stat(GetAIConfigPath()); err != nil {
|
||||
t.Fatalf("expected default ai config file to exist: %v", err)
|
||||
}
|
||||
data, err := os.ReadFile(GetAIConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read default ai config: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), "system_prompt:") {
|
||||
t.Fatalf("expected default ai config to include system_prompt field, got %q", string(data))
|
||||
}
|
||||
if strings.Contains(string(data), "\nmodel:") || strings.HasPrefix(string(data), "model:") {
|
||||
t.Fatalf("expected default ai config to omit legacy model field, got %q", string(data))
|
||||
}
|
||||
if !strings.Contains(string(data), "authorized security testing") {
|
||||
t.Fatalf("expected default ai config to include the default system prompt, got %q", string(data))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIConfigMigratesLegacySharedModelToProviderModels(t *testing.T) {
|
||||
t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
|
||||
|
||||
configPath := GetAIConfigPath()
|
||||
if err := os.MkdirAll(filepath.Dir(configPath), 0700); err != nil {
|
||||
t.Fatalf("failed to create config dir: %v", err)
|
||||
}
|
||||
|
||||
data := []byte(`provider: "openai"
|
||||
model: "gpt-test"
|
||||
openai:
|
||||
api_key: "openai-key"
|
||||
`)
|
||||
if err := os.WriteFile(configPath, data, 0600); err != nil {
|
||||
t.Fatalf("failed to write legacy ai config: %v", err)
|
||||
}
|
||||
|
||||
config := GetAIConfig()
|
||||
if config.Model != "" {
|
||||
t.Fatalf("expected legacy model field to be cleared, got %q", config.Model)
|
||||
}
|
||||
if config.OpenAI == nil || len(config.OpenAI.Models) != 1 || config.OpenAI.Models[0] != "gpt-test" {
|
||||
t.Fatalf("expected legacy model to migrate into openai models, got %#v", config.OpenAI)
|
||||
}
|
||||
|
||||
savedData, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read migrated ai config: %v", err)
|
||||
}
|
||||
if strings.Contains(string(savedData), "\nmodel:") || strings.HasPrefix(string(savedData), "model:") {
|
||||
t.Fatalf("expected migrated ai config to omit legacy model field, got %q", string(savedData))
|
||||
}
|
||||
if !strings.Contains(string(savedData), "models:") {
|
||||
t.Fatalf("expected migrated ai config to include provider models, got %q", string(savedData))
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerConfigMigratesEmbeddedAIConfigFromYAML(t *testing.T) {
|
||||
t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
|
||||
|
||||
serverConfigPath := GetServerConfigPath()
|
||||
if err := os.MkdirAll(filepath.Dir(serverConfigPath), 0700); err != nil {
|
||||
t.Fatalf("failed to create config dir: %v", err)
|
||||
}
|
||||
|
||||
data := []byte(`daemon_mode: true
|
||||
daemon:
|
||||
host: "127.0.0.1"
|
||||
port: 4444
|
||||
ai:
|
||||
provider: "openai"
|
||||
model: "gpt-test"
|
||||
system_prompt: "Stay concise."
|
||||
openai:
|
||||
api_key: "openai-key"
|
||||
`)
|
||||
if err := os.WriteFile(serverConfigPath, data, 0600); err != nil {
|
||||
t.Fatalf("failed to write yaml config: %v", err)
|
||||
}
|
||||
|
||||
config := GetServerConfig()
|
||||
if !config.DaemonMode || config.DaemonConfig == nil || config.DaemonConfig.Port != 4444 {
|
||||
t.Fatalf("expected daemon settings to load, got %#v", config.DaemonConfig)
|
||||
}
|
||||
if config.AI == nil || config.AI.Provider != "openai" || config.AI.SystemPrompt != "Stay concise." {
|
||||
t.Fatalf("expected embedded AI config to migrate, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.Model != "" {
|
||||
t.Fatalf("expected embedded legacy ai model field to be cleared, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.OpenAI == nil || len(config.AI.OpenAI.Models) != 1 || config.AI.OpenAI.Models[0] != "gpt-test" {
|
||||
t.Fatalf("expected embedded ai model to migrate into openai models, got %#v", config.AI)
|
||||
}
|
||||
if config.AI.OpenAI == nil || config.AI.OpenAI.APIKey != "openai-key" {
|
||||
t.Fatalf("expected embedded openai api key to migrate, got %#v", config.AI.OpenAI)
|
||||
}
|
||||
if _, err := os.Stat(GetAIConfigPath()); err != nil {
|
||||
t.Fatalf("expected migrated ai config file to exist: %v", err)
|
||||
}
|
||||
|
||||
serverData, err := os.ReadFile(GetServerConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read migrated server config: %v", err)
|
||||
}
|
||||
if strings.Contains(string(serverData), "\nai:") || strings.HasPrefix(string(serverData), "ai:") {
|
||||
t.Fatalf("expected migrated server config to drop embedded ai block, got %q", string(serverData))
|
||||
}
|
||||
}
|
||||
@@ -197,7 +197,7 @@ type ServerConfig struct {
|
||||
DaemonConfig *DaemonConfig `json:"daemon" yaml:"daemon"`
|
||||
Logs *LogConfig `json:"logs" yaml:"logs"`
|
||||
GRPC *GRPCConfig `json:"grpc" yaml:"grpc"`
|
||||
AI *AIConfig `json:"ai" yaml:"ai"`
|
||||
AI *AIConfig `json:"-" yaml:"-"`
|
||||
Watchtower *WatchTowerConfig `json:"watch_tower" yaml:"watch_tower"`
|
||||
GoProxy string `json:"go_proxy" yaml:"go_proxy"`
|
||||
HTTPDefaults *HttpDefaultConfig `json:"http_default" yaml:"http_default"`
|
||||
@@ -210,6 +210,13 @@ type ServerConfig struct {
|
||||
|
||||
// Save - Save config file to disk
|
||||
func (c *ServerConfig) Save() error {
|
||||
if c != nil {
|
||||
c.AI = normalizeAIConfig(c.AI)
|
||||
if err := c.AI.Save(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
configPath := GetServerConfigPath()
|
||||
configDir := filepath.Dir(configPath)
|
||||
if _, err := os.Stat(configDir); os.IsNotExist(err) {
|
||||
@@ -237,6 +244,7 @@ func GetServerConfig() *ServerConfig {
|
||||
legacyPath := getServerLegacyConfigPath()
|
||||
config := getDefaultServerConfig()
|
||||
migratedLegacy := false
|
||||
var migratedAIConfig *AIConfig
|
||||
if _, err := os.Stat(configPath); !os.IsNotExist(err) {
|
||||
data, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
@@ -248,6 +256,7 @@ func GetServerConfig() *ServerConfig {
|
||||
serverConfigLog.Errorf("Failed to parse config file %s", err)
|
||||
return config
|
||||
}
|
||||
migratedAIConfig = aiConfigFromYAML(data)
|
||||
} else if _, err := os.Stat(legacyPath); !os.IsNotExist(err) {
|
||||
data, err := os.ReadFile(legacyPath)
|
||||
if err != nil {
|
||||
@@ -259,6 +268,7 @@ func GetServerConfig() *ServerConfig {
|
||||
serverConfigLog.Errorf("Failed to parse legacy config file %s", err)
|
||||
return config
|
||||
}
|
||||
migratedAIConfig = aiConfigFromJSON(data)
|
||||
migratedLegacy = true
|
||||
serverConfigLog.Infof("Migrating legacy config %s to %s", legacyPath, configPath)
|
||||
} else {
|
||||
@@ -286,7 +296,7 @@ func GetServerConfig() *ServerConfig {
|
||||
defaultPermit := true
|
||||
config.GRPC.Keepalive.PermitWithoutStream = &defaultPermit
|
||||
}
|
||||
normalizeAIConfig(config)
|
||||
config.AI = getAIConfig(migratedAIConfig)
|
||||
|
||||
err := config.Save() // This updates the config with any missing fields
|
||||
if err != nil {
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
@@ -25,32 +26,6 @@ logs:
|
||||
grpc_unary_payloads: true
|
||||
grpc_stream_payloads: true
|
||||
tls_key_logger: true
|
||||
ai:
|
||||
provider: "openai"
|
||||
model: "gpt-test"
|
||||
thinking_level: "high"
|
||||
max_output_tokens: 2048
|
||||
temperature: 0.25
|
||||
anthropic:
|
||||
api_key: "anthropic-key"
|
||||
base_url: "https://api.anthropic.example"
|
||||
use_bedrock: true
|
||||
openai:
|
||||
api_key: "openai-key"
|
||||
base_url: "https://api.openai.example"
|
||||
organization: "org-test"
|
||||
project: "proj-test"
|
||||
use_responses_api: true
|
||||
google:
|
||||
api_key: "google-key"
|
||||
project: "vertex-project"
|
||||
location: "us-central1"
|
||||
skip_auth: true
|
||||
openai_compat:
|
||||
base_url: "http://127.0.0.1:8080/v1"
|
||||
openrouter:
|
||||
api_key: "openrouter-key"
|
||||
user_agent: "sliver-test/1.0"
|
||||
watch_tower:
|
||||
vt_api_key: "vt"
|
||||
xforce_api_key: "xforce"
|
||||
@@ -84,6 +59,38 @@ cxx:
|
||||
if err := os.WriteFile(configPath, data, 0600); err != nil {
|
||||
t.Fatalf("failed to write yaml config: %v", err)
|
||||
}
|
||||
aiPath := GetAIConfigPath()
|
||||
aiData := []byte(`provider: "openai"
|
||||
thinking_level: "high"
|
||||
max_output_tokens: 2048
|
||||
temperature: 0.25
|
||||
anthropic:
|
||||
api_key: "anthropic-key"
|
||||
base_url: "https://api.anthropic.example"
|
||||
use_bedrock: true
|
||||
openai:
|
||||
models:
|
||||
- "gpt-test"
|
||||
- "gpt-mini-test"
|
||||
api_key: "openai-key"
|
||||
base_url: "https://api.openai.example"
|
||||
organization: "org-test"
|
||||
project: "proj-test"
|
||||
use_responses_api: true
|
||||
google:
|
||||
api_key: "google-key"
|
||||
project: "vertex-project"
|
||||
location: "us-central1"
|
||||
skip_auth: true
|
||||
openai_compat:
|
||||
base_url: "http://127.0.0.1:8080/v1"
|
||||
openrouter:
|
||||
api_key: "openrouter-key"
|
||||
user_agent: "sliver-test/1.0"
|
||||
`)
|
||||
if err := os.WriteFile(aiPath, aiData, 0600); err != nil {
|
||||
t.Fatalf("failed to write ai yaml config: %v", err)
|
||||
}
|
||||
|
||||
config := GetServerConfig()
|
||||
if !config.DaemonMode {
|
||||
@@ -113,8 +120,11 @@ cxx:
|
||||
if config.AI == nil || config.AI.Provider != "openai" {
|
||||
t.Fatalf("expected ai provider %q, got %#v", "openai", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.Model != "gpt-test" {
|
||||
t.Fatalf("expected ai model %q, got %#v", "gpt-test", config.AI)
|
||||
if config.AI == nil || config.AI.Model != "" {
|
||||
t.Fatalf("expected legacy ai model field to remain empty, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.OpenAI == nil || len(config.AI.OpenAI.Models) != 2 || config.AI.OpenAI.Models[0] != "gpt-test" || config.AI.OpenAI.Models[1] != "gpt-mini-test" {
|
||||
t.Fatalf("expected ai openai models to load, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.ThinkingLevel != "high" {
|
||||
t.Fatalf("expected ai thinking level %q, got %#v", "high", config.AI)
|
||||
@@ -205,9 +215,29 @@ func TestServerConfigWritesDefault(t *testing.T) {
|
||||
if config.AI.Provider != "" || config.AI.Model != "" || config.AI.ThinkingLevel != "" {
|
||||
t.Fatalf("expected empty default ai selections, got %#v", config.AI)
|
||||
}
|
||||
if config.AI.SystemPrompt != defaultAISystemPrompt {
|
||||
t.Fatalf("expected default ai system prompt %q, got %q", defaultAISystemPrompt, config.AI.SystemPrompt)
|
||||
}
|
||||
if _, err := os.Stat(GetServerConfigPath()); err != nil {
|
||||
t.Fatalf("expected default config file to exist: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(GetAIConfigPath()); err != nil {
|
||||
t.Fatalf("expected default ai config file to exist: %v", err)
|
||||
}
|
||||
aiData, err := os.ReadFile(GetAIConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read default ai config: %v", err)
|
||||
}
|
||||
if strings.Contains(string(aiData), "\nmodel:") || strings.HasPrefix(string(aiData), "model:") {
|
||||
t.Fatalf("expected default ai config to omit legacy model field, got %q", string(aiData))
|
||||
}
|
||||
data, err := os.ReadFile(GetServerConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read default server config: %v", err)
|
||||
}
|
||||
if strings.Contains(string(data), "\nai:") || strings.HasPrefix(string(data), "ai:") {
|
||||
t.Fatalf("expected default server config to exclude ai block, got %q", string(data))
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerConfigMigratesLegacyJSON(t *testing.T) {
|
||||
@@ -325,9 +355,15 @@ func TestServerConfigMigratesLegacyJSON(t *testing.T) {
|
||||
if config.AI == nil || config.AI.OpenRouter == nil || config.AI.OpenRouter.APIKey != "openrouter-legacy" {
|
||||
t.Fatalf("expected legacy ai openrouter config to load")
|
||||
}
|
||||
if config.AI == nil || config.AI.Provider != "anthropic" || config.AI.Model != "claude-test" || config.AI.ThinkingLevel != "medium" || config.AI.MaxOutputTokens != 1024 {
|
||||
if config.AI == nil || config.AI.Provider != "anthropic" || config.AI.ThinkingLevel != "medium" || config.AI.MaxOutputTokens != 1024 {
|
||||
t.Fatalf("expected legacy ai selections to load, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.Model != "" {
|
||||
t.Fatalf("expected legacy ai model field to be cleared, got %#v", config.AI)
|
||||
}
|
||||
if config.AI == nil || config.AI.Anthropic == nil || len(config.AI.Anthropic.Models) != 1 || config.AI.Anthropic.Models[0] != "claude-test" {
|
||||
t.Fatalf("expected legacy model to migrate into anthropic models, got %#v", config.AI)
|
||||
}
|
||||
if config.Watchtower == nil || config.Watchtower.VTApiKey != "vt-legacy" {
|
||||
t.Fatalf("expected legacy watch_tower to load")
|
||||
}
|
||||
@@ -340,12 +376,22 @@ func TestServerConfigMigratesLegacyJSON(t *testing.T) {
|
||||
if _, err := os.Stat(GetServerConfigPath()); err != nil {
|
||||
t.Fatalf("expected migrated yaml config file to exist: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(GetAIConfigPath()); err != nil {
|
||||
t.Fatalf("expected migrated ai yaml config file to exist: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(legacyPath); !os.IsNotExist(err) {
|
||||
t.Fatalf("expected legacy config to be renamed: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(legacyBackupPath(legacyPath)); err != nil {
|
||||
t.Fatalf("expected legacy backup file to exist: %v", err)
|
||||
}
|
||||
serverData, err := os.ReadFile(GetServerConfigPath())
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read migrated server config: %v", err)
|
||||
}
|
||||
if strings.Contains(string(serverData), "\nai:") || strings.HasPrefix(string(serverData), "ai:") {
|
||||
t.Fatalf("expected migrated server config to exclude ai block, got %q", string(serverData))
|
||||
}
|
||||
}
|
||||
|
||||
func boolPtr(value bool) *bool {
|
||||
|
||||
@@ -42,7 +42,18 @@ func aiConfigCmd(_ *cobra.Command, args []string) {
|
||||
}
|
||||
|
||||
result := currentAIConfigFormResult(serverConfig)
|
||||
if err := forms.AIConfig(result); err != nil {
|
||||
originalProvider := ai.NormalizeProviderName(result.Provider)
|
||||
if err := forms.SelectAIProvider(result); err != nil {
|
||||
if errors.Is(err, forms.ErrUserAborted) {
|
||||
return
|
||||
}
|
||||
fmt.Printf(Warn+"AI configuration form failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
if ai.NormalizeProviderName(result.Provider) != originalProvider {
|
||||
result = currentAIConfigFormResultForProvider(serverConfig, result.Provider)
|
||||
}
|
||||
if err := forms.EditAIConfig(result); err != nil {
|
||||
if errors.Is(err, forms.ErrUserAborted) {
|
||||
return
|
||||
}
|
||||
@@ -56,11 +67,11 @@ func aiConfigCmd(_ *cobra.Command, args []string) {
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf(Info+"Saved AI configuration to %s\n", configs.GetServerConfigPath())
|
||||
fmt.Printf(Info+"Saved AI configuration to %s\n", configs.GetAIConfigPath())
|
||||
fmt.Printf(
|
||||
Info+"Provider=%s, model=%s, thinking=%s, api_key=%s\n",
|
||||
Info+"Provider=%s, default_model=%s, thinking=%s, api_key=%s\n",
|
||||
formatAIValue(serverConfig.AI.Provider, ai.ProviderOpenAI),
|
||||
formatAIValue(serverConfig.AI.Model, "provider default"),
|
||||
formatAIValue(aiConfigDefaultModel(serverConfig.AI, serverConfig.AI.Provider), "provider default"),
|
||||
formatAIValue(serverConfig.AI.ThinkingLevel, "provider default"),
|
||||
apiKeyStatus(serverConfig.AI.Provider, aiProviderConfig(serverConfig.AI, serverConfig.AI.Provider)),
|
||||
)
|
||||
@@ -68,19 +79,35 @@ func aiConfigCmd(_ *cobra.Command, args []string) {
|
||||
}
|
||||
|
||||
func currentAIConfigFormResult(serverConfig *configs.ServerConfig) *forms.AIConfigFormResult {
|
||||
return currentAIConfigFormResultForProvider(serverConfig, selectedAIProvider(serverConfigAIConfig(serverConfig)))
|
||||
}
|
||||
|
||||
func currentAIConfigFormResultForProvider(serverConfig *configs.ServerConfig, provider string) *forms.AIConfigFormResult {
|
||||
if serverConfig == nil {
|
||||
return &forms.AIConfigFormResult{Provider: ai.ProviderOpenAI, UseResponsesAPI: true}
|
||||
provider = ai.NormalizeProviderName(provider)
|
||||
if !ai.IsSupportedProvider(provider) {
|
||||
provider = ai.ProviderOpenAI
|
||||
}
|
||||
return &forms.AIConfigFormResult{
|
||||
Provider: provider,
|
||||
UseResponsesAPI: provider == ai.ProviderOpenAI,
|
||||
}
|
||||
}
|
||||
|
||||
provider := selectedAIProvider(serverConfig.AI)
|
||||
provider = ai.NormalizeProviderName(provider)
|
||||
if !ai.IsSupportedProvider(provider) {
|
||||
provider = selectedAIProvider(serverConfig.AI)
|
||||
}
|
||||
providerConfig := aiProviderConfig(serverConfig.AI, provider)
|
||||
|
||||
result := &forms.AIConfigFormResult{
|
||||
Provider: provider,
|
||||
Provider: provider,
|
||||
UseResponsesAPI: provider == ai.ProviderOpenAI,
|
||||
}
|
||||
if serverConfig.AI != nil {
|
||||
result.Model = strings.TrimSpace(serverConfig.AI.Model)
|
||||
result.Models = aiConfigFormModels(serverConfig.AI, provider)
|
||||
result.ThinkingLevel = strings.TrimSpace(serverConfig.AI.ThinkingLevel)
|
||||
result.SystemPrompt = strings.TrimSpace(serverConfig.AI.SystemPrompt)
|
||||
}
|
||||
if providerConfig != nil {
|
||||
result.APIKey = strings.TrimSpace(providerConfig.APIKey)
|
||||
@@ -99,6 +126,13 @@ func currentAIConfigFormResult(serverConfig *configs.ServerConfig) *forms.AIConf
|
||||
return result
|
||||
}
|
||||
|
||||
func serverConfigAIConfig(serverConfig *configs.ServerConfig) *configs.AIConfig {
|
||||
if serverConfig == nil {
|
||||
return nil
|
||||
}
|
||||
return serverConfig.AI
|
||||
}
|
||||
|
||||
func applyAIConfigFormResult(serverConfig *configs.ServerConfig, result *forms.AIConfigFormResult) {
|
||||
if serverConfig == nil || result == nil {
|
||||
return
|
||||
@@ -109,13 +143,15 @@ func applyAIConfigFormResult(serverConfig *configs.ServerConfig, result *forms.A
|
||||
|
||||
provider := ai.NormalizeProviderName(result.Provider)
|
||||
serverConfig.AI.Provider = provider
|
||||
serverConfig.AI.Model = strings.TrimSpace(result.Model)
|
||||
serverConfig.AI.Model = ""
|
||||
serverConfig.AI.ThinkingLevel = strings.ToLower(strings.TrimSpace(result.ThinkingLevel))
|
||||
serverConfig.AI.SystemPrompt = strings.TrimSpace(result.SystemPrompt)
|
||||
|
||||
providerConfig := aiProviderConfig(serverConfig.AI, provider)
|
||||
if providerConfig == nil {
|
||||
return
|
||||
}
|
||||
providerConfig.Models = append([]string(nil), result.Models...)
|
||||
providerConfig.APIKey = strings.TrimSpace(result.APIKey)
|
||||
providerConfig.BaseURL = strings.TrimSpace(result.BaseURL)
|
||||
providerConfig.UserAgent = strings.TrimSpace(result.UserAgent)
|
||||
@@ -127,6 +163,27 @@ func applyAIConfigFormResult(serverConfig *configs.ServerConfig, result *forms.A
|
||||
providerConfig.UseBedrock = result.UseBedrock
|
||||
}
|
||||
|
||||
func aiConfigFormModels(aiConfig *configs.AIConfig, provider string) []string {
|
||||
if aiConfig == nil {
|
||||
return nil
|
||||
}
|
||||
if providerConfig := aiProviderConfig(aiConfig, provider); providerConfig != nil && len(providerConfig.Models) > 0 {
|
||||
return append([]string(nil), providerConfig.Models...)
|
||||
}
|
||||
if legacyModel := strings.TrimSpace(aiConfig.Model); legacyModel != "" {
|
||||
return []string{legacyModel}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func aiConfigDefaultModel(aiConfig *configs.AIConfig, provider string) string {
|
||||
models := aiConfigFormModels(aiConfig, provider)
|
||||
if len(models) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(models[0])
|
||||
}
|
||||
|
||||
func boolPtr(value bool) *bool {
|
||||
return &value
|
||||
}
|
||||
|
||||
@@ -31,12 +31,88 @@ func TestCurrentAIConfigFormResultUsesConfiguredProviderFallback(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCurrentAIConfigFormResultForProviderLoadsRequestedProviderSettings(t *testing.T) {
|
||||
serverConfig := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ai.ProviderAnthropic,
|
||||
ThinkingLevel: "medium",
|
||||
SystemPrompt: "Stay concise.",
|
||||
Anthropic: &configs.AIProviderConfig{APIKey: "anthropic-key"},
|
||||
Google: &configs.AIProviderConfig{},
|
||||
OpenAI: &configs.AIProviderConfig{
|
||||
Models: []string{"gpt-test", "gpt-test-mini"},
|
||||
APIKey: "openai-key",
|
||||
BaseURL: "https://api.openai.test",
|
||||
Organization: "org-test",
|
||||
Project: "proj-test",
|
||||
UseResponsesAPI: boolPtr(true),
|
||||
},
|
||||
OpenAICompat: &configs.AIProviderConfig{},
|
||||
OpenRouter: &configs.AIProviderConfig{},
|
||||
},
|
||||
}
|
||||
|
||||
result := currentAIConfigFormResultForProvider(serverConfig, ai.ProviderOpenAI)
|
||||
if result.Provider != ai.ProviderOpenAI {
|
||||
t.Fatalf("expected %q provider, got %q", ai.ProviderOpenAI, result.Provider)
|
||||
}
|
||||
if len(result.Models) != 2 || result.Models[0] != "gpt-test" || result.Models[1] != "gpt-test-mini" {
|
||||
t.Fatalf("expected provider models to load, got %#v", result.Models)
|
||||
}
|
||||
if result.SystemPrompt != "Stay concise." {
|
||||
t.Fatalf("expected shared system prompt, got %q", result.SystemPrompt)
|
||||
}
|
||||
if result.APIKey != "openai-key" {
|
||||
t.Fatalf("expected openai API key, got %q", result.APIKey)
|
||||
}
|
||||
if result.BaseURL != "https://api.openai.test" {
|
||||
t.Fatalf("expected openai base url, got %q", result.BaseURL)
|
||||
}
|
||||
if result.Organization != "org-test" {
|
||||
t.Fatalf("expected openai organization, got %q", result.Organization)
|
||||
}
|
||||
if result.Project != "proj-test" {
|
||||
t.Fatalf("expected openai project, got %q", result.Project)
|
||||
}
|
||||
if !result.UseResponsesAPI {
|
||||
t.Fatal("expected openai responses api to remain enabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCurrentAIConfigFormResultForProviderFallsBackToLegacySharedModel(t *testing.T) {
|
||||
serverConfig := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ai.ProviderAnthropic,
|
||||
Model: "legacy-model",
|
||||
ThinkingLevel: "medium",
|
||||
Anthropic: &configs.AIProviderConfig{APIKey: "anthropic-key"},
|
||||
OpenAI: &configs.AIProviderConfig{APIKey: "openai-key"},
|
||||
},
|
||||
}
|
||||
|
||||
result := currentAIConfigFormResultForProvider(serverConfig, ai.ProviderOpenAI)
|
||||
if len(result.Models) != 1 || result.Models[0] != "legacy-model" {
|
||||
t.Fatalf("expected legacy shared model fallback, got %#v", result.Models)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCurrentAIConfigFormResultForProviderDefaultsResponsesAPIForOpenAI(t *testing.T) {
|
||||
result := currentAIConfigFormResultForProvider(&configs.ServerConfig{}, ai.ProviderOpenAI)
|
||||
if result.Provider != ai.ProviderOpenAI {
|
||||
t.Fatalf("expected %q provider, got %q", ai.ProviderOpenAI, result.Provider)
|
||||
}
|
||||
if !result.UseResponsesAPI {
|
||||
t.Fatal("expected openai to default to responses API")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyAIConfigFormResultUpdatesOnlySelectedProvider(t *testing.T) {
|
||||
serverConfig := &configs.ServerConfig{
|
||||
AI: &configs.AIConfig{
|
||||
Provider: ai.ProviderAnthropic,
|
||||
Model: "old-model",
|
||||
ThinkingLevel: "low",
|
||||
SystemPrompt: "Existing default",
|
||||
Anthropic: &configs.AIProviderConfig{APIKey: "anthropic-key"},
|
||||
Google: &configs.AIProviderConfig{},
|
||||
OpenAI: &configs.AIProviderConfig{APIKey: "openai-key", BaseURL: "https://api.openai.test"},
|
||||
@@ -47,8 +123,9 @@ func TestApplyAIConfigFormResultUpdatesOnlySelectedProvider(t *testing.T) {
|
||||
|
||||
applyAIConfigFormResult(serverConfig, &forms.AIConfigFormResult{
|
||||
Provider: ai.ProviderOpenAI,
|
||||
Model: "gpt-test",
|
||||
Models: []string{"gpt-test", "gpt-test-mini"},
|
||||
ThinkingLevel: "high",
|
||||
SystemPrompt: "Stay concise.",
|
||||
APIKey: "new-openai-key",
|
||||
BaseURL: "https://override.openai.test",
|
||||
Organization: "org-test",
|
||||
@@ -59,12 +136,18 @@ func TestApplyAIConfigFormResultUpdatesOnlySelectedProvider(t *testing.T) {
|
||||
if serverConfig.AI.Provider != ai.ProviderOpenAI {
|
||||
t.Fatalf("expected provider %q, got %q", ai.ProviderOpenAI, serverConfig.AI.Provider)
|
||||
}
|
||||
if serverConfig.AI.Model != "gpt-test" {
|
||||
t.Fatalf("expected model %q, got %q", "gpt-test", serverConfig.AI.Model)
|
||||
if serverConfig.AI.Model != "" {
|
||||
t.Fatalf("expected legacy model field to be cleared, got %q", serverConfig.AI.Model)
|
||||
}
|
||||
if len(serverConfig.AI.OpenAI.Models) != 2 || serverConfig.AI.OpenAI.Models[0] != "gpt-test" || serverConfig.AI.OpenAI.Models[1] != "gpt-test-mini" {
|
||||
t.Fatalf("expected models to update, got %#v", serverConfig.AI.OpenAI.Models)
|
||||
}
|
||||
if serverConfig.AI.ThinkingLevel != "high" {
|
||||
t.Fatalf("expected thinking level %q, got %q", "high", serverConfig.AI.ThinkingLevel)
|
||||
}
|
||||
if serverConfig.AI.SystemPrompt != "Stay concise." {
|
||||
t.Fatalf("expected system prompt %q, got %q", "Stay concise.", serverConfig.AI.SystemPrompt)
|
||||
}
|
||||
if serverConfig.AI.OpenAI.APIKey != "new-openai-key" {
|
||||
t.Fatalf("expected updated openai api key, got %q", serverConfig.AI.OpenAI.APIKey)
|
||||
}
|
||||
|
||||
@@ -12,8 +12,9 @@ import (
|
||||
// AIConfigFormResult captures the server-side AI configuration collected from the form.
|
||||
type AIConfigFormResult struct {
|
||||
Provider string
|
||||
Model string
|
||||
Models []string
|
||||
ThinkingLevel string
|
||||
SystemPrompt string
|
||||
APIKey string
|
||||
BaseURL string
|
||||
UserAgent string
|
||||
@@ -25,12 +26,24 @@ type AIConfigFormResult struct {
|
||||
UseBedrock bool
|
||||
}
|
||||
|
||||
// AIConfig prompts for the server-side AI configuration stored in server.yaml.
|
||||
// AIConfig prompts for the server-side AI configuration stored in ai.yaml.
|
||||
func AIConfig(result *AIConfigFormResult) error {
|
||||
if result == nil {
|
||||
return errors.New("AI config result is required")
|
||||
}
|
||||
|
||||
if err := SelectAIProvider(result); err != nil {
|
||||
return err
|
||||
}
|
||||
return EditAIConfig(result)
|
||||
}
|
||||
|
||||
// SelectAIProvider prompts for the AI provider to edit.
|
||||
func SelectAIProvider(result *AIConfigFormResult) error {
|
||||
if result == nil {
|
||||
return errors.New("AI config result is required")
|
||||
}
|
||||
|
||||
normalizeAIConfigResult(result)
|
||||
|
||||
providerOptions := []huh.Option[string]{
|
||||
@@ -40,19 +53,12 @@ func AIConfig(result *AIConfigFormResult) error {
|
||||
huh.NewOption("OpenAI-Compatible", ai.ProviderOpenAICompat),
|
||||
huh.NewOption("OpenRouter", ai.ProviderOpenRouter),
|
||||
}
|
||||
thinkingOptions := []huh.Option[string]{
|
||||
huh.NewOption("Provider default", ""),
|
||||
huh.NewOption("Disabled", "disabled"),
|
||||
huh.NewOption("Low", "low"),
|
||||
huh.NewOption("Medium", "medium"),
|
||||
huh.NewOption("High", "high"),
|
||||
}
|
||||
|
||||
providerForm := huh.NewForm(
|
||||
providerForm := newConsoleForm(
|
||||
huh.NewGroup(
|
||||
huh.NewSelect[string]().
|
||||
Title("AI provider").
|
||||
Description("Choose the default AI provider stored in server.yaml.").
|
||||
Description("Choose the default AI provider stored in ai.yaml.").
|
||||
Options(providerOptions...).
|
||||
Height(listHeight(len(providerOptions))).
|
||||
Value(&result.Provider).
|
||||
@@ -64,11 +70,32 @@ func AIConfig(result *AIConfigFormResult) error {
|
||||
}),
|
||||
),
|
||||
)
|
||||
if err := providerForm.Run(); err != nil {
|
||||
if err := runConsoleForm(providerForm); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
detailsForm := huh.NewForm(
|
||||
normalizeAIConfigResult(result)
|
||||
return nil
|
||||
}
|
||||
|
||||
// EditAIConfig prompts for the provider-independent and provider-specific AI settings.
|
||||
func EditAIConfig(result *AIConfigFormResult) error {
|
||||
if result == nil {
|
||||
return errors.New("AI config result is required")
|
||||
}
|
||||
|
||||
normalizeAIConfigResult(result)
|
||||
|
||||
thinkingOptions := []huh.Option[string]{
|
||||
huh.NewOption("Provider default", ""),
|
||||
huh.NewOption("Disabled", "disabled"),
|
||||
huh.NewOption("Low", "low"),
|
||||
huh.NewOption("Medium", "medium"),
|
||||
huh.NewOption("High", "high"),
|
||||
}
|
||||
modelsValue := formatAIModelList(result.Models)
|
||||
|
||||
detailsForm := newConsoleForm(
|
||||
huh.NewGroup(
|
||||
huh.NewNote().
|
||||
Title("Server AI configuration").
|
||||
@@ -79,16 +106,21 @@ func AIConfig(result *AIConfigFormResult) error {
|
||||
)
|
||||
}, &result.Provider),
|
||||
huh.NewInput().
|
||||
Title("Model").
|
||||
Description("Optional default model identifier. Leave blank to let the provider choose.").
|
||||
Title("Models").
|
||||
Description("Optional ordered model identifiers. Separate entries with commas. The first entry is the default.").
|
||||
Placeholder("provider default").
|
||||
Value(&result.Model),
|
||||
Value(&modelsValue),
|
||||
huh.NewSelect[string]().
|
||||
Title("Thinking level").
|
||||
Description("Optional reasoning/thinking level to store in the ai block.").
|
||||
Description("Optional reasoning/thinking level to store in ai.yaml.").
|
||||
Options(thinkingOptions...).
|
||||
Height(listHeight(len(thinkingOptions))).
|
||||
Value(&result.ThinkingLevel),
|
||||
huh.NewInput().
|
||||
Title("System prompt").
|
||||
Description("Optional default system prompt for new AI conversations.").
|
||||
Placeholder("leave blank for no default system prompt").
|
||||
Value(&result.SystemPrompt),
|
||||
huh.NewInput().
|
||||
TitleFunc(func() string {
|
||||
return fmt.Sprintf("%s API key", providerDisplayName(result.Provider))
|
||||
@@ -116,9 +148,10 @@ func AIConfig(result *AIConfigFormResult) error {
|
||||
Value(&result.UserAgent),
|
||||
),
|
||||
)
|
||||
if err := detailsForm.Run(); err != nil {
|
||||
if err := runConsoleForm(detailsForm); err != nil {
|
||||
return err
|
||||
}
|
||||
result.Models = parseAIModelList(modelsValue)
|
||||
|
||||
if err := providerSpecificAIConfig(result); err != nil {
|
||||
return err
|
||||
@@ -142,7 +175,7 @@ func providerSpecificAIConfig(result *AIConfigFormResult) error {
|
||||
}
|
||||
|
||||
func openAIAIConfig(result *AIConfigFormResult) error {
|
||||
form := huh.NewForm(
|
||||
form := newConsoleForm(
|
||||
huh.NewGroup(
|
||||
huh.NewInput().
|
||||
Title("Organization").
|
||||
@@ -160,11 +193,11 @@ func openAIAIConfig(result *AIConfigFormResult) error {
|
||||
Value(&result.UseResponsesAPI),
|
||||
),
|
||||
)
|
||||
return form.Run()
|
||||
return runConsoleForm(form)
|
||||
}
|
||||
|
||||
func anthropicAIConfig(result *AIConfigFormResult) error {
|
||||
form := huh.NewForm(
|
||||
form := newConsoleForm(
|
||||
huh.NewGroup(
|
||||
huh.NewInput().
|
||||
Title("Vertex project").
|
||||
@@ -186,11 +219,11 @@ func anthropicAIConfig(result *AIConfigFormResult) error {
|
||||
Value(&result.UseBedrock),
|
||||
),
|
||||
)
|
||||
return form.Run()
|
||||
return runConsoleForm(form)
|
||||
}
|
||||
|
||||
func googleAIConfig(result *AIConfigFormResult) error {
|
||||
form := huh.NewForm(
|
||||
form := newConsoleForm(
|
||||
huh.NewGroup(
|
||||
huh.NewInput().
|
||||
Title("Vertex project").
|
||||
@@ -208,7 +241,7 @@ func googleAIConfig(result *AIConfigFormResult) error {
|
||||
Value(&result.SkipAuth),
|
||||
),
|
||||
)
|
||||
return form.Run()
|
||||
return runConsoleForm(form)
|
||||
}
|
||||
|
||||
func normalizeAIConfigResult(result *AIConfigFormResult) {
|
||||
@@ -219,8 +252,9 @@ func normalizeAIConfigResult(result *AIConfigFormResult) {
|
||||
if !ai.IsSupportedProvider(result.Provider) {
|
||||
result.Provider = ai.ProviderOpenAI
|
||||
}
|
||||
result.Model = strings.TrimSpace(result.Model)
|
||||
result.Models = parseAIModelList(formatAIModelList(result.Models))
|
||||
result.ThinkingLevel = strings.ToLower(strings.TrimSpace(result.ThinkingLevel))
|
||||
result.SystemPrompt = strings.TrimSpace(result.SystemPrompt)
|
||||
switch result.ThinkingLevel {
|
||||
case "", "disabled", "low", "medium", "high":
|
||||
default:
|
||||
@@ -254,6 +288,37 @@ func normalizeAIConfigResult(result *AIConfigFormResult) {
|
||||
}
|
||||
}
|
||||
|
||||
func formatAIModelList(models []string) string {
|
||||
return strings.Join(parseAIModelList(strings.Join(models, ",")), ", ")
|
||||
}
|
||||
|
||||
func parseAIModelList(value string) []string {
|
||||
fields := strings.FieldsFunc(value, func(r rune) bool {
|
||||
return r == ',' || r == '\n' || r == '\r'
|
||||
})
|
||||
if len(fields) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(fields))
|
||||
seen := make(map[string]struct{}, len(fields))
|
||||
for _, field := range fields {
|
||||
field = strings.TrimSpace(field)
|
||||
if field == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[field]; exists {
|
||||
continue
|
||||
}
|
||||
seen[field] = struct{}{}
|
||||
models = append(models, field)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
func providerDisplayName(provider string) string {
|
||||
switch ai.NormalizeProviderName(provider) {
|
||||
case ai.ProviderAnthropic:
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
package forms
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"charm.land/huh/v2"
|
||||
)
|
||||
|
||||
func newConsoleForm(groups ...*huh.Group) *huh.Form {
|
||||
return huh.NewForm(groups...)
|
||||
}
|
||||
|
||||
func runConsoleForm(form *huh.Form) error {
|
||||
if form == nil {
|
||||
return errors.New("form is required")
|
||||
}
|
||||
|
||||
ttyIn, ttyOut, err := tea.OpenTTY()
|
||||
if err == nil {
|
||||
if ttyIn != nil {
|
||||
defer ttyIn.Close()
|
||||
}
|
||||
if ttyOut != nil && (ttyIn == nil || ttyOut.Fd() != ttyIn.Fd()) {
|
||||
defer ttyOut.Close()
|
||||
}
|
||||
if ttyIn != nil {
|
||||
form.WithInput(ttyIn)
|
||||
}
|
||||
if ttyOut != nil {
|
||||
form.WithOutput(ttyOut)
|
||||
}
|
||||
}
|
||||
|
||||
return form.Run()
|
||||
}
|
||||
@@ -35,8 +35,8 @@ func SelectOperator(title string, operators []string, value *string) error {
|
||||
return nil
|
||||
})
|
||||
|
||||
form := huh.NewForm(huh.NewGroup(field))
|
||||
return form.Run()
|
||||
form := newConsoleForm(huh.NewGroup(field))
|
||||
return runConsoleForm(form)
|
||||
}
|
||||
|
||||
func ensureSelectedValue(options []string, value *string) {
|
||||
|
||||
+99
-4
@@ -1407,7 +1407,59 @@ func AIConversationByID(id string, operatorName string, includeMessages bool) (*
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conversation.ToProtobuf(), nil
|
||||
pbConversation := conversation.ToProtobuf()
|
||||
if includeMessages {
|
||||
pbConversation = aiConversationWithSystemPromptMessage(pbConversation)
|
||||
}
|
||||
return pbConversation, nil
|
||||
}
|
||||
|
||||
func aiConversationWithSystemPromptMessage(conversation *clientpb.AIConversation) *clientpb.AIConversation {
|
||||
if conversation == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
systemPrompt := strings.TrimSpace(conversation.GetSystemPrompt())
|
||||
if systemPrompt == "" {
|
||||
return conversation
|
||||
}
|
||||
for _, message := range conversation.GetMessages() {
|
||||
if message == nil {
|
||||
continue
|
||||
}
|
||||
if message.GetKind() != clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT {
|
||||
continue
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(message.GetRole()), "system") &&
|
||||
strings.TrimSpace(message.GetContent()) == systemPrompt {
|
||||
return conversation
|
||||
}
|
||||
}
|
||||
|
||||
includeInContext := true
|
||||
systemMessageID := "system-prompt"
|
||||
if conversationID := strings.TrimSpace(conversation.GetID()); conversationID != "" {
|
||||
systemMessageID = conversationID + "-system-prompt"
|
||||
}
|
||||
systemMessage := &clientpb.AIConversationMessage{
|
||||
ID: systemMessageID,
|
||||
ConversationID: strings.TrimSpace(conversation.GetID()),
|
||||
CreatedAt: conversation.GetCreatedAt(),
|
||||
UpdatedAt: conversation.GetUpdatedAt(),
|
||||
OperatorName: strings.TrimSpace(conversation.GetOperatorName()),
|
||||
Provider: strings.TrimSpace(conversation.GetProvider()),
|
||||
Model: strings.TrimSpace(conversation.GetModel()),
|
||||
Sequence: 0,
|
||||
Role: "system",
|
||||
Content: systemPrompt,
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
State: clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED,
|
||||
IncludeInContext: &includeInContext,
|
||||
}
|
||||
|
||||
conversation.Messages = append([]*clientpb.AIConversationMessage{systemMessage}, conversation.GetMessages()...)
|
||||
return conversation
|
||||
}
|
||||
|
||||
// SaveAIConversation - Create or update an AI conversation thread.
|
||||
@@ -1435,9 +1487,19 @@ func SaveAIConversation(conversation *clientpb.AIConversation, operatorName stri
|
||||
}
|
||||
existing.Provider = dbConversation.Provider
|
||||
existing.Model = dbConversation.Model
|
||||
existing.ThinkingLevel = dbConversation.ThinkingLevel
|
||||
existing.Title = dbConversation.Title
|
||||
existing.Summary = dbConversation.Summary
|
||||
existing.SystemPrompt = dbConversation.SystemPrompt
|
||||
existing.ActiveTurnID = dbConversation.ActiveTurnID
|
||||
existing.TurnState = dbConversation.TurnState
|
||||
existing.TargetSessionID = dbConversation.TargetSessionID
|
||||
existing.TargetBeaconID = dbConversation.TargetBeaconID
|
||||
existing.ContextInputTokens = dbConversation.ContextInputTokens
|
||||
existing.ContextOutputTokens = dbConversation.ContextOutputTokens
|
||||
existing.ContextTotalTokens = dbConversation.ContextTotalTokens
|
||||
existing.ContextWindowTokens = dbConversation.ContextWindowTokens
|
||||
existing.ContextWindowTokensEstimated = dbConversation.ContextWindowTokensEstimated
|
||||
|
||||
if err := dbSession.Save(existing).Error; err != nil {
|
||||
return nil, err
|
||||
@@ -1509,6 +1571,20 @@ func SaveAIConversationMessage(message *clientpb.AIConversationMessage, operator
|
||||
dbSession := Session()
|
||||
var savedMessage *models.AIConversationMessage
|
||||
err = dbSession.Transaction(func(tx *gorm.DB) error {
|
||||
if dbMessage.ID == uuid.Nil && strings.TrimSpace(dbMessage.ItemID) != "" {
|
||||
existingByItem := &models.AIConversationMessage{}
|
||||
err := tx.Where("conversation_id = ?", dbMessage.ConversationID).
|
||||
Where("item_id = ?", strings.TrimSpace(dbMessage.ItemID)).
|
||||
First(existingByItem).Error
|
||||
if err != nil {
|
||||
if !errors.Is(err, ErrRecordNotFound) {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
dbMessage.ID = existingByItem.ID
|
||||
}
|
||||
}
|
||||
|
||||
if dbMessage.ID == uuid.Nil {
|
||||
if dbMessage.Sequence == 0 {
|
||||
lastMessage := &models.AIConversationMessage{}
|
||||
@@ -1539,15 +1615,34 @@ func SaveAIConversationMessage(message *clientpb.AIConversationMessage, operator
|
||||
if dbMessage.OperatorName != "" {
|
||||
existing.OperatorName = dbMessage.OperatorName
|
||||
}
|
||||
existing.Provider = dbMessage.Provider
|
||||
existing.Model = dbMessage.Model
|
||||
if dbMessage.Provider != "" {
|
||||
existing.Provider = dbMessage.Provider
|
||||
}
|
||||
if dbMessage.Model != "" {
|
||||
existing.Model = dbMessage.Model
|
||||
}
|
||||
if dbMessage.Sequence != 0 {
|
||||
existing.Sequence = dbMessage.Sequence
|
||||
}
|
||||
existing.Role = dbMessage.Role
|
||||
if dbMessage.Role != "" {
|
||||
existing.Role = dbMessage.Role
|
||||
}
|
||||
existing.Content = dbMessage.Content
|
||||
existing.ProviderMessageID = dbMessage.ProviderMessageID
|
||||
existing.FinishReason = dbMessage.FinishReason
|
||||
existing.Kind = dbMessage.Kind
|
||||
existing.Visibility = dbMessage.Visibility
|
||||
existing.State = dbMessage.State
|
||||
existing.TurnID = dbMessage.TurnID
|
||||
existing.ItemID = dbMessage.ItemID
|
||||
existing.ToolCallID = dbMessage.ToolCallID
|
||||
existing.ToolName = dbMessage.ToolName
|
||||
existing.ToolArguments = dbMessage.ToolArguments
|
||||
existing.ToolResult = dbMessage.ToolResult
|
||||
existing.ErrorText = dbMessage.ErrorText
|
||||
if dbMessage.IncludeInContext != nil {
|
||||
existing.IncludeInContext = dbMessage.IncludeInContext
|
||||
}
|
||||
|
||||
if err := tx.Save(existing).Error; err != nil {
|
||||
return err
|
||||
|
||||
+120
-26
@@ -28,15 +28,25 @@ import (
|
||||
|
||||
// AIConversation - A server-side AI conversation thread.
|
||||
type AIConversation struct {
|
||||
ID uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`
|
||||
CreatedAt time.Time `gorm:"->;<-:create;"`
|
||||
UpdatedAt time.Time
|
||||
OperatorName string `gorm:"index"`
|
||||
Provider string `gorm:"index"`
|
||||
Model string
|
||||
Title string
|
||||
Summary string
|
||||
SystemPrompt string `gorm:"type:text;"`
|
||||
ID uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`
|
||||
CreatedAt time.Time `gorm:"->;<-:create;"`
|
||||
UpdatedAt time.Time
|
||||
OperatorName string `gorm:"index"`
|
||||
Provider string `gorm:"index"`
|
||||
Model string
|
||||
ThinkingLevel string
|
||||
Title string
|
||||
Summary string
|
||||
SystemPrompt string `gorm:"type:text;"`
|
||||
ActiveTurnID string `gorm:"index"`
|
||||
TurnState int32
|
||||
TargetSessionID string `gorm:"index"`
|
||||
TargetBeaconID string `gorm:"index"`
|
||||
ContextInputTokens int64
|
||||
ContextOutputTokens int64
|
||||
ContextTotalTokens int64
|
||||
ContextWindowTokens int64
|
||||
ContextWindowTokensEstimated bool
|
||||
|
||||
Messages []AIConversationMessage `gorm:"foreignKey:ConversationID;constraint:OnDelete:CASCADE;"`
|
||||
}
|
||||
@@ -61,16 +71,28 @@ func (a *AIConversation) ToProtobuf() *clientpb.AIConversation {
|
||||
}
|
||||
|
||||
return &clientpb.AIConversation{
|
||||
ID: a.ID.String(),
|
||||
CreatedAt: a.CreatedAt.Unix(),
|
||||
UpdatedAt: a.UpdatedAt.Unix(),
|
||||
OperatorName: a.OperatorName,
|
||||
Provider: a.Provider,
|
||||
Model: a.Model,
|
||||
Title: a.Title,
|
||||
Summary: a.Summary,
|
||||
SystemPrompt: a.SystemPrompt,
|
||||
Messages: messages,
|
||||
ID: a.ID.String(),
|
||||
CreatedAt: a.CreatedAt.Unix(),
|
||||
UpdatedAt: a.UpdatedAt.Unix(),
|
||||
OperatorName: a.OperatorName,
|
||||
Provider: a.Provider,
|
||||
Model: a.Model,
|
||||
ThinkingLevel: a.ThinkingLevel,
|
||||
Title: a.Title,
|
||||
Summary: a.Summary,
|
||||
SystemPrompt: a.SystemPrompt,
|
||||
Messages: messages,
|
||||
ActiveTurnID: a.ActiveTurnID,
|
||||
TurnState: clientpb.AIConversationTurnState(a.TurnState),
|
||||
TargetSessionID: a.TargetSessionID,
|
||||
TargetBeaconID: a.TargetBeaconID,
|
||||
ContextWindowUsage: aiContextWindowUsageToProtobuf(
|
||||
a.ContextInputTokens,
|
||||
a.ContextOutputTokens,
|
||||
a.ContextTotalTokens,
|
||||
a.ContextWindowTokens,
|
||||
a.ContextWindowTokensEstimated,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,13 +105,36 @@ func AIConversationFromProtobuf(pbConversation *clientpb.AIConversation) *AIConv
|
||||
id, _ := uuid.FromString(pbConversation.ID)
|
||||
|
||||
return &AIConversation{
|
||||
ID: id,
|
||||
OperatorName: pbConversation.OperatorName,
|
||||
Provider: pbConversation.Provider,
|
||||
Model: pbConversation.Model,
|
||||
Title: pbConversation.Title,
|
||||
Summary: pbConversation.Summary,
|
||||
SystemPrompt: pbConversation.SystemPrompt,
|
||||
ID: id,
|
||||
OperatorName: pbConversation.OperatorName,
|
||||
Provider: pbConversation.Provider,
|
||||
Model: pbConversation.Model,
|
||||
ThinkingLevel: pbConversation.ThinkingLevel,
|
||||
Title: pbConversation.Title,
|
||||
Summary: pbConversation.Summary,
|
||||
SystemPrompt: pbConversation.SystemPrompt,
|
||||
ActiveTurnID: pbConversation.ActiveTurnID,
|
||||
TurnState: int32(pbConversation.TurnState),
|
||||
TargetSessionID: pbConversation.TargetSessionID,
|
||||
TargetBeaconID: pbConversation.TargetBeaconID,
|
||||
ContextInputTokens: pbConversation.GetContextWindowUsage().GetInputTokens(),
|
||||
ContextOutputTokens: pbConversation.GetContextWindowUsage().GetOutputTokens(),
|
||||
ContextTotalTokens: pbConversation.GetContextWindowUsage().GetTotalTokens(),
|
||||
ContextWindowTokens: pbConversation.GetContextWindowUsage().GetContextWindowTokens(),
|
||||
ContextWindowTokensEstimated: pbConversation.GetContextWindowUsage().GetContextWindowTokensEstimated(),
|
||||
}
|
||||
}
|
||||
|
||||
func aiContextWindowUsageToProtobuf(inputTokens, outputTokens, totalTokens, windowTokens int64, estimated bool) *clientpb.AIContextWindowUsage {
|
||||
if inputTokens == 0 && outputTokens == 0 && totalTokens == 0 && windowTokens == 0 && !estimated {
|
||||
return nil
|
||||
}
|
||||
return &clientpb.AIContextWindowUsage{
|
||||
InputTokens: inputTokens,
|
||||
OutputTokens: outputTokens,
|
||||
TotalTokens: totalTokens,
|
||||
ContextWindowTokens: windowTokens,
|
||||
ContextWindowTokensEstimated: estimated,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -107,6 +152,17 @@ type AIConversationMessage struct {
|
||||
Content string `gorm:"type:text;"`
|
||||
ProviderMessageID string
|
||||
FinishReason string
|
||||
Kind int32
|
||||
Visibility int32
|
||||
State int32
|
||||
TurnID string `gorm:"index"`
|
||||
ItemID string `gorm:"index"`
|
||||
ToolCallID string `gorm:"index"`
|
||||
ToolName string
|
||||
ToolArguments string `gorm:"type:text;"`
|
||||
ToolResult string `gorm:"type:text;"`
|
||||
ErrorText string `gorm:"type:text;"`
|
||||
IncludeInContext *bool
|
||||
}
|
||||
|
||||
// BeforeCreate - GORM hook.
|
||||
@@ -136,6 +192,17 @@ func (a *AIConversationMessage) ToProtobuf() *clientpb.AIConversationMessage {
|
||||
Content: a.Content,
|
||||
ProviderMessageID: a.ProviderMessageID,
|
||||
FinishReason: a.FinishReason,
|
||||
Kind: clientpb.AIConversationMessageKind(a.Kind),
|
||||
Visibility: clientpb.AIConversationMessageVisibility(a.Visibility),
|
||||
State: clientpb.AIConversationMessageState(a.State),
|
||||
TurnID: a.TurnID,
|
||||
ItemID: a.ItemID,
|
||||
ToolCallID: a.ToolCallID,
|
||||
ToolName: a.ToolName,
|
||||
ToolArguments: a.ToolArguments,
|
||||
ToolResult: a.ToolResult,
|
||||
ErrorText: a.ErrorText,
|
||||
IncludeInContext: resolveAIConversationMessageIncludeInContext(a.IncludeInContext, a.Visibility),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -159,5 +226,32 @@ func AIConversationMessageFromProtobuf(pbMessage *clientpb.AIConversationMessage
|
||||
Content: pbMessage.Content,
|
||||
ProviderMessageID: pbMessage.ProviderMessageID,
|
||||
FinishReason: pbMessage.FinishReason,
|
||||
Kind: int32(pbMessage.Kind),
|
||||
Visibility: int32(pbMessage.Visibility),
|
||||
State: int32(pbMessage.State),
|
||||
TurnID: pbMessage.TurnID,
|
||||
ItemID: pbMessage.ItemID,
|
||||
ToolCallID: pbMessage.ToolCallID,
|
||||
ToolName: pbMessage.ToolName,
|
||||
ToolArguments: pbMessage.ToolArguments,
|
||||
ToolResult: pbMessage.ToolResult,
|
||||
ErrorText: pbMessage.ErrorText,
|
||||
IncludeInContext: cloneOptionalBool(pbMessage.IncludeInContext),
|
||||
}
|
||||
}
|
||||
|
||||
func resolveAIConversationMessageIncludeInContext(value *bool, visibility int32) *bool {
|
||||
if value != nil {
|
||||
return cloneOptionalBool(value)
|
||||
}
|
||||
include := clientpb.AIConversationMessageVisibility(visibility) == clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT
|
||||
return &include
|
||||
}
|
||||
|
||||
func cloneOptionalBool(value *bool) *bool {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *value
|
||||
return &cloned
|
||||
}
|
||||
+42
-19
@@ -143,36 +143,59 @@ func pivotPeerFailureHandler(implantConn *core.ImplantConnection, data []byte) *
|
||||
}
|
||||
pivotLog.Errorf("pivot peer failure received: %v", peerFailure)
|
||||
|
||||
reportingSession := core.Sessions.FromImplantConnection(implantConn)
|
||||
if reportingSession == nil {
|
||||
pivotLog.Warn("ignoring pivot peer failure from unknown session")
|
||||
return nil
|
||||
}
|
||||
|
||||
targetSessionIDs := map[string]struct{}{}
|
||||
targetPivotIDs := map[string]struct{}{}
|
||||
core.PivotSessions.Range(func(key, value interface{}) bool {
|
||||
pivot := value.(*core.Pivot)
|
||||
|
||||
found := pivot.OriginID == peerFailure.PeerID
|
||||
|
||||
if !found {
|
||||
pivotLog.Warnf("Filed peer not found by OriginID, searching by Peers instead")
|
||||
|
||||
for _, peer := range pivot.Peers {
|
||||
if peer.PeerID == peerFailure.PeerID {
|
||||
pivotLog.Warnf("Found session with needed peer!")
|
||||
found = true
|
||||
}
|
||||
}
|
||||
reporterIdx := pivotPeerIndex(pivot.Peers, reportingSession.PeerID)
|
||||
if reporterIdx <= 0 {
|
||||
return true
|
||||
}
|
||||
if pivot.Peers[reporterIdx-1].PeerID != peerFailure.PeerID {
|
||||
return true
|
||||
}
|
||||
|
||||
if found {
|
||||
session := core.Sessions.FromImplantConnection(pivot.ImplantConn)
|
||||
targetPivotIDs[pivot.ID] = struct{}{}
|
||||
if pivot.OriginID != peerFailure.PeerID {
|
||||
return true
|
||||
}
|
||||
|
||||
if session != nil {
|
||||
core.Sessions.Remove(session.ID)
|
||||
}
|
||||
defer core.PivotSessions.Delete(pivot.ID)
|
||||
return false
|
||||
session := core.Sessions.FromImplantConnection(pivot.ImplantConn)
|
||||
if session != nil {
|
||||
targetSessionIDs[session.ID] = struct{}{}
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if len(targetSessionIDs) == 0 {
|
||||
pivotLog.Warnf("ignoring unauthorized pivot peer failure from %d for %d", reportingSession.PeerID, peerFailure.PeerID)
|
||||
return nil
|
||||
}
|
||||
|
||||
for sessionID := range targetSessionIDs {
|
||||
core.Sessions.Remove(sessionID)
|
||||
}
|
||||
for pivotID := range targetPivotIDs {
|
||||
core.PivotSessions.Delete(pivotID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func pivotPeerIndex(peers []*sliverpb.PivotPeer, peerID int64) int {
|
||||
for idx, peer := range peers {
|
||||
if peer.PeerID == peerID {
|
||||
return idx
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func pivotServerPingHandler(pivot *core.Pivot) {
|
||||
pivot.ImplantConn.UpdateLastMessage()
|
||||
pivot.ImmediateImplantConn.UpdateLastMessage()
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
func TestPivotPeerFailureHandlerRemovesAuthorizedImmediateChild(t *testing.T) {
|
||||
resetPivotFailureTestState(t)
|
||||
|
||||
reporterConn, reporterSession := addPivotFailureTestSession(t, 300, "wg")
|
||||
_, unrelatedSession := addPivotFailureTestSession(t, 900, "wg")
|
||||
childSession, childPivot := addPivotFailureTestPivotSession(t, 200, 200, 300)
|
||||
grandchildSession, grandchildPivot := addPivotFailureTestPivotSession(t, 100, 100, 200, 300)
|
||||
unrelatedChildSession, unrelatedChildPivot := addPivotFailureTestPivotSession(t, 800, 800, 900)
|
||||
|
||||
pivotPeerFailureHandler(reporterConn, mustMarshalPivotPeerFailure(t, 200))
|
||||
|
||||
if got := core.Sessions.Get(reporterSession.ID); got == nil {
|
||||
t.Fatalf("reporting session was removed")
|
||||
}
|
||||
if got := core.Sessions.Get(childSession.ID); got != nil {
|
||||
t.Fatalf("child session still present after authorized failure")
|
||||
}
|
||||
if got := core.Sessions.Get(grandchildSession.ID); got != nil {
|
||||
t.Fatalf("grandchild session still present after authorized failure")
|
||||
}
|
||||
if got := core.Sessions.Get(unrelatedSession.ID); got == nil {
|
||||
t.Fatalf("unrelated direct session was removed")
|
||||
}
|
||||
if got := core.Sessions.Get(unrelatedChildSession.ID); got == nil {
|
||||
t.Fatalf("unrelated pivot session was removed")
|
||||
}
|
||||
if pivotSessionExists(childPivot.ID) {
|
||||
t.Fatalf("child pivot session still present after authorized failure")
|
||||
}
|
||||
if pivotSessionExists(grandchildPivot.ID) {
|
||||
t.Fatalf("grandchild pivot session still present after authorized failure")
|
||||
}
|
||||
if !pivotSessionExists(unrelatedChildPivot.ID) {
|
||||
t.Fatalf("unrelated pivot session was removed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPivotPeerFailureHandlerRejectsNonImmediateChild(t *testing.T) {
|
||||
resetPivotFailureTestState(t)
|
||||
|
||||
reporterConn, _ := addPivotFailureTestSession(t, 300, "wg")
|
||||
childSession, childPivot := addPivotFailureTestPivotSession(t, 200, 200, 300)
|
||||
grandchildSession, grandchildPivot := addPivotFailureTestPivotSession(t, 100, 100, 200, 300)
|
||||
|
||||
pivotPeerFailureHandler(reporterConn, mustMarshalPivotPeerFailure(t, 100))
|
||||
|
||||
if got := core.Sessions.Get(childSession.ID); got == nil {
|
||||
t.Fatalf("child session was removed by non-immediate failure report")
|
||||
}
|
||||
if got := core.Sessions.Get(grandchildSession.ID); got == nil {
|
||||
t.Fatalf("grandchild session was removed by non-immediate failure report")
|
||||
}
|
||||
if !pivotSessionExists(childPivot.ID) {
|
||||
t.Fatalf("child pivot session was removed by non-immediate failure report")
|
||||
}
|
||||
if !pivotSessionExists(grandchildPivot.ID) {
|
||||
t.Fatalf("grandchild pivot session was removed by non-immediate failure report")
|
||||
}
|
||||
}
|
||||
|
||||
func addPivotFailureTestSession(t *testing.T, peerID int64, transport string) (*core.ImplantConnection, *core.Session) {
|
||||
t.Helper()
|
||||
|
||||
conn := core.NewImplantConnection(transport, fmt.Sprintf("peer-%d", peerID))
|
||||
session := core.NewSession(conn)
|
||||
session.Name = fmt.Sprintf("peer-%d", peerID)
|
||||
session.PeerID = peerID
|
||||
core.Sessions.Add(session)
|
||||
return conn, session
|
||||
}
|
||||
|
||||
func addPivotFailureTestPivotSession(t *testing.T, originID int64, chain ...int64) (*core.Session, *core.Pivot) {
|
||||
t.Helper()
|
||||
|
||||
peers := make([]*sliverpb.PivotPeer, 0, len(chain))
|
||||
for _, peerID := range chain {
|
||||
peers = append(peers, &sliverpb.PivotPeer{
|
||||
PeerID: peerID,
|
||||
Name: fmt.Sprintf("peer-%d", peerID),
|
||||
})
|
||||
}
|
||||
|
||||
pivot := core.NewPivotSession(peers)
|
||||
pivot.OriginID = originID
|
||||
pivot.ImplantConn = core.NewImplantConnection(core.PivotTransportName, fmt.Sprintf("pivot-%d", originID))
|
||||
core.PivotSessions.Store(pivot.ID, pivot)
|
||||
|
||||
session := core.NewSession(pivot.ImplantConn)
|
||||
session.Name = fmt.Sprintf("peer-%d", originID)
|
||||
session.PeerID = originID
|
||||
core.Sessions.Add(session)
|
||||
|
||||
return session, pivot
|
||||
}
|
||||
|
||||
func mustMarshalPivotPeerFailure(t *testing.T, peerID int64) []byte {
|
||||
t.Helper()
|
||||
|
||||
data, err := proto.Marshal(&sliverpb.PivotPeerFailure{
|
||||
PeerID: peerID,
|
||||
Type: sliverpb.PeerFailureType_SEND_FAILURE,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal peer failure: %v", err)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func pivotSessionExists(pivotID string) bool {
|
||||
_, ok := core.PivotSessions.Load(pivotID)
|
||||
return ok
|
||||
}
|
||||
|
||||
func resetPivotFailureTestState(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
for _, session := range core.Sessions.All() {
|
||||
core.Sessions.Remove(session.ID)
|
||||
}
|
||||
core.PivotSessions.Range(func(key, value interface{}) bool {
|
||||
core.PivotSessions.Delete(key)
|
||||
return true
|
||||
})
|
||||
|
||||
t.Cleanup(func() {
|
||||
for _, session := range core.Sessions.All() {
|
||||
core.Sessions.Remove(session.ID)
|
||||
}
|
||||
core.PivotSessions.Range(func(key, value interface{}) bool {
|
||||
core.PivotSessions.Delete(key)
|
||||
return true
|
||||
})
|
||||
})
|
||||
}
|
||||
@@ -169,10 +169,10 @@ func tunnelCloseHandler(implantConn *core.ImplantConnection, data []byte) *slive
|
||||
}
|
||||
} else {
|
||||
rtunnel := rtunnels.GetRTunnel(tunnelData.TunnelID)
|
||||
if rtunnel != nil && session.ID == tunnel.SessionID {
|
||||
if rtunnel != nil && session.ID == rtunnel.SessionID {
|
||||
rtunnel.Close()
|
||||
rtunnels.RemoveRTunnel(rtunnel.ID)
|
||||
} else if rtunnel != nil && session.ID != tunnel.SessionID {
|
||||
} else if rtunnel != nil && session.ID != rtunnel.SessionID {
|
||||
sessionHandlerLog.Warnf("Warning: Session %s attempted to send data on reverse tunnel it did not own", session.ID)
|
||||
} else {
|
||||
sessionHandlerLog.Warnf("Close sent on nil tunnel %d", tunnelData.TunnelID)
|
||||
@@ -206,7 +206,7 @@ func socksDataHandler(implantConn *core.ImplantConnection, data []byte) *sliverp
|
||||
// core.SocksTunnels.Close(socksData.TunnelID)
|
||||
// return nil
|
||||
//}
|
||||
sessionHandlerLog.Debugf("socksDataHandler:", len(socksData.Data), socksData.Data)
|
||||
sessionHandlerLog.Debugf("socksDataHandler: %d bytes: %v", len(socksData.Data), socksData.Data)
|
||||
socksTunnel := core.SocksTunnels.Get(socksData.TunnelID)
|
||||
if socksTunnel != nil {
|
||||
if session.ID == socksTunnel.SessionID {
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"github.com/bishopfox/sliver/server/core/rtunnels"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
type testWriteCloser struct {
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (t *testWriteCloser) Write(data []byte) (int, error) {
|
||||
return len(data), nil
|
||||
}
|
||||
|
||||
func (t *testWriteCloser) Close() error {
|
||||
t.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func addTestSession(t *testing.T) (*core.ImplantConnection, *core.Session) {
|
||||
t.Helper()
|
||||
|
||||
conn := core.NewImplantConnection("test", "n/a")
|
||||
session := core.NewSession(conn)
|
||||
core.Sessions.Add(session)
|
||||
t.Cleanup(func() {
|
||||
core.Sessions.Remove(session.ID)
|
||||
})
|
||||
|
||||
return conn, session
|
||||
}
|
||||
|
||||
func marshalTunnelCloseData(t *testing.T, tunnelID uint64) []byte {
|
||||
t.Helper()
|
||||
|
||||
data, err := proto.Marshal(&sliverpb.TunnelData{
|
||||
TunnelID: tunnelID,
|
||||
Closed: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal tunnel close data: %v", err)
|
||||
}
|
||||
|
||||
return data
|
||||
}
|
||||
|
||||
func assertNoPanic(t *testing.T, fn func()) {
|
||||
t.Helper()
|
||||
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
t.Fatalf("unexpected panic: %v", recovered)
|
||||
}
|
||||
}()
|
||||
|
||||
fn()
|
||||
}
|
||||
|
||||
func TestTunnelCloseHandlerClosesOwnedReverseTunnel(t *testing.T) {
|
||||
conn, session := addTestSession(t)
|
||||
tunnelID := core.NewTunnelID()
|
||||
writer := &testWriteCloser{}
|
||||
|
||||
rtunnels.AddRTunnel(rtunnels.NewRTunnel(tunnelID, session.ID, writer))
|
||||
t.Cleanup(func() {
|
||||
rtunnels.RemoveRTunnel(tunnelID)
|
||||
})
|
||||
|
||||
assertNoPanic(t, func() {
|
||||
tunnelCloseHandler(conn, marshalTunnelCloseData(t, tunnelID))
|
||||
})
|
||||
|
||||
if rtunnel := rtunnels.GetRTunnel(tunnelID); rtunnel != nil {
|
||||
t.Fatalf("expected reverse tunnel %d to be removed", tunnelID)
|
||||
}
|
||||
if !writer.closed {
|
||||
t.Fatal("expected reverse tunnel writer to be closed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTunnelCloseHandlerKeepsUnownedReverseTunnel(t *testing.T) {
|
||||
ownerConn, ownerSession := addTestSession(t)
|
||||
_ = ownerConn
|
||||
attackerConn, attackerSession := addTestSession(t)
|
||||
tunnelID := core.NewTunnelID()
|
||||
writer := &testWriteCloser{}
|
||||
|
||||
rtunnels.AddRTunnel(rtunnels.NewRTunnel(tunnelID, ownerSession.ID, writer))
|
||||
t.Cleanup(func() {
|
||||
rtunnels.RemoveRTunnel(tunnelID)
|
||||
})
|
||||
|
||||
assertNoPanic(t, func() {
|
||||
tunnelCloseHandler(attackerConn, marshalTunnelCloseData(t, tunnelID))
|
||||
})
|
||||
|
||||
rtunnel := rtunnels.GetRTunnel(tunnelID)
|
||||
if rtunnel == nil {
|
||||
t.Fatalf("expected reverse tunnel %d to remain for owner %s", tunnelID, ownerSession.ID)
|
||||
}
|
||||
if rtunnel.SessionID != ownerSession.ID {
|
||||
t.Fatalf("expected reverse tunnel owner %s, got %s", ownerSession.ID, rtunnel.SessionID)
|
||||
}
|
||||
if attackerSession.ID == ownerSession.ID {
|
||||
t.Fatal("expected distinct owner and attacker sessions")
|
||||
}
|
||||
if writer.closed {
|
||||
t.Fatal("expected unowned reverse tunnel writer to remain open")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,361 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
stdbase32 "encoding/base32"
|
||||
stdbase64 "encoding/base64"
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
utilencoders "github.com/bishopfox/sliver/util/encoders"
|
||||
)
|
||||
|
||||
const aiEncodingDefaultMaxBytes = 64 * 1024
|
||||
|
||||
type encodingInputArgs struct {
|
||||
Text *string `json:"text,omitempty"`
|
||||
DataBase64 *string `json:"data_base64,omitempty"`
|
||||
}
|
||||
|
||||
type encodedValueArgs struct {
|
||||
Value *string `json:"value,omitempty"`
|
||||
MaxBytes *int64 `json:"max_bytes,omitempty"`
|
||||
}
|
||||
|
||||
type gzipDecodeArgs struct {
|
||||
DataBase64 *string `json:"data_base64,omitempty"`
|
||||
MaxBytes *int64 `json:"max_bytes,omitempty"`
|
||||
}
|
||||
|
||||
type encodingValueResult struct {
|
||||
Encoding string `json:"encoding"`
|
||||
Operation string `json:"operation"`
|
||||
Variant string `json:"variant,omitempty"`
|
||||
Value string `json:"value"`
|
||||
ByteLen int `json:"byte_len"`
|
||||
}
|
||||
|
||||
type encodingBytesResult struct {
|
||||
Encoding string `json:"encoding"`
|
||||
Operation string `json:"operation"`
|
||||
Variant string `json:"variant,omitempty"`
|
||||
ByteLen int `json:"byte_len"`
|
||||
SHA256 string `json:"sha256,omitempty"`
|
||||
DataBase64 string `json:"data_base64,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
}
|
||||
|
||||
type standardBase32Encoder struct{}
|
||||
|
||||
type standardBase64Encoder struct{}
|
||||
|
||||
func encodingToolDefinitions() []serverai.AgenticToolDefinition {
|
||||
return []serverai.AgenticToolDefinition{
|
||||
{
|
||||
Name: "base32_encode",
|
||||
Description: "Encode bytes using standard RFC 4648 base32 with padding. Provide either text or data_base64 input.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Optional standard base64 wrapper for arbitrary input bytes."},
|
||||
"text": map[string]any{"type": "string", "description": "Optional UTF-8 text to encode."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "base32_decode",
|
||||
Description: "Decode standard RFC 4648 base32. Returns the decoded bytes as standard data_base64 plus UTF-8 text when valid.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional maximum decoded byte length. Defaults to 65536."},
|
||||
"value": map[string]any{"type": "string", "description": "Base32 value to decode."},
|
||||
},
|
||||
"required": []string{"value"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "base64_encode",
|
||||
Description: "Encode bytes using standard RFC 4648 base64 with padding. Provide either text or data_base64 input.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Optional standard base64 wrapper for arbitrary input bytes."},
|
||||
"text": map[string]any{"type": "string", "description": "Optional UTF-8 text to encode."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "base64_decode",
|
||||
Description: "Decode standard RFC 4648 base64. Returns the decoded bytes as standard data_base64 plus UTF-8 text when valid.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional maximum decoded byte length. Defaults to 65536."},
|
||||
"value": map[string]any{"type": "string", "description": "Base64 value to decode."},
|
||||
},
|
||||
"required": []string{"value"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "hex_encode",
|
||||
Description: "Encode bytes as lowercase hexadecimal. Provide either text or data_base64 input.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Optional standard base64 wrapper for arbitrary input bytes."},
|
||||
"text": map[string]any{"type": "string", "description": "Optional UTF-8 text to encode."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "hex_decode",
|
||||
Description: "Decode a hexadecimal string. Returns the decoded bytes as standard data_base64 plus UTF-8 text when valid.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional maximum decoded byte length. Defaults to 65536."},
|
||||
"value": map[string]any{"type": "string", "description": "Hexadecimal value to decode."},
|
||||
},
|
||||
"required": []string{"value"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "gzip_encode",
|
||||
Description: "Compress bytes with gzip. Provide either text or data_base64 input. Returns the compressed bytes as standard data_base64.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Optional standard base64 wrapper for arbitrary input bytes."},
|
||||
"text": map[string]any{"type": "string", "description": "Optional UTF-8 text to compress."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "gzip_decode",
|
||||
Description: "Decompress gzip bytes supplied as standard data_base64. Returns the decoded bytes as standard data_base64 plus UTF-8 text when valid.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Standard base64 wrapper for gzip-compressed bytes."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional maximum decoded byte length. Defaults to 65536."},
|
||||
},
|
||||
"required": []string{"data_base64"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callEncodingTool(_ context.Context, name string, arguments string) (string, bool, error) {
|
||||
switch strings.TrimSpace(name) {
|
||||
case "base32_encode":
|
||||
var args encodingInputArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callEncodeValueTool("base32", "standard", args, standardBase32Encoder{})
|
||||
return result, true, err
|
||||
case "base32_decode":
|
||||
var args encodedValueArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callDecodeValueTool("base32", "standard", args, standardBase32Encoder{})
|
||||
return result, true, err
|
||||
case "base64_encode":
|
||||
var args encodingInputArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callEncodeValueTool("base64", "standard", args, standardBase64Encoder{})
|
||||
return result, true, err
|
||||
case "base64_decode":
|
||||
var args encodedValueArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callDecodeValueTool("base64", "standard", args, standardBase64Encoder{})
|
||||
return result, true, err
|
||||
case "hex_encode":
|
||||
var args encodingInputArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callEncodeValueTool("hex", "", args, utilencoders.Hex{})
|
||||
return result, true, err
|
||||
case "hex_decode":
|
||||
var args encodedValueArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := callDecodeValueTool("hex", "", args, utilencoders.Hex{})
|
||||
return result, true, err
|
||||
case "gzip_encode":
|
||||
var args encodingInputArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
input, err := resolveEncodingInput(args)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
output, err := new(utilencoders.Gzip).Encode(input)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := marshalToolResult(newEncodingBytesResult("gzip", "encode", "", output))
|
||||
return result, true, err
|
||||
case "gzip_decode":
|
||||
var args gzipDecodeArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
data, err := requireStandardBase64(args.DataBase64, "data_base64")
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
output, err := new(utilencoders.Gzip).DecodeWithMaxLen(data, resolveMaxBytes(args.MaxBytes))
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := marshalToolResult(newEncodingBytesResult("gzip", "decode", "", output))
|
||||
return result, true, err
|
||||
default:
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func callEncodeValueTool(name string, variant string, args encodingInputArgs, encoder utilencoders.Encoder) (string, error) {
|
||||
input, err := resolveEncodingInput(args)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
output, err := encoder.Encode(input)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(encodingValueResult{
|
||||
Encoding: name,
|
||||
Operation: "encode",
|
||||
Variant: variant,
|
||||
Value: string(output),
|
||||
ByteLen: len(output),
|
||||
})
|
||||
}
|
||||
|
||||
func (standardBase32Encoder) Encode(data []byte) ([]byte, error) {
|
||||
return []byte(stdbase32.StdEncoding.EncodeToString(data)), nil
|
||||
}
|
||||
|
||||
func (standardBase32Encoder) Decode(data []byte) ([]byte, error) {
|
||||
value := strings.TrimSpace(string(data))
|
||||
decoded, err := stdbase32.StdEncoding.DecodeString(value)
|
||||
if err == nil {
|
||||
return decoded, nil
|
||||
}
|
||||
return stdbase32.StdEncoding.WithPadding(stdbase32.NoPadding).DecodeString(value)
|
||||
}
|
||||
|
||||
func (standardBase64Encoder) Encode(data []byte) ([]byte, error) {
|
||||
return []byte(stdbase64.StdEncoding.EncodeToString(data)), nil
|
||||
}
|
||||
|
||||
func (standardBase64Encoder) Decode(data []byte) ([]byte, error) {
|
||||
return decodeStandardBase64(strings.TrimSpace(string(data)))
|
||||
}
|
||||
|
||||
func callDecodeValueTool(name string, variant string, args encodedValueArgs, encoder utilencoders.Encoder) (string, error) {
|
||||
value, err := requireString(args.Value, "value")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
output, err := encoder.Decode([]byte(value))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
maxBytes := resolveMaxBytes(args.MaxBytes)
|
||||
if int64(len(output)) > maxBytes {
|
||||
return "", fmt.Errorf("%s decoded payload exceeds %d bytes", name, maxBytes)
|
||||
}
|
||||
|
||||
return marshalToolResult(newEncodingBytesResult(name, "decode", variant, output))
|
||||
}
|
||||
|
||||
func resolveEncodingInput(args encodingInputArgs) ([]byte, error) {
|
||||
hasText := args.Text != nil
|
||||
hasData := args.DataBase64 != nil
|
||||
switch {
|
||||
case hasText && hasData:
|
||||
return nil, fmt.Errorf("provide only one of text or data_base64")
|
||||
case hasText:
|
||||
return []byte(*args.Text), nil
|
||||
case hasData:
|
||||
return decodeStandardBase64(*args.DataBase64)
|
||||
default:
|
||||
return nil, fmt.Errorf("text or data_base64 is required")
|
||||
}
|
||||
}
|
||||
|
||||
func requireString(value *string, field string) (string, error) {
|
||||
if value == nil {
|
||||
return "", fmt.Errorf("%s is required", field)
|
||||
}
|
||||
return *value, nil
|
||||
}
|
||||
|
||||
func requireStandardBase64(value *string, field string) ([]byte, error) {
|
||||
if value == nil {
|
||||
return nil, fmt.Errorf("%s is required", field)
|
||||
}
|
||||
data, err := decodeStandardBase64(*value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid %s: %w", field, err)
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func decodeStandardBase64(value string) ([]byte, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
data, err := stdbase64.StdEncoding.DecodeString(value)
|
||||
if err == nil {
|
||||
return data, nil
|
||||
}
|
||||
return stdbase64.RawStdEncoding.DecodeString(value)
|
||||
}
|
||||
|
||||
func resolveMaxBytes(maxBytes *int64) int64 {
|
||||
if maxBytes != nil && *maxBytes > 0 {
|
||||
return *maxBytes
|
||||
}
|
||||
return aiEncodingDefaultMaxBytes
|
||||
}
|
||||
|
||||
func newEncodingBytesResult(name string, operation string, variant string, data []byte) encodingBytesResult {
|
||||
result := encodingBytesResult{
|
||||
Encoding: name,
|
||||
Operation: operation,
|
||||
Variant: variant,
|
||||
ByteLen: len(data),
|
||||
DataBase64: stdbase64.StdEncoding.EncodeToString(data),
|
||||
}
|
||||
if len(data) > 0 {
|
||||
result.SHA256 = sha256Hex(data)
|
||||
if utf8.Valid(data) {
|
||||
result.Text = string(data)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
stdbase32 "encoding/base32"
|
||||
stdbase64 "encoding/base64"
|
||||
stdhex "encoding/hex"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValueEncodingToolsRoundTrip(t *testing.T) {
|
||||
exec := &executor{}
|
||||
ctx := context.Background()
|
||||
|
||||
cases := []struct {
|
||||
encodeTool string
|
||||
decodeTool string
|
||||
expected string
|
||||
variant string
|
||||
}{
|
||||
{encodeTool: "base32_encode", decodeTool: "base32_decode", expected: stdbase32.StdEncoding.EncodeToString([]byte("hello world")), variant: "standard"},
|
||||
{encodeTool: "base64_encode", decodeTool: "base64_decode", expected: stdbase64.StdEncoding.EncodeToString([]byte("hello world")), variant: "standard"},
|
||||
{encodeTool: "hex_encode", decodeTool: "hex_decode", expected: stdhex.EncodeToString([]byte("hello world")), variant: ""},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.encodeTool, func(t *testing.T) {
|
||||
encodedRaw, err := exec.CallTool(ctx, tc.encodeTool, `{"text":"hello world"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("encode with %s: %v", tc.encodeTool, err)
|
||||
}
|
||||
|
||||
var encoded encodingValueResult
|
||||
if err := json.Unmarshal([]byte(encodedRaw), &encoded); err != nil {
|
||||
t.Fatalf("decode encode result: %v", err)
|
||||
}
|
||||
if encoded.Value == "" {
|
||||
t.Fatalf("expected encoded value from %s", tc.encodeTool)
|
||||
}
|
||||
if encoded.Value != tc.expected {
|
||||
t.Fatalf("unexpected encoded value: got=%q want=%q", encoded.Value, tc.expected)
|
||||
}
|
||||
if encoded.Variant != tc.variant {
|
||||
t.Fatalf("unexpected variant: got=%q want=%q", encoded.Variant, tc.variant)
|
||||
}
|
||||
|
||||
decodedRaw, err := exec.CallTool(ctx, tc.decodeTool, `{"value":"`+encoded.Value+`"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("decode with %s: %v", tc.decodeTool, err)
|
||||
}
|
||||
|
||||
var decoded encodingBytesResult
|
||||
if err := json.Unmarshal([]byte(decodedRaw), &decoded); err != nil {
|
||||
t.Fatalf("decode decode result: %v", err)
|
||||
}
|
||||
if decoded.Text != "hello world" {
|
||||
t.Fatalf("unexpected decoded text: %+v", decoded)
|
||||
}
|
||||
if decoded.Variant != tc.variant {
|
||||
t.Fatalf("unexpected decoded variant: got=%q want=%q", decoded.Variant, tc.variant)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGzipEncodingToolsRoundTrip(t *testing.T) {
|
||||
exec := &executor{}
|
||||
ctx := context.Background()
|
||||
|
||||
encodedRaw, err := exec.CallTool(ctx, "gzip_encode", `{"text":"hello world"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("gzip encode: %v", err)
|
||||
}
|
||||
|
||||
var encoded encodingBytesResult
|
||||
if err := json.Unmarshal([]byte(encodedRaw), &encoded); err != nil {
|
||||
t.Fatalf("decode gzip encode result: %v", err)
|
||||
}
|
||||
if encoded.DataBase64 == "" {
|
||||
t.Fatalf("expected gzip output bytes, got %+v", encoded)
|
||||
}
|
||||
|
||||
decodedRaw, err := exec.CallTool(ctx, "gzip_decode", `{"data_base64":"`+encoded.DataBase64+`"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("gzip decode: %v", err)
|
||||
}
|
||||
|
||||
var decoded encodingBytesResult
|
||||
if err := json.Unmarshal([]byte(decodedRaw), &decoded); err != nil {
|
||||
t.Fatalf("decode gzip decode result: %v", err)
|
||||
}
|
||||
if decoded.Text != "hello world" {
|
||||
t.Fatalf("unexpected gzip decode text: %+v", decoded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValueDecodingToolsRespectMaxBytes(t *testing.T) {
|
||||
exec := &executor{}
|
||||
ctx := context.Background()
|
||||
|
||||
encoded := stdbase64.StdEncoding.EncodeToString([]byte("abc"))
|
||||
if _, err := exec.CallTool(ctx, "gzip_decode", `{"data_base64":"`+encoded+`","max_bytes":2}`); err == nil {
|
||||
t.Fatal("expected gzip decode to reject invalid compressed data")
|
||||
}
|
||||
|
||||
hexEncodedRaw, err := exec.CallTool(ctx, "hex_encode", `{"text":"abc"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("hex encode: %v", err)
|
||||
}
|
||||
var hexEncoded encodingValueResult
|
||||
if err := json.Unmarshal([]byte(hexEncodedRaw), &hexEncoded); err != nil {
|
||||
t.Fatalf("decode hex encode result: %v", err)
|
||||
}
|
||||
|
||||
if _, err := exec.CallTool(ctx, "hex_decode", `{"value":"`+hexEncoded.Value+`","max_bytes":2}`); err == nil {
|
||||
t.Fatal("expected hex decode max_bytes limit to be enforced")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,748 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"github.com/bishopfox/sliver/util/encoders"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const aiToolDefaultWaitSeconds = int64(60)
|
||||
|
||||
// Backend captures the RPC methods the AI tool executor needs.
|
||||
type Backend interface {
|
||||
GetSessions(context.Context, *commonpb.Empty) (*clientpb.Sessions, error)
|
||||
GetBeacons(context.Context, *commonpb.Empty) (*clientpb.Beacons, error)
|
||||
Ls(context.Context, *sliverpb.LsReq) (*sliverpb.Ls, error)
|
||||
Mv(context.Context, *sliverpb.MvReq) (*sliverpb.Mv, error)
|
||||
Cp(context.Context, *sliverpb.CpReq) (*sliverpb.Cp, error)
|
||||
Rm(context.Context, *sliverpb.RmReq) (*sliverpb.Rm, error)
|
||||
Mkdir(context.Context, *sliverpb.MkdirReq) (*sliverpb.Mkdir, error)
|
||||
Cd(context.Context, *sliverpb.CdReq) (*sliverpb.Pwd, error)
|
||||
Download(context.Context, *sliverpb.DownloadReq) (*sliverpb.Download, error)
|
||||
Pwd(context.Context, *sliverpb.PwdReq) (*sliverpb.Pwd, error)
|
||||
Chmod(context.Context, *sliverpb.ChmodReq) (*sliverpb.Chmod, error)
|
||||
Chown(context.Context, *sliverpb.ChownReq) (*sliverpb.Chown, error)
|
||||
Chtimes(context.Context, *sliverpb.ChtimesReq) (*sliverpb.Chtimes, error)
|
||||
Mount(context.Context, *sliverpb.MountReq) (*sliverpb.Mount, error)
|
||||
Ifconfig(context.Context, *sliverpb.IfconfigReq) (*sliverpb.Ifconfig, error)
|
||||
Netstat(context.Context, *sliverpb.NetstatReq) (*sliverpb.Netstat, error)
|
||||
Ps(context.Context, *sliverpb.PsReq) (*sliverpb.Ps, error)
|
||||
GetEnv(context.Context, *sliverpb.EnvReq) (*sliverpb.EnvInfo, error)
|
||||
Ping(context.Context, *sliverpb.Ping) (*sliverpb.Ping, error)
|
||||
Screenshot(context.Context, *sliverpb.ScreenshotReq) (*sliverpb.Screenshot, error)
|
||||
Execute(context.Context, *sliverpb.ExecuteReq) (*sliverpb.Execute, error)
|
||||
ExecuteWindows(context.Context, *sliverpb.ExecuteWindowsReq) (*sliverpb.Execute, error)
|
||||
ExecuteAssembly(context.Context, *sliverpb.ExecuteAssemblyReq) (*sliverpb.ExecuteAssembly, error)
|
||||
Sideload(context.Context, *sliverpb.SideloadReq) (*sliverpb.Sideload, error)
|
||||
SpawnDll(context.Context, *sliverpb.InvokeSpawnDllReq) (*sliverpb.SpawnDll, error)
|
||||
RegisterExtension(context.Context, *sliverpb.RegisterExtensionReq) (*sliverpb.RegisterExtension, error)
|
||||
ListExtensions(context.Context, *sliverpb.ListExtensionsReq) (*sliverpb.ListExtensions, error)
|
||||
CallExtension(context.Context, *sliverpb.CallExtensionReq) (*sliverpb.CallExtension, error)
|
||||
}
|
||||
|
||||
type executor struct {
|
||||
backend Backend
|
||||
conversation *clientpb.AIConversation
|
||||
}
|
||||
|
||||
type targetArgs struct {
|
||||
SessionID string `json:"session_id,omitempty"`
|
||||
BeaconID string `json:"beacon_id,omitempty"`
|
||||
}
|
||||
|
||||
type lsToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
}
|
||||
|
||||
type catToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
MaxBytes int64 `json:"max_bytes,omitempty"`
|
||||
MaxLines int64 `json:"max_lines,omitempty"`
|
||||
}
|
||||
|
||||
type lsFileEntry struct {
|
||||
Name string `json:"name"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
Size int64 `json:"size"`
|
||||
ModTime int64 `json:"mod_time"`
|
||||
Mode string `json:"mode"`
|
||||
Link string `json:"link,omitempty"`
|
||||
UID string `json:"uid,omitempty"`
|
||||
GID string `json:"gid,omitempty"`
|
||||
}
|
||||
|
||||
type lsResult struct {
|
||||
Path string `json:"path"`
|
||||
Exists bool `json:"exists"`
|
||||
Timezone string `json:"timezone"`
|
||||
TimezoneOffset int32 `json:"timezone_offset"`
|
||||
Files []lsFileEntry `json:"files"`
|
||||
}
|
||||
|
||||
type catResult struct {
|
||||
Path string `json:"path"`
|
||||
RequestedPath string `json:"requested_path"`
|
||||
Exists bool `json:"exists"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
Encoder string `json:"encoder,omitempty"`
|
||||
Decoded bool `json:"decoded"`
|
||||
ByteLen int `json:"byte_len"`
|
||||
DataBase64 string `json:"data_base64"`
|
||||
Text string `json:"text,omitempty"`
|
||||
}
|
||||
|
||||
type pathResult struct {
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
type sessionSummary struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
RemoteAddress string `json:"remote_address"`
|
||||
Hostname string `json:"hostname"`
|
||||
Username string `json:"username"`
|
||||
PID int32 `json:"pid"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Locale string `json:"locale"`
|
||||
LastCheckin int64 `json:"last_checkin"`
|
||||
IsDead bool `json:"is_dead"`
|
||||
Integrity string `json:"integrity"`
|
||||
}
|
||||
|
||||
type beaconSummary struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
RemoteAddress string `json:"remote_address"`
|
||||
Hostname string `json:"hostname"`
|
||||
Username string `json:"username"`
|
||||
PID int32 `json:"pid"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Locale string `json:"locale"`
|
||||
LastCheckin int64 `json:"last_checkin"`
|
||||
NextCheckin int64 `json:"next_checkin"`
|
||||
Interval int64 `json:"interval"`
|
||||
Jitter int64 `json:"jitter"`
|
||||
TasksCount int64 `json:"tasks_count"`
|
||||
TasksCountCompleted int64 `json:"tasks_count_completed"`
|
||||
IsDead bool `json:"is_dead"`
|
||||
Integrity string `json:"integrity"`
|
||||
}
|
||||
|
||||
type sessionsAndBeaconsResult struct {
|
||||
Sessions []sessionSummary `json:"sessions"`
|
||||
SessionsCount int `json:"sessions_count"`
|
||||
Beacons []beaconSummary `json:"beacons"`
|
||||
BeaconsCount int `json:"beacons_count"`
|
||||
}
|
||||
|
||||
type toolTarget struct {
|
||||
SessionID string
|
||||
BeaconID string
|
||||
}
|
||||
|
||||
// NewExecutor returns the Sliver-backed tool executor used by the AI agent loop.
|
||||
func NewExecutor(backend Backend, conversation *clientpb.AIConversation) serverai.AgenticToolExecutor {
|
||||
return &executor{
|
||||
backend: backend,
|
||||
conversation: conversation,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) ToolDefinitions() []serverai.AgenticToolDefinition {
|
||||
definitions := []serverai.AgenticToolDefinition{
|
||||
{
|
||||
Name: "list_sessions_and_beacons",
|
||||
Description: "List active sessions and beacons so the model can discover valid target IDs before invoking filesystem tools.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_ls",
|
||||
Description: "List the contents of a remote directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"session_id": map[string]any{"type": "string", "description": "Interactive session ID to run the command against."},
|
||||
"beacon_id": map[string]any{"type": "string", "description": "Beacon ID to run the command against."},
|
||||
"path": map[string]any{"type": "string", "description": "Directory path to list. Defaults to . when omitted."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_cat",
|
||||
Description: "Download and return the contents of a remote file. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"session_id": map[string]any{"type": "string", "description": "Interactive session ID to run the command against."},
|
||||
"beacon_id": map[string]any{"type": "string", "description": "Beacon ID to run the command against."},
|
||||
"path": map[string]any{"type": "string", "description": "Remote file path to download."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional byte limit for the download."},
|
||||
"max_lines": map[string]any{"type": "integer", "description": "Optional line limit for the download."},
|
||||
},
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_pwd",
|
||||
Description: "Return the current working directory of a remote session or beacon. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"session_id": map[string]any{"type": "string", "description": "Interactive session ID to run the command against."},
|
||||
"beacon_id": map[string]any{"type": "string", "description": "Beacon ID to run the command against."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
}
|
||||
definitions = append(definitions, filesystemToolDefinitions()...)
|
||||
definitions = append(definitions, systemToolDefinitions()...)
|
||||
definitions = append(definitions, encodingToolDefinitions()...)
|
||||
definitions = append(definitions, storeToolDefinitions()...)
|
||||
definitions = append(definitions, packageToolDefinitions()...)
|
||||
return definitions
|
||||
}
|
||||
|
||||
func (e *executor) CallTool(ctx context.Context, name string, arguments string) (string, error) {
|
||||
switch strings.TrimSpace(name) {
|
||||
case "list_sessions_and_beacons":
|
||||
return e.callListSessionsAndBeacons(ctx)
|
||||
case "fs_ls":
|
||||
var args lsToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callLS(ctx, args)
|
||||
case "fs_cat":
|
||||
var args catToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callCat(ctx, args)
|
||||
case "fs_pwd":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callPWD(ctx, args)
|
||||
case "fs_mv":
|
||||
var args moveToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callMV(ctx, args)
|
||||
case "fs_cp":
|
||||
var args copyToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callCP(ctx, args)
|
||||
case "fs_rm":
|
||||
var args removeToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callRM(ctx, args)
|
||||
case "fs_mkdir":
|
||||
var args mkdirToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callMkdir(ctx, args)
|
||||
case "fs_cd":
|
||||
var args cdToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callCD(ctx, args)
|
||||
case "fs_head":
|
||||
var args headTailToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callHead(ctx, args)
|
||||
case "fs_tail":
|
||||
var args headTailToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callTail(ctx, args)
|
||||
case "fs_chmod":
|
||||
var args chmodToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callChmod(ctx, args)
|
||||
case "fs_chown":
|
||||
var args chownToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callChown(ctx, args)
|
||||
case "fs_chtimes":
|
||||
var args chtimesToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callChtimes(ctx, args)
|
||||
case "fs_mount":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callMount(ctx, args)
|
||||
case "exec_command":
|
||||
var args execToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callExec(ctx, args)
|
||||
case "ifconfig":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callIfconfig(ctx, args)
|
||||
case "netstat":
|
||||
var args netstatToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callNetstat(ctx, args)
|
||||
case "ps":
|
||||
var args psToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callPS(ctx, args)
|
||||
case "env":
|
||||
var args envToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callEnv(ctx, args)
|
||||
case "getpid":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callGetPID(ctx, args)
|
||||
case "getuid":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callGetUID(ctx, args)
|
||||
case "info":
|
||||
var args targetArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callInfo(ctx, args)
|
||||
case "ping":
|
||||
var args pingToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callPing(ctx, args)
|
||||
case "screenshot":
|
||||
var args screenshotToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return e.callScreenshot(ctx, args)
|
||||
default:
|
||||
if result, handled, err := e.callEncodingTool(ctx, name, arguments); handled {
|
||||
return result, err
|
||||
}
|
||||
if result, handled, err := e.callStoreTool(ctx, name, arguments); handled {
|
||||
return result, err
|
||||
}
|
||||
if result, handled, err := e.callPackageTool(ctx, name, arguments); handled {
|
||||
return result, err
|
||||
}
|
||||
return "", fmt.Errorf("unsupported AI tool %q", name)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callListSessionsAndBeacons(ctx context.Context) (string, error) {
|
||||
if e == nil || e.backend == nil {
|
||||
return "", fmt.Errorf("AI tool executor is unavailable")
|
||||
}
|
||||
|
||||
sessionsResp, err := e.backend.GetSessions(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
beaconsResp, err := e.backend.GetBeacons(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := sessionsAndBeaconsResult{
|
||||
Sessions: make([]sessionSummary, 0, len(sessionsResp.GetSessions())),
|
||||
Beacons: make([]beaconSummary, 0, len(beaconsResp.GetBeacons())),
|
||||
}
|
||||
|
||||
for _, session := range sessionsResp.GetSessions() {
|
||||
if session == nil {
|
||||
continue
|
||||
}
|
||||
result.Sessions = append(result.Sessions, sessionSummary{
|
||||
ID: session.ID,
|
||||
Name: session.Name,
|
||||
Transport: session.Transport,
|
||||
RemoteAddress: session.RemoteAddress,
|
||||
Hostname: session.Hostname,
|
||||
Username: session.Username,
|
||||
PID: session.PID,
|
||||
OS: session.OS,
|
||||
Arch: session.Arch,
|
||||
Locale: session.Locale,
|
||||
LastCheckin: session.LastCheckin,
|
||||
IsDead: session.IsDead,
|
||||
Integrity: session.Integrity,
|
||||
})
|
||||
}
|
||||
|
||||
for _, beacon := range beaconsResp.GetBeacons() {
|
||||
if beacon == nil {
|
||||
continue
|
||||
}
|
||||
result.Beacons = append(result.Beacons, beaconSummary{
|
||||
ID: beacon.ID,
|
||||
Name: beacon.Name,
|
||||
Transport: beacon.Transport,
|
||||
RemoteAddress: beacon.RemoteAddress,
|
||||
Hostname: beacon.Hostname,
|
||||
Username: beacon.Username,
|
||||
PID: beacon.PID,
|
||||
OS: beacon.OS,
|
||||
Arch: beacon.Arch,
|
||||
Locale: beacon.Locale,
|
||||
LastCheckin: beacon.LastCheckin,
|
||||
NextCheckin: beacon.NextCheckin,
|
||||
Interval: beacon.Interval,
|
||||
Jitter: beacon.Jitter,
|
||||
TasksCount: beacon.TasksCount,
|
||||
TasksCountCompleted: beacon.TasksCountCompleted,
|
||||
IsDead: beacon.IsDead,
|
||||
Integrity: beacon.Integrity,
|
||||
})
|
||||
}
|
||||
|
||||
result.SessionsCount = len(result.Sessions)
|
||||
result.BeaconsCount = len(result.Beacons)
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callLS(ctx context.Context, args lsToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(args.Path) == "" {
|
||||
args.Path = "."
|
||||
}
|
||||
|
||||
req, isBeacon, callCtx, cancel, err := buildRequestContext(ctx, target)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if cancel != nil {
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
resp, err := e.backend.Ls(callCtx, &sliverpb.LsReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if isBeacon && resp.Response != nil && resp.Response.Async {
|
||||
resolved := &sliverpb.Ls{}
|
||||
if err := waitForBeaconTaskResponse(callCtx, resp.Response.TaskID, resolved); err != nil {
|
||||
return "", err
|
||||
}
|
||||
resp = resolved
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
result := lsResult{
|
||||
Path: resp.Path,
|
||||
Exists: resp.Exists,
|
||||
Timezone: resp.Timezone,
|
||||
TimezoneOffset: resp.TimezoneOffset,
|
||||
Files: make([]lsFileEntry, 0, len(resp.Files)),
|
||||
}
|
||||
for _, fileInfo := range resp.Files {
|
||||
if fileInfo == nil {
|
||||
continue
|
||||
}
|
||||
result.Files = append(result.Files, lsFileEntry{
|
||||
Name: fileInfo.Name,
|
||||
IsDir: fileInfo.IsDir,
|
||||
Size: fileInfo.Size,
|
||||
ModTime: fileInfo.ModTime,
|
||||
Mode: fileInfo.Mode,
|
||||
Link: fileInfo.Link,
|
||||
UID: fileInfo.Uid,
|
||||
GID: fileInfo.Gid,
|
||||
})
|
||||
}
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callCat(ctx context.Context, args catToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(args.Path) == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
|
||||
req, isBeacon, callCtx, cancel, err := buildRequestContext(ctx, target)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if cancel != nil {
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
resp, err := e.backend.Download(callCtx, &sliverpb.DownloadReq{
|
||||
Request: req,
|
||||
RestrictedToFile: true,
|
||||
Path: args.Path,
|
||||
MaxBytes: args.MaxBytes,
|
||||
MaxLines: args.MaxLines,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if isBeacon && resp.Response != nil && resp.Response.Async {
|
||||
resolved := &sliverpb.Download{}
|
||||
if err := waitForBeaconTaskResponse(callCtx, resp.Response.TaskID, resolved); err != nil {
|
||||
return "", err
|
||||
}
|
||||
resp = resolved
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
result, err := buildCatResult(resp, args.Path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callPWD(ctx context.Context, args targetArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
req, isBeacon, callCtx, cancel, err := buildRequestContext(ctx, target)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if cancel != nil {
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
resp, err := e.backend.Pwd(callCtx, &sliverpb.PwdReq{Request: req})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if isBeacon && resp.Response != nil && resp.Response.Async {
|
||||
resolved := &sliverpb.Pwd{}
|
||||
if err := waitForBeaconTaskResponse(callCtx, resp.Response.TaskID, resolved); err != nil {
|
||||
return "", err
|
||||
}
|
||||
resp = resolved
|
||||
if err := genericResponseError(resp.Response); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) resolveTarget(sessionID, beaconID string) (toolTarget, error) {
|
||||
sessionID = strings.TrimSpace(sessionID)
|
||||
beaconID = strings.TrimSpace(beaconID)
|
||||
if sessionID != "" && beaconID != "" {
|
||||
return toolTarget{}, fmt.Errorf("provide only one of session_id or beacon_id")
|
||||
}
|
||||
if sessionID == "" && beaconID == "" && e != nil && e.conversation != nil {
|
||||
sessionID = strings.TrimSpace(e.conversation.GetTargetSessionID())
|
||||
beaconID = strings.TrimSpace(e.conversation.GetTargetBeaconID())
|
||||
}
|
||||
if sessionID == "" && beaconID == "" {
|
||||
return toolTarget{}, fmt.Errorf("session_id or beacon_id is required; call list_sessions_and_beacons or start a conversation from an active target")
|
||||
}
|
||||
return toolTarget{
|
||||
SessionID: sessionID,
|
||||
BeaconID: beaconID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func buildRequestContext(ctx context.Context, target toolTarget) (*commonpb.Request, bool, context.Context, context.CancelFunc, error) {
|
||||
if strings.TrimSpace(target.SessionID) != "" && strings.TrimSpace(target.BeaconID) != "" {
|
||||
return nil, false, nil, nil, fmt.Errorf("provide only one of session_id or beacon_id")
|
||||
}
|
||||
|
||||
timeoutSeconds := aiToolDefaultWaitSeconds
|
||||
req := &commonpb.Request{
|
||||
Timeout: (int64(time.Second) * timeoutSeconds) - 1,
|
||||
}
|
||||
isBeacon := false
|
||||
|
||||
if strings.TrimSpace(target.SessionID) != "" {
|
||||
req.SessionID = strings.TrimSpace(target.SessionID)
|
||||
req.Async = false
|
||||
} else {
|
||||
req.BeaconID = strings.TrimSpace(target.BeaconID)
|
||||
req.Async = true
|
||||
isBeacon = true
|
||||
}
|
||||
|
||||
callCtx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
|
||||
return req, isBeacon, callCtx, cancel, nil
|
||||
}
|
||||
|
||||
func decodeToolArgs(raw string, out any) error {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
raw = "{}"
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), out); err != nil {
|
||||
return fmt.Errorf("invalid tool arguments: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func marshalToolResult(result any) (string, error) {
|
||||
data, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func genericResponseError(resp *commonpb.Response) error {
|
||||
if resp == nil || strings.TrimSpace(resp.Err) == "" {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("%s", strings.TrimSpace(resp.Err))
|
||||
}
|
||||
|
||||
func waitForBeaconTaskResponse(ctx context.Context, taskID string, out proto.Message) error {
|
||||
task, err := waitForBeaconTask(ctx, taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
state := strings.ToLower(strings.TrimSpace(task.GetState()))
|
||||
if state != "completed" {
|
||||
if state == "" {
|
||||
state = "unknown"
|
||||
}
|
||||
return fmt.Errorf("beacon task %s %s", task.GetID(), state)
|
||||
}
|
||||
if len(task.GetResponse()) == 0 {
|
||||
return fmt.Errorf("beacon task %s returned empty response", task.GetID())
|
||||
}
|
||||
return proto.Unmarshal(task.GetResponse(), out)
|
||||
}
|
||||
|
||||
func waitForBeaconTask(ctx context.Context, taskID string) (*clientpb.BeaconTask, error) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
task, err := db.BeaconTaskByID(taskID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(task.GetState())) {
|
||||
case "completed", "failed", "canceled":
|
||||
return task, nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func buildCatResult(download *sliverpb.Download, requestedPath string) (*catResult, error) {
|
||||
data := download.Data
|
||||
decoded := false
|
||||
if download.Encoder == "gzip" {
|
||||
decoded = true
|
||||
var err error
|
||||
data, err = new(encoders.Gzip).Decode(download.Data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
result := &catResult{
|
||||
Path: download.Path,
|
||||
RequestedPath: requestedPath,
|
||||
Exists: download.Exists,
|
||||
IsDir: download.IsDir,
|
||||
Encoder: download.Encoder,
|
||||
Decoded: decoded,
|
||||
ByteLen: len(data),
|
||||
DataBase64: base64.StdEncoding.EncodeToString(data),
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
if utf8.Valid(data) {
|
||||
result.Text = string(data)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,574 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
)
|
||||
|
||||
type moveToolArgs struct {
|
||||
targetArgs
|
||||
Src string `json:"src,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
}
|
||||
|
||||
type copyToolArgs struct {
|
||||
targetArgs
|
||||
Src string `json:"src,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
}
|
||||
|
||||
type removeToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
Recursive bool `json:"recursive,omitempty"`
|
||||
Force bool `json:"force,omitempty"`
|
||||
}
|
||||
|
||||
type mkdirToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
}
|
||||
|
||||
type cdToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
}
|
||||
|
||||
type headTailToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
MaxBytes int64 `json:"max_bytes,omitempty"`
|
||||
MaxLines int64 `json:"max_lines,omitempty"`
|
||||
}
|
||||
|
||||
type chmodToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
FileMode string `json:"file_mode,omitempty"`
|
||||
Recursive bool `json:"recursive,omitempty"`
|
||||
}
|
||||
|
||||
type chownToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
UID string `json:"uid,omitempty"`
|
||||
GID string `json:"gid,omitempty"`
|
||||
Recursive bool `json:"recursive,omitempty"`
|
||||
}
|
||||
|
||||
type chtimesToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
ATimeUnix int64 `json:"atime_unix,omitempty"`
|
||||
MTimeUnix int64 `json:"mtime_unix,omitempty"`
|
||||
}
|
||||
|
||||
type moveResult struct {
|
||||
Src string `json:"src"`
|
||||
Dst string `json:"dst"`
|
||||
}
|
||||
|
||||
type copyResult struct {
|
||||
Src string `json:"src"`
|
||||
Dst string `json:"dst"`
|
||||
BytesWritten int64 `json:"bytes_written"`
|
||||
}
|
||||
|
||||
type contentResult struct {
|
||||
Operation string `json:"operation"`
|
||||
RequestedMaxBytes int64 `json:"requested_max_bytes,omitempty"`
|
||||
RequestedMaxLines int64 `json:"requested_max_lines,omitempty"`
|
||||
catResult
|
||||
}
|
||||
|
||||
type chtimesResult struct {
|
||||
Path string `json:"path"`
|
||||
ATimeUnix int64 `json:"atime_unix"`
|
||||
MTimeUnix int64 `json:"mtime_unix"`
|
||||
}
|
||||
|
||||
type mountInfoResult struct {
|
||||
VolumeName string `json:"volume_name"`
|
||||
VolumeType string `json:"volume_type"`
|
||||
MountPoint string `json:"mount_point"`
|
||||
Label string `json:"label"`
|
||||
FileSystem string `json:"file_system"`
|
||||
UsedSpace uint64 `json:"used_space"`
|
||||
FreeSpace uint64 `json:"free_space"`
|
||||
TotalSpace uint64 `json:"total_space"`
|
||||
MountOptions string `json:"mount_options"`
|
||||
}
|
||||
|
||||
type mountResult struct {
|
||||
Mounts []mountInfoResult `json:"mounts"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
func filesystemToolDefinitions() []serverai.AgenticToolDefinition {
|
||||
return []serverai.AgenticToolDefinition{
|
||||
{
|
||||
Name: "fs_mv",
|
||||
Description: "Move or rename a remote file or directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"src": map[string]any{"type": "string", "description": "Source path to move."},
|
||||
"dst": map[string]any{"type": "string", "description": "Destination path."},
|
||||
}),
|
||||
"required": []string{"src", "dst"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_cp",
|
||||
Description: "Copy a remote file. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"src": map[string]any{"type": "string", "description": "Source path to copy."},
|
||||
"dst": map[string]any{"type": "string", "description": "Destination path."},
|
||||
}),
|
||||
"required": []string{"src", "dst"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_rm",
|
||||
Description: "Remove a remote file or directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Path to remove."},
|
||||
"recursive": map[string]any{"type": "boolean", "description": "Recursively remove directories."},
|
||||
"force": map[string]any{"type": "boolean", "description": "Force removal."},
|
||||
}),
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_mkdir",
|
||||
Description: "Create a remote directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Directory path to create."},
|
||||
}),
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_cd",
|
||||
Description: "Change the remote working directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Directory to change into. Defaults to . when omitted."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_head",
|
||||
Description: "Return the first lines or bytes from a remote file. Defaults to the first 10 lines when no limit is provided.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Remote file path to read."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional number of bytes to read from the start of the file."},
|
||||
"max_lines": map[string]any{"type": "integer", "description": "Optional number of lines to read from the start of the file."},
|
||||
}),
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_tail",
|
||||
Description: "Return the last lines or bytes from a remote file. Defaults to the last 10 lines when no limit is provided.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Remote file path to read."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional number of bytes to read from the end of the file."},
|
||||
"max_lines": map[string]any{"type": "integer", "description": "Optional number of lines to read from the end of the file."},
|
||||
}),
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_chmod",
|
||||
Description: "Change permissions on a remote file or directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Path to change permissions on."},
|
||||
"file_mode": map[string]any{"type": "string", "description": "Mode string such as 0644 or 755."},
|
||||
"recursive": map[string]any{"type": "boolean", "description": "Apply recursively."},
|
||||
}),
|
||||
"required": []string{"path", "file_mode"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_chown",
|
||||
Description: "Change ownership on a remote file or directory. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Path to change ownership on."},
|
||||
"uid": map[string]any{"type": "string", "description": "User ID to assign."},
|
||||
"gid": map[string]any{"type": "string", "description": "Group ID to assign."},
|
||||
"recursive": map[string]any{"type": "boolean", "description": "Apply recursively."},
|
||||
}),
|
||||
"required": []string{"path", "uid", "gid"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_chtimes",
|
||||
Description: "Change access and modification times on a remote file or directory using Unix timestamps in seconds.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Path to modify."},
|
||||
"atime_unix": map[string]any{"type": "integer", "description": "Access time as a Unix timestamp in seconds."},
|
||||
"mtime_unix": map[string]any{"type": "integer", "description": "Modification time as a Unix timestamp in seconds."},
|
||||
}),
|
||||
"required": []string{"path", "atime_unix", "mtime_unix"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "fs_mount",
|
||||
Description: "List mounted filesystems or drives on the remote host. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": targetSchemaProperties(),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callMV(ctx context.Context, args moveToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Src == "" || args.Dst == "" {
|
||||
return "", fmt.Errorf("src and dst are required")
|
||||
}
|
||||
|
||||
_, err = callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Mv, error) {
|
||||
return e.backend.Mv(callCtx, &sliverpb.MvReq{Request: req, Src: args.Src, Dst: args.Dst})
|
||||
},
|
||||
func() *sliverpb.Mv { return &sliverpb.Mv{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(moveResult{Src: args.Src, Dst: args.Dst})
|
||||
}
|
||||
|
||||
func (e *executor) callCP(ctx context.Context, args copyToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Src == "" || args.Dst == "" {
|
||||
return "", fmt.Errorf("src and dst are required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Cp, error) {
|
||||
return e.backend.Cp(callCtx, &sliverpb.CpReq{Request: req, Src: args.Src, Dst: args.Dst})
|
||||
},
|
||||
func() *sliverpb.Cp { return &sliverpb.Cp{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(copyResult{Src: args.Src, Dst: args.Dst, BytesWritten: resp.BytesWritten})
|
||||
}
|
||||
|
||||
func (e *executor) callRM(ctx context.Context, args removeToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Rm, error) {
|
||||
return e.backend.Rm(callCtx, &sliverpb.RmReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
Recursive: args.Recursive,
|
||||
Force: args.Force,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Rm { return &sliverpb.Rm{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) callMkdir(ctx context.Context, args mkdirToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Mkdir, error) {
|
||||
return e.backend.Mkdir(callCtx, &sliverpb.MkdirReq{Request: req, Path: args.Path})
|
||||
},
|
||||
func() *sliverpb.Mkdir { return &sliverpb.Mkdir{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) callCD(ctx context.Context, args cdToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
args.Path = "."
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Pwd, error) {
|
||||
return e.backend.Cd(callCtx, &sliverpb.CdReq{Request: req, Path: args.Path})
|
||||
},
|
||||
func() *sliverpb.Pwd { return &sliverpb.Pwd{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) callHead(ctx context.Context, args headTailToolArgs) (string, error) {
|
||||
return e.callHeadTail(ctx, args, true)
|
||||
}
|
||||
|
||||
func (e *executor) callTail(ctx context.Context, args headTailToolArgs) (string, error) {
|
||||
return e.callHeadTail(ctx, args, false)
|
||||
}
|
||||
|
||||
func (e *executor) callHeadTail(ctx context.Context, args headTailToolArgs, head bool) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
if args.MaxBytes < 0 || args.MaxLines < 0 {
|
||||
return "", fmt.Errorf("max_bytes and max_lines must be non-negative")
|
||||
}
|
||||
if args.MaxBytes != 0 && args.MaxLines != 0 {
|
||||
return "", fmt.Errorf("provide only one of max_bytes or max_lines")
|
||||
}
|
||||
if args.MaxBytes == 0 && args.MaxLines == 0 {
|
||||
args.MaxLines = 10
|
||||
}
|
||||
|
||||
requestedBytes := args.MaxBytes
|
||||
requestedLines := args.MaxLines
|
||||
if !head {
|
||||
requestedBytes *= -1
|
||||
requestedLines *= -1
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Download, error) {
|
||||
return e.backend.Download(callCtx, &sliverpb.DownloadReq{
|
||||
Request: req,
|
||||
RestrictedToFile: true,
|
||||
Path: args.Path,
|
||||
MaxBytes: requestedBytes,
|
||||
MaxLines: requestedLines,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Download { return &sliverpb.Download{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
built, err := buildCatResult(resp, args.Path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
operation := "head"
|
||||
if !head {
|
||||
operation = "tail"
|
||||
}
|
||||
return marshalToolResult(contentResult{
|
||||
Operation: operation,
|
||||
RequestedMaxBytes: args.MaxBytes,
|
||||
RequestedMaxLines: args.MaxLines,
|
||||
catResult: *built,
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callChmod(ctx context.Context, args chmodToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" || args.FileMode == "" {
|
||||
return "", fmt.Errorf("path and file_mode are required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Chmod, error) {
|
||||
return e.backend.Chmod(callCtx, &sliverpb.ChmodReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
FileMode: args.FileMode,
|
||||
Recursive: args.Recursive,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Chmod { return &sliverpb.Chmod{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) callChown(ctx context.Context, args chownToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" || args.UID == "" || args.GID == "" {
|
||||
return "", fmt.Errorf("path, uid, and gid are required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Chown, error) {
|
||||
return e.backend.Chown(callCtx, &sliverpb.ChownReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
Uid: args.UID,
|
||||
Gid: args.GID,
|
||||
Recursive: args.Recursive,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Chown { return &sliverpb.Chown{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pathResult{Path: resp.Path})
|
||||
}
|
||||
|
||||
func (e *executor) callChtimes(ctx context.Context, args chtimesToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Chtimes, error) {
|
||||
return e.backend.Chtimes(callCtx, &sliverpb.ChtimesReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
ATime: args.ATimeUnix,
|
||||
MTime: args.MTimeUnix,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Chtimes { return &sliverpb.Chtimes{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(chtimesResult{
|
||||
Path: resp.Path,
|
||||
ATimeUnix: args.ATimeUnix,
|
||||
MTimeUnix: args.MTimeUnix,
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callMount(ctx context.Context, args targetArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Mount, error) {
|
||||
return e.backend.Mount(callCtx, &sliverpb.MountReq{Request: req})
|
||||
},
|
||||
func() *sliverpb.Mount { return &sliverpb.Mount{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := mountResult{
|
||||
Mounts: make([]mountInfoResult, 0, len(resp.Info)),
|
||||
}
|
||||
for _, mountInfo := range resp.Info {
|
||||
if mountInfo == nil {
|
||||
continue
|
||||
}
|
||||
result.Mounts = append(result.Mounts, mountInfoResult{
|
||||
VolumeName: mountInfo.VolumeName,
|
||||
VolumeType: mountInfo.VolumeType,
|
||||
MountPoint: mountInfo.MountPoint,
|
||||
Label: mountInfo.Label,
|
||||
FileSystem: mountInfo.FileSystem,
|
||||
UsedSpace: mountInfo.UsedSpace,
|
||||
FreeSpace: mountInfo.FreeSpace,
|
||||
TotalSpace: mountInfo.TotalSpace,
|
||||
MountOptions: mountInfo.MountOptions,
|
||||
})
|
||||
}
|
||||
result.Count = len(result.Mounts)
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
type responseMessage interface {
|
||||
proto.Message
|
||||
GetResponse() *commonpb.Response
|
||||
}
|
||||
|
||||
func targetSchemaProperties() map[string]any {
|
||||
return map[string]any{
|
||||
"session_id": map[string]any{"type": "string", "description": "Interactive session ID to run the command against."},
|
||||
"beacon_id": map[string]any{"type": "string", "description": "Beacon ID to run the command against."},
|
||||
}
|
||||
}
|
||||
|
||||
func mergeSchemaProperties(extra map[string]any) map[string]any {
|
||||
properties := targetSchemaProperties()
|
||||
for key, value := range extra {
|
||||
properties[key] = value
|
||||
}
|
||||
return properties
|
||||
}
|
||||
|
||||
func callTargetRPC[T responseMessage](
|
||||
ctx context.Context,
|
||||
target toolTarget,
|
||||
invoke func(context.Context, *commonpb.Request) (T, error),
|
||||
newResponse func() T,
|
||||
) (T, error) {
|
||||
var zero T
|
||||
|
||||
req, isBeacon, callCtx, cancel, err := buildRequestContext(ctx, target)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
if cancel != nil {
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
resp, err := invoke(callCtx, req)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
if err := genericResponseError(resp.GetResponse()); err != nil {
|
||||
return zero, err
|
||||
}
|
||||
if isBeacon && resp.GetResponse() != nil && resp.GetResponse().Async {
|
||||
resolved := newResponse()
|
||||
if err := waitForBeaconTaskResponse(callCtx, resp.GetResponse().TaskID, resolved); err != nil {
|
||||
return zero, err
|
||||
}
|
||||
resp = resolved
|
||||
if err := genericResponseError(resp.GetResponse()); err != nil {
|
||||
return zero, err
|
||||
}
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (e *executor) lookupTargetMetadata(ctx context.Context, sessionID, beaconID string) (*clientpb.Session, *clientpb.Beacon, error) {
|
||||
target, err := e.resolveTarget(sessionID, beaconID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if e == nil || e.backend == nil {
|
||||
return nil, nil, fmt.Errorf("AI tool executor is unavailable")
|
||||
}
|
||||
|
||||
if target.SessionID != "" {
|
||||
sessionsResp, err := e.backend.GetSessions(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
for _, session := range sessionsResp.GetSessions() {
|
||||
if session != nil && session.ID == target.SessionID {
|
||||
return session, nil, nil
|
||||
}
|
||||
}
|
||||
for _, session := range sessionsResp.GetSessions() {
|
||||
if session != nil && strings.HasPrefix(session.ID, target.SessionID) {
|
||||
return session, nil, nil
|
||||
}
|
||||
}
|
||||
return nil, nil, fmt.Errorf("session %q not found", target.SessionID)
|
||||
}
|
||||
|
||||
beaconsResp, err := e.backend.GetBeacons(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
for _, beacon := range beaconsResp.GetBeacons() {
|
||||
if beacon != nil && beacon.ID == target.BeaconID {
|
||||
return nil, beacon, nil
|
||||
}
|
||||
}
|
||||
for _, beacon := range beaconsResp.GetBeacons() {
|
||||
if beacon != nil && strings.HasPrefix(beacon.ID, target.BeaconID) {
|
||||
return nil, beacon, nil
|
||||
}
|
||||
}
|
||||
return nil, nil, fmt.Errorf("beacon %q not found", target.BeaconID)
|
||||
}
|
||||
|
||||
func bytesToTextAndBase64(data []byte) (string, string) {
|
||||
if len(data) == 0 {
|
||||
return "", ""
|
||||
}
|
||||
if utf8.Valid(data) {
|
||||
return string(data), ""
|
||||
}
|
||||
return "", base64.StdEncoding.EncodeToString(data)
|
||||
}
|
||||
|
||||
func sha256Hex(data []byte) string {
|
||||
sum := sha256.Sum256(data)
|
||||
return fmt.Sprintf("%x", sum)
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/core"
|
||||
)
|
||||
|
||||
func buildBOFExtensionArgs(command *aiExtensionCommand, bofData []byte, args []string) ([]byte, error) {
|
||||
argsBuffer := core.BOFArgsBuffer{Buffer: new(bytes.Buffer)}
|
||||
if err := argsBuffer.AddString(command.Entrypoint); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := argsBuffer.AddData(bofData); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
parsedArgs, err := parseBOFArguments(command, args)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := argsBuffer.AddData(parsedArgs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return argsBuffer.GetBuffer()
|
||||
}
|
||||
|
||||
func parseBOFArguments(command *aiExtensionCommand, args []string) ([]byte, error) {
|
||||
fs, stringValues, wstringValues, intValues, shortValues, fileValues, err := setupBOFFlags(args, command)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
argsBuffer := core.BOFArgsBuffer{Buffer: new(bytes.Buffer)}
|
||||
missingRequired := []string{}
|
||||
|
||||
for _, argDef := range command.Arguments {
|
||||
if argDef == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
var provided bool
|
||||
switch argDef.Type {
|
||||
case "string":
|
||||
provided, err = processBOFStringArg(fs, argDef, stringValues, argsBuffer)
|
||||
case "wstring":
|
||||
provided, err = processBOFWStringArg(fs, argDef, wstringValues, argsBuffer)
|
||||
case "int", "integer":
|
||||
provided, err = processBOFIntArg(fs, argDef, intValues, argsBuffer)
|
||||
case "short":
|
||||
provided, err = processBOFShortArg(fs, argDef, shortValues, argsBuffer)
|
||||
case "file":
|
||||
provided, err = processBOFFileArg(fs, argDef, fileValues, argsBuffer)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported BOF argument type %q", argDef.Type)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if provided {
|
||||
continue
|
||||
}
|
||||
if argDef.Optional {
|
||||
if err := applyBOFDefaultValue(argDef, argsBuffer); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
missingRequired = append(missingRequired, argDef.Name)
|
||||
}
|
||||
|
||||
if len(missingRequired) > 0 {
|
||||
return nil, fmt.Errorf("required BOF arguments were not provided: %s", strings.Join(missingRequired, ", "))
|
||||
}
|
||||
return argsBuffer.GetBuffer()
|
||||
}
|
||||
|
||||
func setupBOFFlags(args []string, command *aiExtensionCommand) (
|
||||
*flag.FlagSet,
|
||||
map[string]*string,
|
||||
map[string]*string,
|
||||
map[string]*int,
|
||||
map[string]*int,
|
||||
map[string]*string,
|
||||
error,
|
||||
) {
|
||||
fs := flag.NewFlagSet("sliver-ai-bof", flag.ContinueOnError)
|
||||
fs.SetOutput(io.Discard)
|
||||
|
||||
stringValues := map[string]*string{}
|
||||
wstringValues := map[string]*string{}
|
||||
intValues := map[string]*int{}
|
||||
shortValues := map[string]*int{}
|
||||
fileValues := map[string]*string{}
|
||||
|
||||
for _, arg := range command.Arguments {
|
||||
if arg == nil {
|
||||
continue
|
||||
}
|
||||
switch arg.Type {
|
||||
case "string":
|
||||
stringValues[arg.Name] = fs.String(arg.Name, "", arg.Desc)
|
||||
case "wstring":
|
||||
wstringValues[arg.Name] = fs.String(arg.Name, "", arg.Desc)
|
||||
case "int", "integer":
|
||||
intValues[arg.Name] = fs.Int(arg.Name, 0, arg.Desc)
|
||||
case "short":
|
||||
shortValues[arg.Name] = fs.Int(arg.Name, 0, arg.Desc)
|
||||
case "file":
|
||||
fileValues[arg.Name] = fs.String(arg.Name, "", arg.Desc)
|
||||
default:
|
||||
return nil, nil, nil, nil, nil, nil, fmt.Errorf("unsupported BOF argument type %q", arg.Type)
|
||||
}
|
||||
}
|
||||
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
return fs, stringValues, wstringValues, intValues, shortValues, fileValues, nil
|
||||
}
|
||||
|
||||
func processBOFStringArg(fs *flag.FlagSet, argDef *aiExtensionArgument, values map[string]*string, argsBuffer core.BOFArgsBuffer) (bool, error) {
|
||||
ptr := values[argDef.Name]
|
||||
if ptr == nil || !bofFlagWasProvided(fs, argDef.Name) {
|
||||
return false, nil
|
||||
}
|
||||
if err := argsBuffer.AddString(*ptr); err != nil {
|
||||
return false, fmt.Errorf("failed to add string argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func processBOFWStringArg(fs *flag.FlagSet, argDef *aiExtensionArgument, values map[string]*string, argsBuffer core.BOFArgsBuffer) (bool, error) {
|
||||
ptr := values[argDef.Name]
|
||||
if ptr == nil || !bofFlagWasProvided(fs, argDef.Name) {
|
||||
return false, nil
|
||||
}
|
||||
if err := argsBuffer.AddWString(*ptr); err != nil {
|
||||
return false, fmt.Errorf("failed to add wstring argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func processBOFIntArg(fs *flag.FlagSet, argDef *aiExtensionArgument, values map[string]*int, argsBuffer core.BOFArgsBuffer) (bool, error) {
|
||||
ptr := values[argDef.Name]
|
||||
if ptr == nil || !bofFlagWasProvided(fs, argDef.Name) {
|
||||
return false, nil
|
||||
}
|
||||
if err := argsBuffer.AddInt(uint32(*ptr)); err != nil {
|
||||
return false, fmt.Errorf("failed to add integer argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func processBOFShortArg(fs *flag.FlagSet, argDef *aiExtensionArgument, values map[string]*int, argsBuffer core.BOFArgsBuffer) (bool, error) {
|
||||
ptr := values[argDef.Name]
|
||||
if ptr == nil || !bofFlagWasProvided(fs, argDef.Name) {
|
||||
return false, nil
|
||||
}
|
||||
if err := argsBuffer.AddShort(uint16(*ptr)); err != nil {
|
||||
return false, fmt.Errorf("failed to add short argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func processBOFFileArg(fs *flag.FlagSet, argDef *aiExtensionArgument, values map[string]*string, argsBuffer core.BOFArgsBuffer) (bool, error) {
|
||||
ptr := values[argDef.Name]
|
||||
if ptr == nil || !bofFlagWasProvided(fs, argDef.Name) {
|
||||
return false, nil
|
||||
}
|
||||
data, err := os.ReadFile(*ptr)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to read file argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
if err := argsBuffer.AddData(data); err != nil {
|
||||
return false, fmt.Errorf("failed to add file argument %s: %w", argDef.Name, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func applyBOFDefaultValue(argDef *aiExtensionArgument, argsBuffer core.BOFArgsBuffer) error {
|
||||
if argDef == nil {
|
||||
return nil
|
||||
}
|
||||
if argDef.Default == nil {
|
||||
return addBOFZeroValue(argDef.Type, argsBuffer)
|
||||
}
|
||||
|
||||
switch argDef.Type {
|
||||
case "string":
|
||||
return argsBuffer.AddString(fmt.Sprint(argDef.Default))
|
||||
case "wstring":
|
||||
return argsBuffer.AddWString(fmt.Sprint(argDef.Default))
|
||||
case "int", "integer":
|
||||
value, err := bofDefaultIntValue(argDef.Default)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid default integer value for %s: %w", argDef.Name, err)
|
||||
}
|
||||
return argsBuffer.AddInt(uint32(value))
|
||||
case "short":
|
||||
value, err := bofDefaultIntValue(argDef.Default)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid default short value for %s: %w", argDef.Name, err)
|
||||
}
|
||||
return argsBuffer.AddShort(uint16(value))
|
||||
case "file":
|
||||
path := fmt.Sprint(argDef.Default)
|
||||
if path == "" {
|
||||
return addBOFZeroValue(argDef.Type, argsBuffer)
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read default file for %s: %w", argDef.Name, err)
|
||||
}
|
||||
return argsBuffer.AddData(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported BOF argument type %q", argDef.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func addBOFZeroValue(argType string, argsBuffer core.BOFArgsBuffer) error {
|
||||
switch argType {
|
||||
case "string":
|
||||
return argsBuffer.AddString("")
|
||||
case "wstring":
|
||||
return argsBuffer.AddWString("")
|
||||
case "int", "integer":
|
||||
return argsBuffer.AddInt(0)
|
||||
case "short":
|
||||
return argsBuffer.AddShort(0)
|
||||
case "file":
|
||||
return argsBuffer.AddData([]byte{})
|
||||
default:
|
||||
return fmt.Errorf("unsupported BOF argument type %q", argType)
|
||||
}
|
||||
}
|
||||
|
||||
func bofDefaultIntValue(raw any) (int64, error) {
|
||||
switch value := raw.(type) {
|
||||
case float64:
|
||||
return int64(value), nil
|
||||
case float32:
|
||||
return int64(value), nil
|
||||
case int:
|
||||
return int64(value), nil
|
||||
case int8:
|
||||
return int64(value), nil
|
||||
case int16:
|
||||
return int64(value), nil
|
||||
case int32:
|
||||
return int64(value), nil
|
||||
case int64:
|
||||
return value, nil
|
||||
case uint:
|
||||
return int64(value), nil
|
||||
case uint8:
|
||||
return int64(value), nil
|
||||
case uint16:
|
||||
return int64(value), nil
|
||||
case uint32:
|
||||
return int64(value), nil
|
||||
case uint64:
|
||||
return int64(value), nil
|
||||
case string:
|
||||
return strconv.ParseInt(strings.TrimSpace(value), 10, 64)
|
||||
default:
|
||||
return 0, fmt.Errorf("unsupported numeric default type %T", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func bofFlagWasProvided(fs *flag.FlagSet, name string) bool {
|
||||
provided := false
|
||||
fs.Visit(func(f *flag.Flag) {
|
||||
if f.Name == name {
|
||||
provided = true
|
||||
}
|
||||
})
|
||||
return provided
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,463 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
serverassets "github.com/bishopfox/sliver/server/assets"
|
||||
)
|
||||
|
||||
func TestSearchAliasesIncludesTargetCompatibility(t *testing.T) {
|
||||
rootDir := t.TempDir()
|
||||
t.Setenv("SLIVER_ROOT_DIR", rootDir)
|
||||
|
||||
aliasDir := filepath.Join(serverassets.GetAIAliasesDir(), "Rubeus")
|
||||
if err := os.MkdirAll(aliasDir, 0o700); err != nil {
|
||||
t.Fatalf("mkdir alias dir: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(aliasDir, "Rubeus.exe"), []byte("alias-binary"), 0o600); err != nil {
|
||||
t.Fatalf("write alias artifact: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(aliasDir, aiAliasManifestFileName), []byte(`{
|
||||
"name":"Rubeus",
|
||||
"version":"1.0.0",
|
||||
"command_name":"rubeus",
|
||||
"original_author":"GhostPack",
|
||||
"repo_url":"https://example.test/rubeus",
|
||||
"help":"Kerberos abuse helper",
|
||||
"entrypoint":"Main",
|
||||
"allow_args":true,
|
||||
"default_args":"",
|
||||
"is_reflective":false,
|
||||
"is_assembly":true,
|
||||
"files":[{"os":"windows","arch":"amd64","path":"Rubeus.exe"}]
|
||||
}`), 0o600); err != nil {
|
||||
t.Fatalf("write alias manifest: %v", err)
|
||||
}
|
||||
|
||||
backend := &fakePackageBackend{
|
||||
sessions: &clientpb.Sessions{
|
||||
Sessions: []*clientpb.Session{
|
||||
{ID: "session-1", OS: "windows", Arch: "amd64", Hostname: "winhost"},
|
||||
},
|
||||
},
|
||||
}
|
||||
executor := &executor{
|
||||
backend: backend,
|
||||
conversation: &clientpb.AIConversation{
|
||||
TargetSessionID: "session-1",
|
||||
},
|
||||
}
|
||||
|
||||
raw, err := executor.callSearchAliases(context.Background(), searchPackagesArgs{
|
||||
Query: "kerberos",
|
||||
OnlyCompatible: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("search aliases: %v", err)
|
||||
}
|
||||
|
||||
var resp aliasSearchResponse
|
||||
if err := json.Unmarshal([]byte(raw), &resp); err != nil {
|
||||
t.Fatalf("unmarshal search result: %v", err)
|
||||
}
|
||||
if resp.ReturnedCount != 1 || resp.TotalMatches != 1 {
|
||||
t.Fatalf("unexpected alias search counts: %+v", resp)
|
||||
}
|
||||
if resp.Target == nil || resp.Target.SessionID != "session-1" {
|
||||
t.Fatalf("expected session target in response, got %+v", resp.Target)
|
||||
}
|
||||
if len(resp.Results) != 1 {
|
||||
t.Fatalf("expected one alias result, got %+v", resp.Results)
|
||||
}
|
||||
result := resp.Results[0]
|
||||
if !result.Compatible || !result.CompatibilityChecked {
|
||||
t.Fatalf("expected compatible alias result, got %+v", result)
|
||||
}
|
||||
if result.ExecutionMode != "assembly" {
|
||||
t.Fatalf("unexpected alias execution mode: %+v", result)
|
||||
}
|
||||
if !strings.HasSuffix(result.ArtifactPath, filepath.Join("Rubeus", "Rubeus.exe")) {
|
||||
t.Fatalf("unexpected alias artifact path: %q", result.ArtifactPath)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteAliasRunsExecuteAssembly(t *testing.T) {
|
||||
rootDir := t.TempDir()
|
||||
t.Setenv("SLIVER_ROOT_DIR", rootDir)
|
||||
|
||||
aliasDir := filepath.Join(serverassets.GetAIAliasesDir(), "Seatbelt")
|
||||
if err := os.MkdirAll(aliasDir, 0o700); err != nil {
|
||||
t.Fatalf("mkdir alias dir: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(aliasDir, "Seatbelt.exe"), []byte("seatbelt-binary"), 0o600); err != nil {
|
||||
t.Fatalf("write alias artifact: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(aliasDir, aiAliasManifestFileName), []byte(`{
|
||||
"name":"Seatbelt",
|
||||
"version":"1.0.0",
|
||||
"command_name":"seatbelt",
|
||||
"original_author":"GhostPack",
|
||||
"repo_url":"https://example.test/seatbelt",
|
||||
"help":"Host survey helper",
|
||||
"entrypoint":"Main",
|
||||
"allow_args":true,
|
||||
"default_args":"",
|
||||
"is_reflective":false,
|
||||
"is_assembly":true,
|
||||
"files":[{"os":"windows","arch":"amd64","path":"Seatbelt.exe"}]
|
||||
}`), 0o600); err != nil {
|
||||
t.Fatalf("write alias manifest: %v", err)
|
||||
}
|
||||
|
||||
backend := &fakePackageBackend{
|
||||
sessions: &clientpb.Sessions{
|
||||
Sessions: []*clientpb.Session{
|
||||
{ID: "session-1", OS: "windows", Arch: "amd64", Hostname: "winhost"},
|
||||
},
|
||||
},
|
||||
executeAssemblyFn: func(_ context.Context, req *sliverpb.ExecuteAssemblyReq) (*sliverpb.ExecuteAssembly, error) {
|
||||
return &sliverpb.ExecuteAssembly{
|
||||
Output: []byte("assembly-output"),
|
||||
Response: &commonpb.Response{},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
executor := &executor{
|
||||
backend: backend,
|
||||
conversation: &clientpb.AIConversation{
|
||||
TargetSessionID: "session-1",
|
||||
},
|
||||
}
|
||||
|
||||
raw, err := executor.callExecuteAlias(context.Background(), executeAliasArgs{
|
||||
CommandName: "seatbelt",
|
||||
Args: []string{"WindowsCredentialFiles"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("execute alias: %v", err)
|
||||
}
|
||||
|
||||
if len(backend.executeAssemblyReqs) != 1 {
|
||||
t.Fatalf("expected execute-assembly request, got %d", len(backend.executeAssemblyReqs))
|
||||
}
|
||||
req := backend.executeAssemblyReqs[0]
|
||||
if req.GetRequest().GetSessionID() != "session-1" {
|
||||
t.Fatalf("unexpected target request: %+v", req.GetRequest())
|
||||
}
|
||||
if req.GetProcess() != aiAliasDefaultHostProcess["windows"] {
|
||||
t.Fatalf("unexpected default process: %q", req.GetProcess())
|
||||
}
|
||||
if len(req.GetArguments()) != 1 || req.GetArguments()[0] != "WindowsCredentialFiles" {
|
||||
t.Fatalf("unexpected assembly args: %+v", req.GetArguments())
|
||||
}
|
||||
|
||||
var resp aliasExecutionResult
|
||||
if err := json.Unmarshal([]byte(raw), &resp); err != nil {
|
||||
t.Fatalf("unmarshal alias execution result: %v", err)
|
||||
}
|
||||
if resp.OutputText != "assembly-output" {
|
||||
t.Fatalf("unexpected alias output: %+v", resp)
|
||||
}
|
||||
if resp.ExecutionMode != "assembly" {
|
||||
t.Fatalf("unexpected alias execution mode: %+v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteExtensionRegistersDependencyForBOF(t *testing.T) {
|
||||
rootDir := t.TempDir()
|
||||
t.Setenv("SLIVER_ROOT_DIR", rootDir)
|
||||
|
||||
coffLoaderDir := filepath.Join(serverassets.GetAIExtensionsDir(), "coff-loader")
|
||||
if err := os.MkdirAll(coffLoaderDir, 0o700); err != nil {
|
||||
t.Fatalf("mkdir coff-loader dir: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(coffLoaderDir, "coff-loader.x64.dll"), []byte("coff-loader-binary"), 0o600); err != nil {
|
||||
t.Fatalf("write dependency artifact: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(coffLoaderDir, aiExtensionManifestFileName), []byte(`{
|
||||
"name":"coff-loader",
|
||||
"package_name":"coff-loader",
|
||||
"version":"1.0.0",
|
||||
"extension_author":"sliver",
|
||||
"original_author":"sliver",
|
||||
"repo_url":"https://example.test/coff-loader",
|
||||
"commands":[
|
||||
{
|
||||
"command_name":"coff-loader",
|
||||
"help":"Load and run COFFs",
|
||||
"entrypoint":"LoadAndRun",
|
||||
"files":[{"os":"windows","arch":"amd64","path":"coff-loader.x64.dll"}]
|
||||
}
|
||||
]
|
||||
}`), 0o600); err != nil {
|
||||
t.Fatalf("write dependency manifest: %v", err)
|
||||
}
|
||||
|
||||
nanodumpDir := filepath.Join(serverassets.GetAIExtensionsDir(), "nanodump")
|
||||
if err := os.MkdirAll(nanodumpDir, 0o700); err != nil {
|
||||
t.Fatalf("mkdir nanodump dir: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(nanodumpDir, "nanodump.x64.o"), []byte("nanodump-bof"), 0o600); err != nil {
|
||||
t.Fatalf("write bof artifact: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(nanodumpDir, aiExtensionManifestFileName), []byte(`{
|
||||
"name":"nanodump",
|
||||
"package_name":"nanodump",
|
||||
"version":"1.0.0",
|
||||
"extension_author":"sliver",
|
||||
"original_author":"sliver",
|
||||
"repo_url":"https://example.test/nanodump",
|
||||
"commands":[
|
||||
{
|
||||
"command_name":"nanodump",
|
||||
"help":"Dump LSASS",
|
||||
"entrypoint":"go",
|
||||
"depends_on":"coff-loader",
|
||||
"files":[{"os":"windows","arch":"amd64","path":"nanodump.x64.o"}],
|
||||
"arguments":[
|
||||
{"name":"pid","type":"int","desc":"PID to dump","optional":false}
|
||||
]
|
||||
}
|
||||
]
|
||||
}`), 0o600); err != nil {
|
||||
t.Fatalf("write bof manifest: %v", err)
|
||||
}
|
||||
|
||||
backend := &fakePackageBackend{
|
||||
sessions: &clientpb.Sessions{
|
||||
Sessions: []*clientpb.Session{
|
||||
{ID: "session-1", OS: "windows", Arch: "amd64", Hostname: "winhost"},
|
||||
},
|
||||
},
|
||||
listExtensionsFn: func(_ context.Context, _ *sliverpb.ListExtensionsReq) (*sliverpb.ListExtensions, error) {
|
||||
return &sliverpb.ListExtensions{
|
||||
Names: []string{},
|
||||
Response: &commonpb.Response{},
|
||||
}, nil
|
||||
},
|
||||
registerExtensionFn: func(_ context.Context, _ *sliverpb.RegisterExtensionReq) (*sliverpb.RegisterExtension, error) {
|
||||
return &sliverpb.RegisterExtension{Response: &commonpb.Response{}}, nil
|
||||
},
|
||||
callExtensionFn: func(_ context.Context, _ *sliverpb.CallExtensionReq) (*sliverpb.CallExtension, error) {
|
||||
return &sliverpb.CallExtension{
|
||||
Output: []byte("extension-output"),
|
||||
Response: &commonpb.Response{},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
executor := &executor{
|
||||
backend: backend,
|
||||
conversation: &clientpb.AIConversation{
|
||||
TargetSessionID: "session-1",
|
||||
},
|
||||
}
|
||||
|
||||
raw, err := executor.callExecuteExtension(context.Background(), executeExtensionArgs{
|
||||
CommandName: "nanodump",
|
||||
NamedArgs: map[string]any{
|
||||
"pid": 1337,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("execute extension: %v", err)
|
||||
}
|
||||
|
||||
if len(backend.listExtensionsReqs) != 1 {
|
||||
t.Fatalf("expected list-extensions request, got %d", len(backend.listExtensionsReqs))
|
||||
}
|
||||
if len(backend.registerExtensionReqs) != 1 {
|
||||
t.Fatalf("expected one dependency registration, got %d", len(backend.registerExtensionReqs))
|
||||
}
|
||||
registerReq := backend.registerExtensionReqs[0]
|
||||
if registerReq.GetOS() != "windows" {
|
||||
t.Fatalf("unexpected dependency registration target os: %+v", registerReq)
|
||||
}
|
||||
if string(registerReq.GetData()) != "coff-loader-binary" {
|
||||
t.Fatalf("expected dependency bytes to be registered, got %q", string(registerReq.GetData()))
|
||||
}
|
||||
|
||||
if len(backend.callExtensionReqs) != 1 {
|
||||
t.Fatalf("expected one call-extension request, got %d", len(backend.callExtensionReqs))
|
||||
}
|
||||
callReq := backend.callExtensionReqs[0]
|
||||
if callReq.GetName() != registerReq.GetName() {
|
||||
t.Fatalf("expected BOF call to use dependency hash, got register=%q call=%q", registerReq.GetName(), callReq.GetName())
|
||||
}
|
||||
if callReq.GetExport() != "LoadAndRun" {
|
||||
t.Fatalf("unexpected BOF dependency export: %+v", callReq)
|
||||
}
|
||||
if len(callReq.GetArgs()) == 0 {
|
||||
t.Fatalf("expected packed BOF arguments, got empty buffer")
|
||||
}
|
||||
|
||||
var resp extensionExecutionResult
|
||||
if err := json.Unmarshal([]byte(raw), &resp); err != nil {
|
||||
t.Fatalf("unmarshal extension execution result: %v", err)
|
||||
}
|
||||
if resp.OutputText != "extension-output" {
|
||||
t.Fatalf("unexpected extension output: %+v", resp)
|
||||
}
|
||||
if resp.ExecutionMode != "bof" {
|
||||
t.Fatalf("expected bof execution mode, got %+v", resp)
|
||||
}
|
||||
if resp.DependencyRootPath == "" || resp.DependencyArtifactPath == "" {
|
||||
t.Fatalf("expected dependency metadata in response, got %+v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
type fakePackageBackend struct {
|
||||
sessions *clientpb.Sessions
|
||||
beacons *clientpb.Beacons
|
||||
|
||||
executeAssemblyFn func(context.Context, *sliverpb.ExecuteAssemblyReq) (*sliverpb.ExecuteAssembly, error)
|
||||
listExtensionsFn func(context.Context, *sliverpb.ListExtensionsReq) (*sliverpb.ListExtensions, error)
|
||||
registerExtensionFn func(context.Context, *sliverpb.RegisterExtensionReq) (*sliverpb.RegisterExtension, error)
|
||||
callExtensionFn func(context.Context, *sliverpb.CallExtensionReq) (*sliverpb.CallExtension, error)
|
||||
|
||||
executeAssemblyReqs []*sliverpb.ExecuteAssemblyReq
|
||||
listExtensionsReqs []*sliverpb.ListExtensionsReq
|
||||
registerExtensionReqs []*sliverpb.RegisterExtensionReq
|
||||
callExtensionReqs []*sliverpb.CallExtensionReq
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) GetSessions(context.Context, *commonpb.Empty) (*clientpb.Sessions, error) {
|
||||
if f.sessions == nil {
|
||||
return &clientpb.Sessions{}, nil
|
||||
}
|
||||
return f.sessions, nil
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) GetBeacons(context.Context, *commonpb.Empty) (*clientpb.Beacons, error) {
|
||||
if f.beacons == nil {
|
||||
return &clientpb.Beacons{}, nil
|
||||
}
|
||||
return f.beacons, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Ls(context.Context, *sliverpb.LsReq) (*sliverpb.Ls, error) {
|
||||
return &sliverpb.Ls{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Mv(context.Context, *sliverpb.MvReq) (*sliverpb.Mv, error) {
|
||||
return &sliverpb.Mv{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Cp(context.Context, *sliverpb.CpReq) (*sliverpb.Cp, error) {
|
||||
return &sliverpb.Cp{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Rm(context.Context, *sliverpb.RmReq) (*sliverpb.Rm, error) {
|
||||
return &sliverpb.Rm{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Mkdir(context.Context, *sliverpb.MkdirReq) (*sliverpb.Mkdir, error) {
|
||||
return &sliverpb.Mkdir{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Cd(context.Context, *sliverpb.CdReq) (*sliverpb.Pwd, error) {
|
||||
return &sliverpb.Pwd{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Download(context.Context, *sliverpb.DownloadReq) (*sliverpb.Download, error) {
|
||||
return &sliverpb.Download{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Pwd(context.Context, *sliverpb.PwdReq) (*sliverpb.Pwd, error) {
|
||||
return &sliverpb.Pwd{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Chmod(context.Context, *sliverpb.ChmodReq) (*sliverpb.Chmod, error) {
|
||||
return &sliverpb.Chmod{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Chown(context.Context, *sliverpb.ChownReq) (*sliverpb.Chown, error) {
|
||||
return &sliverpb.Chown{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Chtimes(context.Context, *sliverpb.ChtimesReq) (*sliverpb.Chtimes, error) {
|
||||
return &sliverpb.Chtimes{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Mount(context.Context, *sliverpb.MountReq) (*sliverpb.Mount, error) {
|
||||
return &sliverpb.Mount{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Ifconfig(context.Context, *sliverpb.IfconfigReq) (*sliverpb.Ifconfig, error) {
|
||||
return &sliverpb.Ifconfig{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Netstat(context.Context, *sliverpb.NetstatReq) (*sliverpb.Netstat, error) {
|
||||
return &sliverpb.Netstat{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Ps(context.Context, *sliverpb.PsReq) (*sliverpb.Ps, error) {
|
||||
return &sliverpb.Ps{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) GetEnv(context.Context, *sliverpb.EnvReq) (*sliverpb.EnvInfo, error) {
|
||||
return &sliverpb.EnvInfo{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Ping(context.Context, *sliverpb.Ping) (*sliverpb.Ping, error) {
|
||||
return &sliverpb.Ping{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Screenshot(context.Context, *sliverpb.ScreenshotReq) (*sliverpb.Screenshot, error) {
|
||||
return &sliverpb.Screenshot{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Execute(context.Context, *sliverpb.ExecuteReq) (*sliverpb.Execute, error) {
|
||||
return &sliverpb.Execute{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) ExecuteWindows(context.Context, *sliverpb.ExecuteWindowsReq) (*sliverpb.Execute, error) {
|
||||
return &sliverpb.Execute{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) ExecuteAssembly(ctx context.Context, req *sliverpb.ExecuteAssemblyReq) (*sliverpb.ExecuteAssembly, error) {
|
||||
f.executeAssemblyReqs = append(f.executeAssemblyReqs, req)
|
||||
if f.executeAssemblyFn != nil {
|
||||
return f.executeAssemblyFn(ctx, req)
|
||||
}
|
||||
return &sliverpb.ExecuteAssembly{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) Sideload(context.Context, *sliverpb.SideloadReq) (*sliverpb.Sideload, error) {
|
||||
return &sliverpb.Sideload{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (*fakePackageBackend) SpawnDll(context.Context, *sliverpb.InvokeSpawnDllReq) (*sliverpb.SpawnDll, error) {
|
||||
return &sliverpb.SpawnDll{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) RegisterExtension(ctx context.Context, req *sliverpb.RegisterExtensionReq) (*sliverpb.RegisterExtension, error) {
|
||||
f.registerExtensionReqs = append(f.registerExtensionReqs, req)
|
||||
if f.registerExtensionFn != nil {
|
||||
return f.registerExtensionFn(ctx, req)
|
||||
}
|
||||
return &sliverpb.RegisterExtension{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) ListExtensions(ctx context.Context, req *sliverpb.ListExtensionsReq) (*sliverpb.ListExtensions, error) {
|
||||
f.listExtensionsReqs = append(f.listExtensionsReqs, req)
|
||||
if f.listExtensionsFn != nil {
|
||||
return f.listExtensionsFn(ctx, req)
|
||||
}
|
||||
return &sliverpb.ListExtensions{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
|
||||
func (f *fakePackageBackend) CallExtension(ctx context.Context, req *sliverpb.CallExtensionReq) (*sliverpb.CallExtension, error) {
|
||||
f.callExtensionReqs = append(f.callExtensionReqs, req)
|
||||
if f.callExtensionFn != nil {
|
||||
return f.callExtensionFn(ctx, req)
|
||||
}
|
||||
return &sliverpb.CallExtension{Response: &commonpb.Response{}}, nil
|
||||
}
|
||||
@@ -0,0 +1,537 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/packages"
|
||||
serverassets "github.com/bishopfox/sliver/server/assets"
|
||||
"github.com/bishopfox/sliver/util"
|
||||
)
|
||||
|
||||
const (
|
||||
aiAliasManifestFileName = "alias.json"
|
||||
aiExtensionManifestFileName = "extension.json"
|
||||
)
|
||||
|
||||
type aiAliasFile struct {
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
type aiAliasArgument struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Desc string `json:"desc"`
|
||||
Optional bool `json:"optional"`
|
||||
Default any `json:"default,omitempty"`
|
||||
Choices []string `json:"choices,omitempty"`
|
||||
}
|
||||
|
||||
type aiAliasManifest struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
CommandName string `json:"command_name"`
|
||||
OriginalAuthor string `json:"original_author"`
|
||||
RepoURL string `json:"repo_url"`
|
||||
Help string `json:"help"`
|
||||
LongHelp string `json:"long_help"`
|
||||
Entrypoint string `json:"entrypoint"`
|
||||
AllowArgs bool `json:"allow_args"`
|
||||
DefaultArgs string `json:"default_args"`
|
||||
IsReflective bool `json:"is_reflective"`
|
||||
IsAssembly bool `json:"is_assembly"`
|
||||
|
||||
Arguments []*aiAliasArgument `json:"arguments"`
|
||||
Files []*aiAliasFile `json:"files"`
|
||||
Schema *packages.OutputSchema `json:"schema"`
|
||||
|
||||
RootPath string `json:"-"`
|
||||
}
|
||||
|
||||
type aiExtensionFile struct {
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
type aiExtensionArgument struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Desc string `json:"desc"`
|
||||
Optional bool `json:"optional"`
|
||||
Default any `json:"default,omitempty"`
|
||||
Choices []string `json:"choices,omitempty"`
|
||||
}
|
||||
|
||||
type aiExtensionCommand struct {
|
||||
CommandName string `json:"command_name"`
|
||||
Help string `json:"help"`
|
||||
LongHelp string `json:"long_help"`
|
||||
Entrypoint string `json:"entrypoint"`
|
||||
DependsOn string `json:"depends_on"`
|
||||
Init string `json:"init"`
|
||||
Files []*aiExtensionFile `json:"files"`
|
||||
Arguments []*aiExtensionArgument `json:"arguments"`
|
||||
Schema *packages.OutputSchema `json:"schema"`
|
||||
|
||||
Manifest *aiExtensionManifest `json:"-"`
|
||||
}
|
||||
|
||||
type aiExtensionManifest struct {
|
||||
Name string `json:"name"`
|
||||
PackageName string `json:"package_name"`
|
||||
Version string `json:"version"`
|
||||
ExtensionAuthor string `json:"extension_author"`
|
||||
OriginalAuthor string `json:"original_author"`
|
||||
RepoURL string `json:"repo_url"`
|
||||
|
||||
ExtCommand []*aiExtensionCommand `json:"commands"`
|
||||
|
||||
RootPath string `json:"-"`
|
||||
}
|
||||
|
||||
type aiLegacyExtensionManifest struct {
|
||||
Name string `json:"name"`
|
||||
CommandName string `json:"command_name"`
|
||||
Version string `json:"version"`
|
||||
ExtensionAuthor string `json:"extension_author"`
|
||||
OriginalAuthor string `json:"original_author"`
|
||||
RepoURL string `json:"repo_url"`
|
||||
Help string `json:"help"`
|
||||
LongHelp string `json:"long_help"`
|
||||
Entrypoint string `json:"entrypoint"`
|
||||
DependsOn string `json:"depends_on"`
|
||||
Init string `json:"init"`
|
||||
Files []*aiExtensionFile `json:"files"`
|
||||
Arguments []*aiExtensionArgument `json:"arguments"`
|
||||
Schema *packages.OutputSchema `json:"schema"`
|
||||
|
||||
RootPath string `json:"-"`
|
||||
}
|
||||
|
||||
type extensionCandidate struct {
|
||||
Command *aiExtensionCommand
|
||||
RootPath string
|
||||
}
|
||||
|
||||
func loadAIAliases() ([]*aiAliasManifest, []string) {
|
||||
storeDir := serverassets.GetAIAliasesDir()
|
||||
dirEntries, err := os.ReadDir(storeDir)
|
||||
if err != nil {
|
||||
return nil, []string{fmt.Sprintf("failed to read alias store %s: %s", storeDir, err)}
|
||||
}
|
||||
|
||||
aliases := []*aiAliasManifest{}
|
||||
warnings := []string{}
|
||||
for _, entry := range dirEntries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
manifestPath := filepath.Join(storeDir, entry.Name(), aiAliasManifestFileName)
|
||||
manifest, err := loadAIAliasManifest(manifestPath)
|
||||
if err != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("skipping alias manifest %s: %s", manifestPath, err))
|
||||
continue
|
||||
}
|
||||
aliases = append(aliases, manifest)
|
||||
}
|
||||
|
||||
sort.Slice(aliases, func(i, j int) bool {
|
||||
if aliases[i].CommandName == aliases[j].CommandName {
|
||||
return aliases[i].RootPath < aliases[j].RootPath
|
||||
}
|
||||
return aliases[i].CommandName < aliases[j].CommandName
|
||||
})
|
||||
sort.Strings(warnings)
|
||||
return aliases, warnings
|
||||
}
|
||||
|
||||
func loadAIExtensions() ([]*aiExtensionManifest, []string) {
|
||||
storeDir := serverassets.GetAIExtensionsDir()
|
||||
dirEntries, err := os.ReadDir(storeDir)
|
||||
if err != nil {
|
||||
return nil, []string{fmt.Sprintf("failed to read extension store %s: %s", storeDir, err)}
|
||||
}
|
||||
|
||||
extensions := []*aiExtensionManifest{}
|
||||
warnings := []string{}
|
||||
for _, entry := range dirEntries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
manifestPath := filepath.Join(storeDir, entry.Name(), aiExtensionManifestFileName)
|
||||
manifest, err := loadAIExtensionManifest(manifestPath)
|
||||
if err != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("skipping extension manifest %s: %s", manifestPath, err))
|
||||
continue
|
||||
}
|
||||
extensions = append(extensions, manifest)
|
||||
}
|
||||
|
||||
sort.Slice(extensions, func(i, j int) bool {
|
||||
if extensionDisplayName(extensions[i]) == extensionDisplayName(extensions[j]) {
|
||||
return extensions[i].RootPath < extensions[j].RootPath
|
||||
}
|
||||
return extensionDisplayName(extensions[i]) < extensionDisplayName(extensions[j])
|
||||
})
|
||||
sort.Strings(warnings)
|
||||
return extensions, warnings
|
||||
}
|
||||
|
||||
func loadAIAliasManifest(manifestPath string) (*aiAliasManifest, error) {
|
||||
data, err := os.ReadFile(manifestPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
manifest, err := parseAIAliasManifest(data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
manifest.RootPath = filepath.Dir(manifestPath)
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func loadAIExtensionManifest(manifestPath string) (*aiExtensionManifest, error) {
|
||||
data, err := os.ReadFile(manifestPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
manifest, err := parseAIExtensionManifest(data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
manifest.RootPath = filepath.Dir(manifestPath)
|
||||
for _, command := range manifest.ExtCommand {
|
||||
command.Manifest = manifest
|
||||
}
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func parseAIAliasManifest(data []byte) (*aiAliasManifest, error) {
|
||||
manifest := &aiAliasManifest{}
|
||||
if err := json.Unmarshal(data, manifest); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(manifest.Name) == "" {
|
||||
return nil, fmt.Errorf("missing alias name in manifest")
|
||||
}
|
||||
if strings.TrimSpace(manifest.CommandName) == "" {
|
||||
return nil, fmt.Errorf("missing command_name in alias manifest")
|
||||
}
|
||||
if strings.TrimSpace(manifest.Help) == "" {
|
||||
return nil, fmt.Errorf("missing help in alias manifest")
|
||||
}
|
||||
|
||||
for _, aliasFile := range manifest.Files {
|
||||
if aliasFile == nil {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(aliasFile.OS) == "" {
|
||||
return nil, fmt.Errorf("missing files.os in alias manifest")
|
||||
}
|
||||
if strings.TrimSpace(aliasFile.Arch) == "" {
|
||||
return nil, fmt.Errorf("missing files.arch in alias manifest")
|
||||
}
|
||||
aliasFile.OS = strings.ToLower(strings.TrimSpace(aliasFile.OS))
|
||||
aliasFile.Arch = strings.ToLower(strings.TrimSpace(aliasFile.Arch))
|
||||
aliasFile.Path = util.ResolvePath(aliasFile.Path)
|
||||
if aliasFile.Path == "" || aliasFile.Path == string(os.PathSeparator) {
|
||||
return nil, fmt.Errorf("missing files.path in alias manifest")
|
||||
}
|
||||
}
|
||||
|
||||
if manifest.Schema != nil {
|
||||
if !packages.IsValidSchemaType(manifest.Schema.Name) {
|
||||
return nil, fmt.Errorf("%s is not a valid schema type", manifest.Schema.Name)
|
||||
}
|
||||
manifest.Schema.IngestColumns()
|
||||
}
|
||||
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func parseAIExtensionManifest(data []byte) (*aiExtensionManifest, error) {
|
||||
manifest := &aiExtensionManifest{}
|
||||
if err := json.Unmarshal(data, manifest); err != nil || len(manifest.ExtCommand) == 0 {
|
||||
legacy := &aiLegacyExtensionManifest{}
|
||||
if legacyErr := json.Unmarshal(data, legacy); legacyErr != nil {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, legacyErr
|
||||
}
|
||||
manifest = convertLegacyAIExtensionManifest(legacy)
|
||||
}
|
||||
|
||||
if strings.TrimSpace(manifest.Name) == "" {
|
||||
return nil, fmt.Errorf("missing name field in extension manifest")
|
||||
}
|
||||
|
||||
for _, command := range manifest.ExtCommand {
|
||||
if command == nil {
|
||||
continue
|
||||
}
|
||||
command.Manifest = manifest
|
||||
if strings.TrimSpace(command.CommandName) == "" {
|
||||
return nil, fmt.Errorf("missing command_name field in extension manifest")
|
||||
}
|
||||
if strings.TrimSpace(command.Help) == "" {
|
||||
return nil, fmt.Errorf("missing help field in extension manifest")
|
||||
}
|
||||
if len(command.Files) == 0 {
|
||||
return nil, fmt.Errorf("missing files field in extension manifest")
|
||||
}
|
||||
for _, extFile := range command.Files {
|
||||
if extFile == nil {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(extFile.OS) == "" {
|
||||
return nil, fmt.Errorf("missing files.os field in extension manifest")
|
||||
}
|
||||
if strings.TrimSpace(extFile.Arch) == "" {
|
||||
return nil, fmt.Errorf("missing files.arch field in extension manifest")
|
||||
}
|
||||
extFile.OS = strings.ToLower(strings.TrimSpace(extFile.OS))
|
||||
extFile.Arch = strings.ToLower(strings.TrimSpace(extFile.Arch))
|
||||
extFile.Path = util.ResolvePath(extFile.Path)
|
||||
if extFile.Path == "" || extFile.Path == string(os.PathSeparator) {
|
||||
return nil, fmt.Errorf("missing files.path field in extension manifest")
|
||||
}
|
||||
}
|
||||
if command.Schema != nil {
|
||||
if !packages.IsValidSchemaType(command.Schema.Name) {
|
||||
return nil, fmt.Errorf("%s is not a valid schema type", command.Schema.Name)
|
||||
}
|
||||
command.Schema.IngestColumns()
|
||||
}
|
||||
}
|
||||
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func convertLegacyAIExtensionManifest(legacy *aiLegacyExtensionManifest) *aiExtensionManifest {
|
||||
return &aiExtensionManifest{
|
||||
Name: legacy.CommandName,
|
||||
Version: legacy.Version,
|
||||
ExtensionAuthor: legacy.ExtensionAuthor,
|
||||
OriginalAuthor: legacy.OriginalAuthor,
|
||||
RepoURL: legacy.RepoURL,
|
||||
RootPath: legacy.RootPath,
|
||||
ExtCommand: []*aiExtensionCommand{
|
||||
{
|
||||
CommandName: legacy.CommandName,
|
||||
Help: legacy.Help,
|
||||
LongHelp: legacy.LongHelp,
|
||||
Entrypoint: legacy.Entrypoint,
|
||||
DependsOn: legacy.DependsOn,
|
||||
Init: legacy.Init,
|
||||
Files: legacy.Files,
|
||||
Arguments: legacy.Arguments,
|
||||
Schema: legacy.Schema,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func selectAIAlias(manifests []*aiAliasManifest, commandName string, rootPath string, targetOS string, targetArch string) (*aiAliasManifest, error) {
|
||||
commandName = normalizePackageCommandName(commandName)
|
||||
if commandName == "" {
|
||||
return nil, fmt.Errorf("command_name is required")
|
||||
}
|
||||
|
||||
candidates := []*aiAliasManifest{}
|
||||
for _, manifest := range manifests {
|
||||
if manifest == nil {
|
||||
continue
|
||||
}
|
||||
if normalizePackageCommandName(manifest.CommandName) != commandName {
|
||||
continue
|
||||
}
|
||||
if rootPath != "" && !samePackageRootPath(manifest.RootPath, rootPath) {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, manifest)
|
||||
}
|
||||
|
||||
switch len(candidates) {
|
||||
case 0:
|
||||
if rootPath != "" {
|
||||
return nil, fmt.Errorf("alias command %q not found at root_path %q", commandName, rootPath)
|
||||
}
|
||||
return nil, fmt.Errorf("alias command %q not found in %s", commandName, serverassets.GetAIAliasesDir())
|
||||
case 1:
|
||||
return candidates[0], nil
|
||||
}
|
||||
|
||||
if targetOS != "" && targetArch != "" {
|
||||
compatible := []*aiAliasManifest{}
|
||||
for _, manifest := range candidates {
|
||||
if _, _, err := manifest.artifactForTarget(targetOS, targetArch); err == nil {
|
||||
compatible = append(compatible, manifest)
|
||||
}
|
||||
}
|
||||
if len(compatible) == 1 {
|
||||
return compatible[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
paths := make([]string, 0, len(candidates))
|
||||
for _, manifest := range candidates {
|
||||
paths = append(paths, manifest.RootPath)
|
||||
}
|
||||
sort.Strings(paths)
|
||||
return nil, fmt.Errorf("alias command %q is ambiguous; specify root_path from search_aliases: %s", commandName, strings.Join(paths, ", "))
|
||||
}
|
||||
|
||||
func selectAIExtensionCommand(manifests []*aiExtensionManifest, commandName string, rootPath string, targetOS string, targetArch string) (*aiExtensionCommand, error) {
|
||||
commandName = normalizePackageCommandName(commandName)
|
||||
if commandName == "" {
|
||||
return nil, fmt.Errorf("command_name is required")
|
||||
}
|
||||
|
||||
candidates := []extensionCandidate{}
|
||||
for _, manifest := range manifests {
|
||||
if manifest == nil {
|
||||
continue
|
||||
}
|
||||
if rootPath != "" && !samePackageRootPath(manifest.RootPath, rootPath) {
|
||||
continue
|
||||
}
|
||||
for _, command := range manifest.ExtCommand {
|
||||
if command == nil {
|
||||
continue
|
||||
}
|
||||
if normalizePackageCommandName(command.CommandName) == commandName {
|
||||
candidates = append(candidates, extensionCandidate{Command: command, RootPath: manifest.RootPath})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
switch len(candidates) {
|
||||
case 0:
|
||||
if rootPath != "" {
|
||||
return nil, fmt.Errorf("extension command %q not found at root_path %q", commandName, rootPath)
|
||||
}
|
||||
return nil, fmt.Errorf("extension command %q not found in %s", commandName, serverassets.GetAIExtensionsDir())
|
||||
case 1:
|
||||
return candidates[0].Command, nil
|
||||
}
|
||||
|
||||
if targetOS != "" && targetArch != "" {
|
||||
compatible := []*aiExtensionCommand{}
|
||||
for _, candidate := range candidates {
|
||||
if _, _, err := candidate.Command.artifactForTarget(targetOS, targetArch); err == nil {
|
||||
compatible = append(compatible, candidate.Command)
|
||||
}
|
||||
}
|
||||
if len(compatible) == 1 {
|
||||
return compatible[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
paths := make([]string, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
paths = append(paths, candidate.RootPath)
|
||||
}
|
||||
sort.Strings(paths)
|
||||
return nil, fmt.Errorf("extension command %q is ambiguous; specify root_path from search_extensions: %s", commandName, strings.Join(paths, ", "))
|
||||
}
|
||||
|
||||
func (manifest *aiAliasManifest) artifactForTarget(targetOS string, targetArch string) (string, *aiAliasFile, error) {
|
||||
for _, aliasFile := range manifest.Files {
|
||||
if aliasFile == nil {
|
||||
continue
|
||||
}
|
||||
if aliasFile.OS != strings.ToLower(strings.TrimSpace(targetOS)) || aliasFile.Arch != strings.ToLower(strings.TrimSpace(targetArch)) {
|
||||
continue
|
||||
}
|
||||
artifactPath := joinPackageArtifactPath(manifest.RootPath, aliasFile.Path)
|
||||
if _, err := os.Stat(artifactPath); err != nil {
|
||||
return "", aliasFile, fmt.Errorf("alias artifact not found: %s", artifactPath)
|
||||
}
|
||||
return artifactPath, aliasFile, nil
|
||||
}
|
||||
return "", nil, fmt.Errorf("no alias artifact found for %s/%s", targetOS, targetArch)
|
||||
}
|
||||
|
||||
func (command *aiExtensionCommand) artifactForTarget(targetOS string, targetArch string) (string, *aiExtensionFile, error) {
|
||||
for _, extFile := range command.Files {
|
||||
if extFile == nil {
|
||||
continue
|
||||
}
|
||||
if extFile.OS != strings.ToLower(strings.TrimSpace(targetOS)) || extFile.Arch != strings.ToLower(strings.TrimSpace(targetArch)) {
|
||||
continue
|
||||
}
|
||||
artifactPath := joinPackageArtifactPath(command.Manifest.RootPath, extFile.Path)
|
||||
if _, err := os.Stat(artifactPath); err != nil {
|
||||
return "", extFile, fmt.Errorf("extension artifact not found: %s", artifactPath)
|
||||
}
|
||||
return artifactPath, extFile, nil
|
||||
}
|
||||
return "", nil, fmt.Errorf("no extension artifact found for %s/%s", targetOS, targetArch)
|
||||
}
|
||||
|
||||
func aliasPlatforms(manifest *aiAliasManifest) []string {
|
||||
platforms := map[string]struct{}{}
|
||||
for _, aliasFile := range manifest.Files {
|
||||
if aliasFile == nil {
|
||||
continue
|
||||
}
|
||||
platforms[fmt.Sprintf("%s/%s", aliasFile.OS, aliasFile.Arch)] = struct{}{}
|
||||
}
|
||||
return sortedMapKeys(platforms)
|
||||
}
|
||||
|
||||
func extensionPlatforms(command *aiExtensionCommand) []string {
|
||||
platforms := map[string]struct{}{}
|
||||
for _, extFile := range command.Files {
|
||||
if extFile == nil {
|
||||
continue
|
||||
}
|
||||
platforms[fmt.Sprintf("%s/%s", extFile.OS, extFile.Arch)] = struct{}{}
|
||||
}
|
||||
return sortedMapKeys(platforms)
|
||||
}
|
||||
|
||||
func extensionDisplayName(manifest *aiExtensionManifest) string {
|
||||
if manifest == nil {
|
||||
return ""
|
||||
}
|
||||
if strings.TrimSpace(manifest.PackageName) != "" {
|
||||
return strings.TrimSpace(manifest.PackageName)
|
||||
}
|
||||
return strings.TrimSpace(manifest.Name)
|
||||
}
|
||||
|
||||
func normalizePackageCommandName(value string) string {
|
||||
return strings.ToLower(strings.TrimSpace(value))
|
||||
}
|
||||
|
||||
func samePackageRootPath(left string, right string) bool {
|
||||
if left == "" || right == "" {
|
||||
return false
|
||||
}
|
||||
return filepath.Clean(left) == filepath.Clean(right)
|
||||
}
|
||||
|
||||
func joinPackageArtifactPath(rootPath string, manifestPath string) string {
|
||||
trimmed := strings.TrimLeft(strings.ReplaceAll(manifestPath, "\\", "/"), "/")
|
||||
return filepath.Join(rootPath, filepath.FromSlash(trimmed))
|
||||
}
|
||||
|
||||
func sortedMapKeys(values map[string]struct{}) []string {
|
||||
keys := make([]string, 0, len(values))
|
||||
for key := range values {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,805 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
clientcredentials "github.com/bishopfox/sliver/client/credentials"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
serverloot "github.com/bishopfox/sliver/server/loot"
|
||||
"github.com/gofrs/uuid"
|
||||
)
|
||||
|
||||
const aiLootReadDefaultMaxBytes = 64 * 1024
|
||||
|
||||
var (
|
||||
hashTypeAliases = map[string]string{
|
||||
"NTLM": "MD4",
|
||||
"SHA224": "SHA2_224",
|
||||
"SHA256": "SHA2_256",
|
||||
"SHA384": "SHA2_384",
|
||||
"SHA512": "SHA2_512",
|
||||
}
|
||||
fileTypeAliases = map[string]string{
|
||||
"BIN": "BINARY",
|
||||
"TXT": "TEXT",
|
||||
}
|
||||
)
|
||||
|
||||
type lootAddToolArgs struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
FileName string `json:"file_name,omitempty"`
|
||||
FileTypeName string `json:"file_type_name,omitempty"`
|
||||
FileTypeValue *int32 `json:"file_type_value,omitempty"`
|
||||
OriginHostUUID string `json:"origin_host_uuid,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
DataBase64 string `json:"data_base64,omitempty"`
|
||||
}
|
||||
|
||||
type lootReadToolArgs struct {
|
||||
LootID string `json:"loot_id,omitempty"`
|
||||
MaxBytes int64 `json:"max_bytes,omitempty"`
|
||||
}
|
||||
|
||||
type lootDeleteToolArgs struct {
|
||||
LootID string `json:"loot_id,omitempty"`
|
||||
}
|
||||
|
||||
type credentialsAddToolArgs struct {
|
||||
Collection string `json:"collection,omitempty"`
|
||||
Username string `json:"username,omitempty"`
|
||||
Plaintext string `json:"plaintext,omitempty"`
|
||||
Hash string `json:"hash,omitempty"`
|
||||
HashTypeName string `json:"hash_type_name,omitempty"`
|
||||
HashTypeValue *int32 `json:"hash_type_value,omitempty"`
|
||||
OriginHostUUID string `json:"origin_host_uuid,omitempty"`
|
||||
}
|
||||
|
||||
type credentialReadToolArgs struct {
|
||||
CredentialID string `json:"credential_id,omitempty"`
|
||||
}
|
||||
|
||||
type credentialDeleteToolArgs struct {
|
||||
CredentialID string `json:"credential_id,omitempty"`
|
||||
}
|
||||
|
||||
type lootSummary struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
FileTypeName string `json:"file_type_name"`
|
||||
FileTypeValue int32 `json:"file_type_value"`
|
||||
OriginHostUUID string `json:"origin_host_uuid,omitempty"`
|
||||
Size int64 `json:"size"`
|
||||
HasFile bool `json:"has_file"`
|
||||
FileName string `json:"file_name,omitempty"`
|
||||
}
|
||||
|
||||
type lootListResult struct {
|
||||
Loot []lootSummary `json:"loot"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type lootReadResult struct {
|
||||
RequestedID string `json:"requested_id"`
|
||||
RequestedMaxBytes int64 `json:"requested_max_bytes"`
|
||||
TotalByteLen int `json:"total_byte_len"`
|
||||
ReturnedByteLen int `json:"returned_byte_len"`
|
||||
Truncated bool `json:"truncated"`
|
||||
SHA256 string `json:"sha256,omitempty"`
|
||||
DataBase64 string `json:"data_base64,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
lootSummary
|
||||
}
|
||||
|
||||
type lootDeleteResult struct {
|
||||
RequestedID string `json:"requested_id"`
|
||||
Deleted lootSummary `json:"deleted"`
|
||||
}
|
||||
|
||||
type credentialSummary struct {
|
||||
ID string `json:"id"`
|
||||
Collection string `json:"collection,omitempty"`
|
||||
Username string `json:"username,omitempty"`
|
||||
OriginHostUUID string `json:"origin_host_uuid,omitempty"`
|
||||
IsCracked bool `json:"is_cracked"`
|
||||
HasPlaintext bool `json:"has_plaintext"`
|
||||
HasHash bool `json:"has_hash"`
|
||||
HashTypeName string `json:"hash_type_name,omitempty"`
|
||||
HashTypeValue *int32 `json:"hash_type_value,omitempty"`
|
||||
}
|
||||
|
||||
type credentialsListResult struct {
|
||||
Credentials []credentialSummary `json:"credentials"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type credentialReadResult struct {
|
||||
Plaintext string `json:"plaintext,omitempty"`
|
||||
Hash string `json:"hash,omitempty"`
|
||||
credentialSummary
|
||||
}
|
||||
|
||||
type credentialDeleteResult struct {
|
||||
RequestedID string `json:"requested_id"`
|
||||
Deleted credentialSummary `json:"deleted"`
|
||||
}
|
||||
|
||||
func storeToolDefinitions() []serverai.AgenticToolDefinition {
|
||||
return []serverai.AgenticToolDefinition{
|
||||
{
|
||||
Name: "loot_list",
|
||||
Description: "List loot entries from the server loot store. Results include metadata only, not file content.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "loot_read",
|
||||
Description: "Read one loot entry from the server loot store and return metadata plus up to max_bytes of file content. Defaults to 65536 bytes when max_bytes is omitted.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"loot_id": map[string]any{"type": "string", "description": "Loot ID or unique ID prefix to read."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Optional maximum number of bytes to return from the loot file content."},
|
||||
},
|
||||
"required": []string{"loot_id"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "loot_add",
|
||||
Description: "Add a new entry to the server loot store from inline text or base64 bytes. Provide either text or data_base64. file_name defaults to name, and file_type defaults to TEXT for text or BINARY for base64 content.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"data_base64": map[string]any{"type": "string", "description": "Binary loot content encoded as base64. Provide either text or data_base64."},
|
||||
"file_name": map[string]any{"type": "string", "description": "Optional filename stored with the loot entry. Defaults to name."},
|
||||
"file_type_name": map[string]any{"type": "string", "description": "Optional file type name such as TEXT or BINARY."},
|
||||
"file_type_value": map[string]any{"type": "integer", "description": "Optional numeric file type enum value."},
|
||||
"name": map[string]any{"type": "string", "description": "Display name for the loot entry."},
|
||||
"origin_host_uuid": map[string]any{"type": "string", "description": "Optional origin host UUID for the loot entry."},
|
||||
"text": map[string]any{"type": "string", "description": "Text loot content. Provide either text or data_base64."},
|
||||
},
|
||||
"required": []string{"name"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "loot_delete",
|
||||
Description: "Delete one loot entry from the server loot store using a full ID or unique prefix.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"loot_id": map[string]any{"type": "string", "description": "Loot ID or unique ID prefix to delete."},
|
||||
},
|
||||
"required": []string{"loot_id"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "credentials_list",
|
||||
Description: "List credentials from the server credentials store. Results include metadata only, not plaintext or hash bodies.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "credentials_read",
|
||||
Description: "Read one credential from the server credentials store using a full ID or unique prefix. Returns the stored plaintext and hash values when present.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"credential_id": map[string]any{"type": "string", "description": "Credential ID or unique ID prefix to read."},
|
||||
},
|
||||
"required": []string{"credential_id"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "credentials_add",
|
||||
Description: "Add one credential to the server credentials store. If hash_type is omitted and hash is provided, the server sniffs the hash type automatically.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"collection": map[string]any{"type": "string", "description": "Optional credential collection or source label."},
|
||||
"hash": map[string]any{"type": "string", "description": "Optional password hash."},
|
||||
"hash_type_name": map[string]any{"type": "string", "description": "Optional hash type name such as MD4, NTLM, or SHA2_256."},
|
||||
"hash_type_value": map[string]any{"type": "integer", "description": "Optional numeric hash type enum value."},
|
||||
"origin_host_uuid": map[string]any{"type": "string", "description": "Optional origin host UUID for the credential."},
|
||||
"plaintext": map[string]any{"type": "string", "description": "Optional plaintext password or secret."},
|
||||
"username": map[string]any{"type": "string", "description": "Optional username or account identifier."},
|
||||
},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "credentials_delete",
|
||||
Description: "Delete one credential from the server credentials store using a full ID or unique prefix.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"credential_id": map[string]any{"type": "string", "description": "Credential ID or unique ID prefix to delete."},
|
||||
},
|
||||
"required": []string{"credential_id"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callStoreTool(ctx context.Context, name string, arguments string) (string, bool, error) {
|
||||
switch strings.TrimSpace(name) {
|
||||
case "loot_list":
|
||||
result, err := e.callLootList(ctx)
|
||||
return result, true, err
|
||||
case "loot_read":
|
||||
var args lootReadToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callLootRead(ctx, args)
|
||||
return result, true, err
|
||||
case "loot_add":
|
||||
var args lootAddToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callLootAdd(ctx, args)
|
||||
return result, true, err
|
||||
case "loot_delete":
|
||||
var args lootDeleteToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callLootDelete(ctx, args)
|
||||
return result, true, err
|
||||
case "credentials_list":
|
||||
result, err := e.callCredentialsList(ctx)
|
||||
return result, true, err
|
||||
case "credentials_read":
|
||||
var args credentialReadToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callCredentialsRead(ctx, args)
|
||||
return result, true, err
|
||||
case "credentials_add":
|
||||
var args credentialsAddToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callCredentialsAdd(ctx, args)
|
||||
return result, true, err
|
||||
case "credentials_delete":
|
||||
var args credentialDeleteToolArgs
|
||||
if err := decodeToolArgs(arguments, &args); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
result, err := e.callCredentialsDelete(ctx, args)
|
||||
return result, true, err
|
||||
default:
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callLootList(_ context.Context) (string, error) {
|
||||
allLoot := serverloot.GetLootStore().All()
|
||||
if allLoot == nil {
|
||||
return "", fmt.Errorf("failed to load loot store")
|
||||
}
|
||||
|
||||
summaries := make([]lootSummary, 0, len(allLoot.GetLoot()))
|
||||
for _, entry := range allLoot.GetLoot() {
|
||||
if entry == nil {
|
||||
continue
|
||||
}
|
||||
summaries = append(summaries, lootSummaryFromProtobuf(entry))
|
||||
}
|
||||
sort.Slice(summaries, func(i, j int) bool {
|
||||
if summaries[i].Name == summaries[j].Name {
|
||||
return summaries[i].ID < summaries[j].ID
|
||||
}
|
||||
return summaries[i].Name < summaries[j].Name
|
||||
})
|
||||
|
||||
return marshalToolResult(lootListResult{
|
||||
Loot: summaries,
|
||||
Count: len(summaries),
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callLootAdd(_ context.Context, args lootAddToolArgs) (string, error) {
|
||||
name := strings.TrimSpace(args.Name)
|
||||
if name == "" {
|
||||
return "", fmt.Errorf("name is required")
|
||||
}
|
||||
|
||||
data, err := decodeLootContent(args.Text, args.DataBase64)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
fileType, err := resolveLootFileType(args.FileTypeName, args.FileTypeValue, args.Text, args.DataBase64)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
fileName := strings.TrimSpace(args.FileName)
|
||||
if fileName == "" {
|
||||
fileName = name
|
||||
}
|
||||
|
||||
entry, err := serverloot.GetLootStore().Add(&clientpb.Loot{
|
||||
Name: name,
|
||||
FileType: fileType,
|
||||
OriginHostUUID: strings.TrimSpace(args.OriginHostUUID),
|
||||
File: &commonpb.File{
|
||||
Name: fileName,
|
||||
Data: data,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(lootSummaryFromProtobuf(entry))
|
||||
}
|
||||
|
||||
func (e *executor) callLootRead(_ context.Context, args lootReadToolArgs) (string, error) {
|
||||
record, err := resolveLootRecord(args.LootID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
entry, err := serverloot.GetLootStore().GetContent(record.ID.String(), true)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
maxBytes := args.MaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = aiLootReadDefaultMaxBytes
|
||||
}
|
||||
|
||||
data := []byte(nil)
|
||||
if entry.GetFile() != nil {
|
||||
data = entry.GetFile().GetData()
|
||||
}
|
||||
totalLen := len(data)
|
||||
returned := data
|
||||
if int64(len(returned)) > maxBytes {
|
||||
returned = returned[:maxBytes]
|
||||
}
|
||||
|
||||
result := lootReadResult{
|
||||
RequestedID: strings.TrimSpace(args.LootID),
|
||||
RequestedMaxBytes: maxBytes,
|
||||
TotalByteLen: totalLen,
|
||||
ReturnedByteLen: len(returned),
|
||||
Truncated: len(returned) < totalLen,
|
||||
lootSummary: lootSummaryFromProtobuf(entry),
|
||||
}
|
||||
if totalLen > 0 {
|
||||
result.SHA256 = sha256Hex(data)
|
||||
result.DataBase64 = base64.StdEncoding.EncodeToString(returned)
|
||||
if len(returned) > 0 {
|
||||
result.Text, _ = bytesToTextAndBase64(returned)
|
||||
}
|
||||
}
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callLootDelete(_ context.Context, args lootDeleteToolArgs) (string, error) {
|
||||
record, err := resolveLootRecord(args.LootID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
entry, err := serverloot.GetLootStore().GetContent(record.ID.String(), false)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := serverloot.GetLootStore().Rm(record.ID.String()); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return marshalToolResult(lootDeleteResult{
|
||||
RequestedID: strings.TrimSpace(args.LootID),
|
||||
Deleted: lootSummaryFromProtobuf(entry),
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callCredentialsList(_ context.Context) (string, error) {
|
||||
dbCreds := []*models.Credential{}
|
||||
if err := db.Session().Where(&models.Credential{}).Find(&dbCreds).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
summaries := make([]credentialSummary, 0, len(dbCreds))
|
||||
for _, cred := range dbCreds {
|
||||
if cred == nil {
|
||||
continue
|
||||
}
|
||||
summaries = append(summaries, credentialSummaryFromProtobuf(cred.ToProtobuf()))
|
||||
}
|
||||
sort.Slice(summaries, func(i, j int) bool {
|
||||
if summaries[i].Collection == summaries[j].Collection {
|
||||
if summaries[i].Username == summaries[j].Username {
|
||||
return summaries[i].ID < summaries[j].ID
|
||||
}
|
||||
return summaries[i].Username < summaries[j].Username
|
||||
}
|
||||
return summaries[i].Collection < summaries[j].Collection
|
||||
})
|
||||
|
||||
return marshalToolResult(credentialsListResult{
|
||||
Credentials: summaries,
|
||||
Count: len(summaries),
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callCredentialsAdd(_ context.Context, args credentialsAddToolArgs) (string, error) {
|
||||
if strings.TrimSpace(args.Username) == "" && args.Plaintext == "" && args.Hash == "" {
|
||||
return "", fmt.Errorf("at least one of username, plaintext, or hash is required")
|
||||
}
|
||||
|
||||
hashType, err := resolveHashType(args.HashTypeName, args.HashTypeValue, args.Hash)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
record := &models.Credential{
|
||||
Collection: strings.TrimSpace(args.Collection),
|
||||
Username: strings.TrimSpace(args.Username),
|
||||
Plaintext: args.Plaintext,
|
||||
Hash: args.Hash,
|
||||
HashType: int32(hashType),
|
||||
IsCracked: args.Plaintext != "" && args.Hash != "",
|
||||
OriginHostUUID: uuid.FromStringOrNil(strings.TrimSpace(args.OriginHostUUID)),
|
||||
}
|
||||
if err := db.Session().Create(record).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return marshalToolResult(credentialReadResultFromProtobuf(record.ToProtobuf()))
|
||||
}
|
||||
|
||||
func (e *executor) callCredentialsRead(_ context.Context, args credentialReadToolArgs) (string, error) {
|
||||
record, err := resolveCredentialRecord(args.CredentialID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(credentialReadResultFromProtobuf(record.ToProtobuf()))
|
||||
}
|
||||
|
||||
func (e *executor) callCredentialsDelete(_ context.Context, args credentialDeleteToolArgs) (string, error) {
|
||||
record, err := resolveCredentialRecord(args.CredentialID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
summary := credentialSummaryFromProtobuf(record.ToProtobuf())
|
||||
if err := db.Session().Delete(record).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return marshalToolResult(credentialDeleteResult{
|
||||
RequestedID: strings.TrimSpace(args.CredentialID),
|
||||
Deleted: summary,
|
||||
})
|
||||
}
|
||||
|
||||
func decodeLootContent(text string, dataBase64 string) ([]byte, error) {
|
||||
hasText := text != ""
|
||||
hasData := strings.TrimSpace(dataBase64) != ""
|
||||
switch {
|
||||
case hasText && hasData:
|
||||
return nil, fmt.Errorf("provide only one of text or data_base64")
|
||||
case hasText:
|
||||
return []byte(text), nil
|
||||
case hasData:
|
||||
data, err := base64.StdEncoding.DecodeString(strings.TrimSpace(dataBase64))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid data_base64: %w", err)
|
||||
}
|
||||
return data, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("text or data_base64 is required")
|
||||
}
|
||||
}
|
||||
|
||||
func resolveLootFileType(name string, value *int32, text string, dataBase64 string) (clientpb.FileType, error) {
|
||||
valueType, hasValue, err := parseFileTypeValue(value)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
nameType, hasName, err := parseFileTypeName(name)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if hasValue && hasName {
|
||||
if valueType != nameType {
|
||||
return 0, fmt.Errorf("file_type_name and file_type_value refer to different file types")
|
||||
}
|
||||
return valueType, nil
|
||||
}
|
||||
if hasValue {
|
||||
return valueType, nil
|
||||
}
|
||||
if hasName {
|
||||
return nameType, nil
|
||||
}
|
||||
if text != "" {
|
||||
return clientpb.FileType_TEXT, nil
|
||||
}
|
||||
if strings.TrimSpace(dataBase64) != "" {
|
||||
return clientpb.FileType_BINARY, nil
|
||||
}
|
||||
return 0, fmt.Errorf("file content is required")
|
||||
}
|
||||
|
||||
func resolveHashType(name string, value *int32, hash string) (clientpb.HashType, error) {
|
||||
valueType, hasValue, err := parseHashTypeValue(value)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
nameType, hasName, err := parseHashTypeName(name)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if hasValue && hasName {
|
||||
if valueType != nameType {
|
||||
return 0, fmt.Errorf("hash_type_name and hash_type_value refer to different hash types")
|
||||
}
|
||||
return valueType, nil
|
||||
}
|
||||
if hasValue {
|
||||
return valueType, nil
|
||||
}
|
||||
if hasName {
|
||||
return nameType, nil
|
||||
}
|
||||
if hash != "" {
|
||||
return clientcredentials.SniffHashType(hash), nil
|
||||
}
|
||||
return clientpb.HashType_MD5, nil
|
||||
}
|
||||
|
||||
func parseFileTypeName(name string) (clientpb.FileType, bool, error) {
|
||||
normalized := normalizeEnumName(name)
|
||||
if normalized == "" {
|
||||
return 0, false, nil
|
||||
}
|
||||
if alias, ok := fileTypeAliases[normalized]; ok {
|
||||
normalized = alias
|
||||
}
|
||||
value, ok := clientpb.FileType_value[normalized]
|
||||
if !ok {
|
||||
return 0, false, fmt.Errorf("unsupported file_type_name %q", name)
|
||||
}
|
||||
return clientpb.FileType(value), true, nil
|
||||
}
|
||||
|
||||
func parseFileTypeValue(value *int32) (clientpb.FileType, bool, error) {
|
||||
if value == nil {
|
||||
return 0, false, nil
|
||||
}
|
||||
if _, ok := clientpb.FileType_name[*value]; !ok {
|
||||
return 0, false, fmt.Errorf("unsupported file_type_value %d", *value)
|
||||
}
|
||||
return clientpb.FileType(*value), true, nil
|
||||
}
|
||||
|
||||
func parseHashTypeName(name string) (clientpb.HashType, bool, error) {
|
||||
normalized := normalizeEnumName(name)
|
||||
if normalized == "" {
|
||||
return 0, false, nil
|
||||
}
|
||||
if alias, ok := hashTypeAliases[normalized]; ok {
|
||||
normalized = alias
|
||||
}
|
||||
value, ok := clientpb.HashType_value[normalized]
|
||||
if !ok {
|
||||
return 0, false, fmt.Errorf("unsupported hash_type_name %q", name)
|
||||
}
|
||||
return clientpb.HashType(value), true, nil
|
||||
}
|
||||
|
||||
func parseHashTypeValue(value *int32) (clientpb.HashType, bool, error) {
|
||||
if value == nil {
|
||||
return 0, false, nil
|
||||
}
|
||||
if _, ok := clientpb.HashType_name[*value]; !ok {
|
||||
return 0, false, fmt.Errorf("unsupported hash_type_value %d", *value)
|
||||
}
|
||||
return clientpb.HashType(*value), true, nil
|
||||
}
|
||||
|
||||
func normalizeEnumName(name string) string {
|
||||
name = strings.ToUpper(strings.TrimSpace(name))
|
||||
if name == "" {
|
||||
return ""
|
||||
}
|
||||
replacer := strings.NewReplacer("-", "_", " ", "_")
|
||||
return replacer.Replace(name)
|
||||
}
|
||||
|
||||
func resolveLootRecord(id string) (*models.Loot, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return nil, fmt.Errorf("loot_id is required")
|
||||
}
|
||||
|
||||
if lootID := uuid.FromStringOrNil(id); lootID != uuid.Nil {
|
||||
record := &models.Loot{}
|
||||
err := db.Session().Where(&models.Loot{ID: lootID}).First(record).Error
|
||||
if err == nil {
|
||||
return record, nil
|
||||
}
|
||||
if !errors.Is(err, db.ErrRecordNotFound) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
records := []*models.Loot{}
|
||||
if err := db.Session().Where(&models.Loot{}).Find(&records).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
matches := make([]*models.Loot, 0, 1)
|
||||
for _, record := range records {
|
||||
if record != nil && strings.HasPrefix(record.ID.String(), id) {
|
||||
matches = append(matches, record)
|
||||
}
|
||||
}
|
||||
switch len(matches) {
|
||||
case 0:
|
||||
return nil, fmt.Errorf("loot %q not found", id)
|
||||
case 1:
|
||||
return matches[0], nil
|
||||
default:
|
||||
return nil, fmt.Errorf("loot ID prefix %q is ambiguous: %s", id, joinRecordIDsLoot(matches))
|
||||
}
|
||||
}
|
||||
|
||||
func resolveCredentialRecord(id string) (*models.Credential, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return nil, fmt.Errorf("credential_id is required")
|
||||
}
|
||||
|
||||
if credentialID := uuid.FromStringOrNil(id); credentialID != uuid.Nil {
|
||||
record := &models.Credential{}
|
||||
err := db.Session().Where(&models.Credential{ID: credentialID}).First(record).Error
|
||||
if err == nil {
|
||||
return record, nil
|
||||
}
|
||||
if !errors.Is(err, db.ErrRecordNotFound) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
records := []*models.Credential{}
|
||||
if err := db.Session().Where(&models.Credential{}).Find(&records).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
matches := make([]*models.Credential, 0, 1)
|
||||
for _, record := range records {
|
||||
if record != nil && strings.HasPrefix(record.ID.String(), id) {
|
||||
matches = append(matches, record)
|
||||
}
|
||||
}
|
||||
switch len(matches) {
|
||||
case 0:
|
||||
return nil, fmt.Errorf("credential %q not found", id)
|
||||
case 1:
|
||||
return matches[0], nil
|
||||
default:
|
||||
return nil, fmt.Errorf("credential ID prefix %q is ambiguous: %s", id, joinRecordIDsCredential(matches))
|
||||
}
|
||||
}
|
||||
|
||||
func joinRecordIDsLoot(records []*models.Loot) string {
|
||||
ids := make([]string, 0, len(records))
|
||||
for _, record := range records {
|
||||
if record == nil || record.ID == uuid.Nil {
|
||||
continue
|
||||
}
|
||||
ids = append(ids, record.ID.String())
|
||||
}
|
||||
return joinRecordIDs(ids)
|
||||
}
|
||||
|
||||
func joinRecordIDsCredential(records []*models.Credential) string {
|
||||
ids := make([]string, 0, len(records))
|
||||
for _, record := range records {
|
||||
if record == nil || record.ID == uuid.Nil {
|
||||
continue
|
||||
}
|
||||
ids = append(ids, record.ID.String())
|
||||
}
|
||||
return joinRecordIDs(ids)
|
||||
}
|
||||
|
||||
func joinRecordIDs(ids []string) string {
|
||||
sort.Strings(ids)
|
||||
if len(ids) > 5 {
|
||||
ids = append(ids[:5], "...")
|
||||
}
|
||||
return strings.Join(ids, ", ")
|
||||
}
|
||||
|
||||
func lootSummaryFromProtobuf(entry *clientpb.Loot) lootSummary {
|
||||
summary := lootSummary{
|
||||
ID: strings.TrimSpace(entry.GetID()),
|
||||
Name: strings.TrimSpace(entry.GetName()),
|
||||
FileTypeName: fileTypeName(entry.GetFileType()),
|
||||
FileTypeValue: int32(entry.GetFileType()),
|
||||
OriginHostUUID: cleanUUIDString(entry.GetOriginHostUUID()),
|
||||
Size: entry.GetSize(),
|
||||
HasFile: entry.GetFile() != nil,
|
||||
}
|
||||
if entry.GetFile() != nil {
|
||||
summary.FileName = strings.TrimSpace(entry.GetFile().GetName())
|
||||
}
|
||||
return summary
|
||||
}
|
||||
|
||||
func credentialSummaryFromProtobuf(cred *clientpb.Credential) credentialSummary {
|
||||
summary := credentialSummary{
|
||||
ID: strings.TrimSpace(cred.GetID()),
|
||||
Collection: strings.TrimSpace(cred.GetCollection()),
|
||||
Username: strings.TrimSpace(cred.GetUsername()),
|
||||
OriginHostUUID: cleanUUIDString(cred.GetOriginHostUUID()),
|
||||
IsCracked: cred.GetIsCracked(),
|
||||
HasPlaintext: cred.GetPlaintext() != "",
|
||||
HasHash: cred.GetHash() != "",
|
||||
}
|
||||
if summary.HasHash {
|
||||
value := int32(cred.GetHashType())
|
||||
summary.HashTypeName = hashTypeName(cred.GetHashType())
|
||||
summary.HashTypeValue = &value
|
||||
}
|
||||
return summary
|
||||
}
|
||||
|
||||
func credentialReadResultFromProtobuf(cred *clientpb.Credential) credentialReadResult {
|
||||
return credentialReadResult{
|
||||
Plaintext: cred.GetPlaintext(),
|
||||
Hash: cred.GetHash(),
|
||||
credentialSummary: credentialSummaryFromProtobuf(cred),
|
||||
}
|
||||
}
|
||||
|
||||
func fileTypeName(fileType clientpb.FileType) string {
|
||||
if name, ok := clientpb.FileType_name[int32(fileType)]; ok {
|
||||
return name
|
||||
}
|
||||
return fmt.Sprintf("%d", fileType)
|
||||
}
|
||||
|
||||
func hashTypeName(hashType clientpb.HashType) string {
|
||||
if name, ok := clientpb.HashType_name[int32(hashType)]; ok {
|
||||
return name
|
||||
}
|
||||
return fmt.Sprintf("%d", hashType)
|
||||
}
|
||||
|
||||
func cleanUUIDString(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == uuid.Nil.String() {
|
||||
return ""
|
||||
}
|
||||
return raw
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestCredentialsStoreToolsRoundTrip(t *testing.T) {
|
||||
setupStoreToolsTestEnv(t)
|
||||
|
||||
exec := &executor{}
|
||||
ctx := context.Background()
|
||||
|
||||
addedRaw, err := exec.CallTool(ctx, "credentials_add", `{
|
||||
"collection":"corp",
|
||||
"username":"alice",
|
||||
"plaintext":"Password123!",
|
||||
"hash":"8846f7eaee8fb117ad06bdd830b7586c",
|
||||
"hash_type_name":"NTLM"
|
||||
}`)
|
||||
if err != nil {
|
||||
t.Fatalf("add credential: %v", err)
|
||||
}
|
||||
|
||||
var added credentialReadResult
|
||||
if err := json.Unmarshal([]byte(addedRaw), &added); err != nil {
|
||||
t.Fatalf("decode add result: %v", err)
|
||||
}
|
||||
if added.ID == "" {
|
||||
t.Fatal("expected credential ID to be returned")
|
||||
}
|
||||
if added.HashTypeName != "MD4" {
|
||||
t.Fatalf("expected NTLM alias to resolve to MD4, got %q", added.HashTypeName)
|
||||
}
|
||||
if !added.HasPlaintext || !added.HasHash {
|
||||
t.Fatalf("expected stored credential material to be reflected in add result, got %+v", added)
|
||||
}
|
||||
|
||||
listedRaw, err := exec.CallTool(ctx, "credentials_list", `{}`)
|
||||
if err != nil {
|
||||
t.Fatalf("list credentials: %v", err)
|
||||
}
|
||||
|
||||
var listed credentialsListResult
|
||||
if err := json.Unmarshal([]byte(listedRaw), &listed); err != nil {
|
||||
t.Fatalf("decode list result: %v", err)
|
||||
}
|
||||
if listed.Count != 1 || len(listed.Credentials) != 1 {
|
||||
t.Fatalf("unexpected credential count: %+v", listed)
|
||||
}
|
||||
|
||||
prefix := added.ID[:8]
|
||||
readRaw, err := exec.CallTool(ctx, "credentials_read", fmt.Sprintf(`{"credential_id":%q}`, prefix))
|
||||
if err != nil {
|
||||
t.Fatalf("read credential: %v", err)
|
||||
}
|
||||
|
||||
var read credentialReadResult
|
||||
if err := json.Unmarshal([]byte(readRaw), &read); err != nil {
|
||||
t.Fatalf("decode read result: %v", err)
|
||||
}
|
||||
if read.ID != added.ID {
|
||||
t.Fatalf("expected credential prefix lookup to resolve %q, got %+v", added.ID, read)
|
||||
}
|
||||
if read.Plaintext != "Password123!" || read.Hash != "8846f7eaee8fb117ad06bdd830b7586c" {
|
||||
t.Fatalf("unexpected credential body: %+v", read)
|
||||
}
|
||||
|
||||
deletedRaw, err := exec.CallTool(ctx, "credentials_delete", fmt.Sprintf(`{"credential_id":%q}`, prefix))
|
||||
if err != nil {
|
||||
t.Fatalf("delete credential: %v", err)
|
||||
}
|
||||
|
||||
var deleted credentialDeleteResult
|
||||
if err := json.Unmarshal([]byte(deletedRaw), &deleted); err != nil {
|
||||
t.Fatalf("decode delete result: %v", err)
|
||||
}
|
||||
if deleted.Deleted.ID != added.ID {
|
||||
t.Fatalf("unexpected deleted credential: %+v", deleted)
|
||||
}
|
||||
|
||||
listedRaw, err = exec.CallTool(ctx, "credentials_list", `{}`)
|
||||
if err != nil {
|
||||
t.Fatalf("list credentials after delete: %v", err)
|
||||
}
|
||||
if err := json.Unmarshal([]byte(listedRaw), &listed); err != nil {
|
||||
t.Fatalf("decode list-after-delete result: %v", err)
|
||||
}
|
||||
if listed.Count != 0 || len(listed.Credentials) != 0 {
|
||||
t.Fatalf("expected empty credentials store after delete, got %+v", listed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLootStoreToolsRoundTrip(t *testing.T) {
|
||||
setupStoreToolsTestEnv(t)
|
||||
|
||||
exec := &executor{}
|
||||
ctx := context.Background()
|
||||
content := "top secret\nline two"
|
||||
|
||||
addedRaw, err := exec.CallTool(ctx, "loot_add", `{
|
||||
"name":"operator-notes",
|
||||
"file_name":"notes.txt",
|
||||
"text":"top secret\nline two"
|
||||
}`)
|
||||
if err != nil {
|
||||
t.Fatalf("add loot: %v", err)
|
||||
}
|
||||
|
||||
var added lootSummary
|
||||
if err := json.Unmarshal([]byte(addedRaw), &added); err != nil {
|
||||
t.Fatalf("decode add result: %v", err)
|
||||
}
|
||||
if added.ID == "" {
|
||||
t.Fatal("expected loot ID to be returned")
|
||||
}
|
||||
if added.FileTypeName != "TEXT" {
|
||||
t.Fatalf("expected text loot type, got %+v", added)
|
||||
}
|
||||
|
||||
listedRaw, err := exec.CallTool(ctx, "loot_list", `{}`)
|
||||
if err != nil {
|
||||
t.Fatalf("list loot: %v", err)
|
||||
}
|
||||
|
||||
var listed lootListResult
|
||||
if err := json.Unmarshal([]byte(listedRaw), &listed); err != nil {
|
||||
t.Fatalf("decode loot list result: %v", err)
|
||||
}
|
||||
if listed.Count != 1 || len(listed.Loot) != 1 {
|
||||
t.Fatalf("unexpected loot list result: %+v", listed)
|
||||
}
|
||||
|
||||
prefix := added.ID[:8]
|
||||
readRaw, err := exec.CallTool(ctx, "loot_read", fmt.Sprintf(`{"loot_id":%q,"max_bytes":4}`, prefix))
|
||||
if err != nil {
|
||||
t.Fatalf("read loot: %v", err)
|
||||
}
|
||||
|
||||
var read lootReadResult
|
||||
if err := json.Unmarshal([]byte(readRaw), &read); err != nil {
|
||||
t.Fatalf("decode loot read result: %v", err)
|
||||
}
|
||||
if read.ID != added.ID {
|
||||
t.Fatalf("expected loot prefix lookup to resolve %q, got %+v", added.ID, read)
|
||||
}
|
||||
if !read.Truncated || read.ReturnedByteLen != 4 || read.TotalByteLen != len(content) {
|
||||
t.Fatalf("unexpected loot read sizing: %+v", read)
|
||||
}
|
||||
if read.Text != content[:4] {
|
||||
t.Fatalf("unexpected loot text snippet: %+v", read)
|
||||
}
|
||||
decoded, err := base64.StdEncoding.DecodeString(read.DataBase64)
|
||||
if err != nil {
|
||||
t.Fatalf("decode loot data: %v", err)
|
||||
}
|
||||
if string(decoded) != content[:4] {
|
||||
t.Fatalf("unexpected loot data snippet: %q", string(decoded))
|
||||
}
|
||||
|
||||
deletedRaw, err := exec.CallTool(ctx, "loot_delete", fmt.Sprintf(`{"loot_id":%q}`, prefix))
|
||||
if err != nil {
|
||||
t.Fatalf("delete loot: %v", err)
|
||||
}
|
||||
|
||||
var deleted lootDeleteResult
|
||||
if err := json.Unmarshal([]byte(deletedRaw), &deleted); err != nil {
|
||||
t.Fatalf("decode loot delete result: %v", err)
|
||||
}
|
||||
if deleted.Deleted.ID != added.ID {
|
||||
t.Fatalf("unexpected deleted loot: %+v", deleted)
|
||||
}
|
||||
|
||||
listedRaw, err = exec.CallTool(ctx, "loot_list", `{}`)
|
||||
if err != nil {
|
||||
t.Fatalf("list loot after delete: %v", err)
|
||||
}
|
||||
if err := json.Unmarshal([]byte(listedRaw), &listed); err != nil {
|
||||
t.Fatalf("decode loot list-after-delete result: %v", err)
|
||||
}
|
||||
if listed.Count != 0 || len(listed.Loot) != 0 {
|
||||
t.Fatalf("expected empty loot store after delete, got %+v", listed)
|
||||
}
|
||||
}
|
||||
|
||||
func setupStoreToolsTestEnv(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
rootDir := t.TempDir()
|
||||
t.Setenv("SLIVER_ROOT_DIR", rootDir)
|
||||
|
||||
originalDB := db.Client
|
||||
testDB, err := gorm.Open(sqlite.Open(filepath.Join(rootDir, "store-tools-test.db")), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open test db: %v", err)
|
||||
}
|
||||
if err := testDB.AutoMigrate(&models.Credential{}, &models.Loot{}); err != nil {
|
||||
t.Fatalf("migrate test db: %v", err)
|
||||
}
|
||||
db.Client = testDB
|
||||
|
||||
t.Cleanup(func() {
|
||||
sqlDB, err := testDB.DB()
|
||||
if err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
db.Client = originalDB
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
package aitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
)
|
||||
|
||||
type netstatToolArgs struct {
|
||||
targetArgs
|
||||
TCP bool `json:"tcp,omitempty"`
|
||||
UDP bool `json:"udp,omitempty"`
|
||||
IP4 bool `json:"ip4,omitempty"`
|
||||
IP6 bool `json:"ip6,omitempty"`
|
||||
Listening bool `json:"listening,omitempty"`
|
||||
}
|
||||
|
||||
type psToolArgs struct {
|
||||
targetArgs
|
||||
FullInfo bool `json:"full_info,omitempty"`
|
||||
PID int32 `json:"pid,omitempty"`
|
||||
Exe string `json:"exe,omitempty"`
|
||||
Owner string `json:"owner,omitempty"`
|
||||
}
|
||||
|
||||
type envToolArgs struct {
|
||||
targetArgs
|
||||
Name string `json:"name,omitempty"`
|
||||
}
|
||||
|
||||
type execToolArgs struct {
|
||||
targetArgs
|
||||
Path string `json:"path,omitempty"`
|
||||
Args []string `json:"args,omitempty"`
|
||||
CaptureOutput *bool `json:"capture_output,omitempty"`
|
||||
Background bool `json:"background,omitempty"`
|
||||
Stdout string `json:"stdout,omitempty"`
|
||||
Stderr string `json:"stderr,omitempty"`
|
||||
EnvInheritance bool `json:"env_inheritance,omitempty"`
|
||||
Env map[string]string `json:"env,omitempty"`
|
||||
UseToken bool `json:"use_token,omitempty"`
|
||||
HideWindow bool `json:"hide_window,omitempty"`
|
||||
PPID uint32 `json:"ppid,omitempty"`
|
||||
}
|
||||
|
||||
type pingToolArgs struct {
|
||||
targetArgs
|
||||
Nonce *int32 `json:"nonce,omitempty"`
|
||||
}
|
||||
|
||||
type screenshotToolArgs struct {
|
||||
targetArgs
|
||||
IncludeData bool `json:"include_data,omitempty"`
|
||||
}
|
||||
|
||||
type ifconfigInterfaceResult struct {
|
||||
Index int32 `json:"index"`
|
||||
Name string `json:"name"`
|
||||
MAC string `json:"mac,omitempty"`
|
||||
IPAddresses []string `json:"ip_addresses"`
|
||||
}
|
||||
|
||||
type ifconfigResult struct {
|
||||
Interfaces []ifconfigInterfaceResult `json:"interfaces"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type processResult struct {
|
||||
PID int32 `json:"pid"`
|
||||
PPID int32 `json:"ppid"`
|
||||
Executable string `json:"executable"`
|
||||
Owner string `json:"owner,omitempty"`
|
||||
Architecture string `json:"architecture,omitempty"`
|
||||
SessionID int32 `json:"session_id,omitempty"`
|
||||
CmdLine []string `json:"cmd_line,omitempty"`
|
||||
}
|
||||
|
||||
type psResult struct {
|
||||
Processes []processResult `json:"processes"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type sockAddrResult struct {
|
||||
IP string `json:"ip"`
|
||||
Port uint32 `json:"port"`
|
||||
}
|
||||
|
||||
type netstatEntryResult struct {
|
||||
LocalAddr sockAddrResult `json:"local_addr"`
|
||||
RemoteAddr sockAddrResult `json:"remote_addr"`
|
||||
State string `json:"state"`
|
||||
UID uint32 `json:"uid"`
|
||||
Protocol string `json:"protocol"`
|
||||
Process *processResult `json:"process,omitempty"`
|
||||
}
|
||||
|
||||
type netstatResult struct {
|
||||
Entries []netstatEntryResult `json:"entries"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type envVarResult struct {
|
||||
Key string `json:"key"`
|
||||
Value string `json:"value"`
|
||||
}
|
||||
|
||||
type envResult struct {
|
||||
Variables []envVarResult `json:"variables"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type targetInfoResult struct {
|
||||
TargetType string `json:"target_type"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Hostname string `json:"hostname"`
|
||||
UUID string `json:"uuid"`
|
||||
Username string `json:"username"`
|
||||
UID string `json:"uid"`
|
||||
GID string `json:"gid"`
|
||||
PID int32 `json:"pid"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Transport string `json:"transport"`
|
||||
ActiveC2 string `json:"active_c2"`
|
||||
RemoteAddress string `json:"remote_address"`
|
||||
ProxyURL string `json:"proxy_url"`
|
||||
Version string `json:"version"`
|
||||
Locale string `json:"locale"`
|
||||
Integrity string `json:"integrity"`
|
||||
IsDead bool `json:"is_dead"`
|
||||
FirstContact int64 `json:"first_contact"`
|
||||
LastCheckin int64 `json:"last_checkin"`
|
||||
ReconnectInterval int64 `json:"reconnect_interval"`
|
||||
Interval int64 `json:"interval,omitempty"`
|
||||
Jitter int64 `json:"jitter,omitempty"`
|
||||
NextCheckin int64 `json:"next_checkin,omitempty"`
|
||||
TasksCount int64 `json:"tasks_count,omitempty"`
|
||||
TasksCountCompleted int64 `json:"tasks_count_completed,omitempty"`
|
||||
}
|
||||
|
||||
type pidResult struct {
|
||||
PID int32 `json:"pid"`
|
||||
}
|
||||
|
||||
type uidResult struct {
|
||||
UID string `json:"uid"`
|
||||
}
|
||||
|
||||
type pingResult struct {
|
||||
Nonce int32 `json:"nonce"`
|
||||
}
|
||||
|
||||
type screenshotResult struct {
|
||||
ByteLen int `json:"byte_len"`
|
||||
SHA256 string `json:"sha256"`
|
||||
DataBase64 string `json:"data_base64,omitempty"`
|
||||
}
|
||||
|
||||
type execResult struct {
|
||||
PID uint32 `json:"pid"`
|
||||
Status uint32 `json:"status"`
|
||||
StdoutLen int `json:"stdout_len"`
|
||||
StderrLen int `json:"stderr_len"`
|
||||
StdoutText string `json:"stdout_text,omitempty"`
|
||||
StdoutBase64 string `json:"stdout_base64,omitempty"`
|
||||
StderrText string `json:"stderr_text,omitempty"`
|
||||
StderrBase64 string `json:"stderr_base64,omitempty"`
|
||||
}
|
||||
|
||||
func systemToolDefinitions() []serverai.AgenticToolDefinition {
|
||||
return []serverai.AgenticToolDefinition{
|
||||
{
|
||||
Name: "exec_command",
|
||||
Description: "Execute a remote command and optionally capture stdout/stderr. If session_id and beacon_id are both omitted, the conversation target is used when available.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"path": map[string]any{"type": "string", "description": "Executable or command path."},
|
||||
"args": map[string]any{"type": "array", "description": "Optional argument list.", "items": map[string]any{"type": "string"}},
|
||||
"capture_output": map[string]any{"type": "boolean", "description": "Capture stdout/stderr inline. Defaults to true unless background is true."},
|
||||
"background": map[string]any{"type": "boolean", "description": "Run the command in the background."},
|
||||
"stdout": map[string]any{"type": "string", "description": "Optional remote path where stdout should be written."},
|
||||
"stderr": map[string]any{"type": "string", "description": "Optional remote path where stderr should be written."},
|
||||
"env_inheritance": map[string]any{"type": "boolean", "description": "Inherit the target environment."},
|
||||
"env": map[string]any{
|
||||
"type": "object",
|
||||
"description": "Optional environment variables to set for the process.",
|
||||
"additionalProperties": map[string]any{"type": "string"},
|
||||
},
|
||||
"use_token": map[string]any{"type": "boolean", "description": "Windows only: use the current token."},
|
||||
"hide_window": map[string]any{"type": "boolean", "description": "Windows only: hide the window."},
|
||||
"ppid": map[string]any{"type": "integer", "description": "Windows only: spoof the parent PID."},
|
||||
}),
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "ifconfig",
|
||||
Description: "List network interfaces and assigned IP addresses for the target.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": targetSchemaProperties(),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "netstat",
|
||||
Description: "List active network connections on the target.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"tcp": map[string]any{"type": "boolean", "description": "Include TCP sockets."},
|
||||
"udp": map[string]any{"type": "boolean", "description": "Include UDP sockets."},
|
||||
"ip4": map[string]any{"type": "boolean", "description": "Include IPv4 sockets."},
|
||||
"ip6": map[string]any{"type": "boolean", "description": "Include IPv6 sockets."},
|
||||
"listening": map[string]any{"type": "boolean", "description": "Only show listening sockets."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "ps",
|
||||
Description: "List target processes. Set full_info when owner or command line details are needed.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"full_info": map[string]any{"type": "boolean", "description": "Request owner, architecture, session, and command line information."},
|
||||
"pid": map[string]any{"type": "integer", "description": "Optional PID filter applied after retrieval."},
|
||||
"exe": map[string]any{"type": "string", "description": "Optional substring filter on the executable name."},
|
||||
"owner": map[string]any{"type": "string", "description": "Optional substring filter on the process owner."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "env",
|
||||
Description: "Return environment variables for the target, or one variable when name is provided.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"name": map[string]any{"type": "string", "description": "Optional environment variable name to retrieve."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "getpid",
|
||||
Description: "Return the current target process PID from session or beacon metadata.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": targetSchemaProperties(),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "getuid",
|
||||
Description: "Return the current target process UID from session or beacon metadata.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": targetSchemaProperties(),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "info",
|
||||
Description: "Return detailed session or beacon metadata for the selected target.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": targetSchemaProperties(),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "ping",
|
||||
Description: "Send a round-trip C2 ping to the target and return the echoed nonce.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"nonce": map[string]any{"type": "integer", "description": "Optional nonce to send. A random one is generated when omitted."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "screenshot",
|
||||
Description: "Capture a screenshot from the target. Set include_data to true to include the PNG bytes as base64.",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": mergeSchemaProperties(map[string]any{
|
||||
"include_data": map[string]any{"type": "boolean", "description": "Include the raw screenshot bytes as base64 in the result."},
|
||||
}),
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *executor) callIfconfig(ctx context.Context, args targetArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Ifconfig, error) {
|
||||
return e.backend.Ifconfig(callCtx, &sliverpb.IfconfigReq{Request: req})
|
||||
},
|
||||
func() *sliverpb.Ifconfig { return &sliverpb.Ifconfig{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := ifconfigResult{
|
||||
Interfaces: make([]ifconfigInterfaceResult, 0, len(resp.NetInterfaces)),
|
||||
}
|
||||
for _, iface := range resp.NetInterfaces {
|
||||
if iface == nil {
|
||||
continue
|
||||
}
|
||||
result.Interfaces = append(result.Interfaces, ifconfigInterfaceResult{
|
||||
Index: iface.Index,
|
||||
Name: iface.Name,
|
||||
MAC: iface.MAC,
|
||||
IPAddresses: append([]string(nil), iface.IPAddresses...),
|
||||
})
|
||||
}
|
||||
sort.Slice(result.Interfaces, func(i, j int) bool {
|
||||
return result.Interfaces[i].Index < result.Interfaces[j].Index
|
||||
})
|
||||
result.Count = len(result.Interfaces)
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callNetstat(ctx context.Context, args netstatToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Netstat, error) {
|
||||
return e.backend.Netstat(callCtx, &sliverpb.NetstatReq{
|
||||
Request: req,
|
||||
TCP: args.TCP,
|
||||
UDP: args.UDP,
|
||||
IP4: args.IP4,
|
||||
IP6: args.IP6,
|
||||
Listening: args.Listening,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Netstat { return &sliverpb.Netstat{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := netstatResult{
|
||||
Entries: make([]netstatEntryResult, 0, len(resp.Entries)),
|
||||
}
|
||||
for _, entry := range resp.Entries {
|
||||
if entry == nil {
|
||||
continue
|
||||
}
|
||||
resultEntry := netstatEntryResult{
|
||||
State: entry.SkState,
|
||||
UID: entry.UID,
|
||||
Protocol: entry.Protocol,
|
||||
}
|
||||
if entry.LocalAddr != nil {
|
||||
resultEntry.LocalAddr = sockAddrResult{IP: entry.LocalAddr.Ip, Port: entry.LocalAddr.Port}
|
||||
}
|
||||
if entry.RemoteAddr != nil {
|
||||
resultEntry.RemoteAddr = sockAddrResult{IP: entry.RemoteAddr.Ip, Port: entry.RemoteAddr.Port}
|
||||
}
|
||||
if entry.Process != nil {
|
||||
resultEntry.Process = processToResult(entry.Process)
|
||||
}
|
||||
result.Entries = append(result.Entries, resultEntry)
|
||||
}
|
||||
result.Count = len(result.Entries)
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callPS(ctx context.Context, args psToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Ps, error) {
|
||||
return e.backend.Ps(callCtx, &sliverpb.PsReq{
|
||||
Request: req,
|
||||
FullInfo: args.FullInfo,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Ps { return &sliverpb.Ps{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
processes := make([]processResult, 0, len(resp.Processes))
|
||||
for _, process := range resp.Processes {
|
||||
if process == nil {
|
||||
continue
|
||||
}
|
||||
if args.PID != 0 && process.Pid != args.PID {
|
||||
continue
|
||||
}
|
||||
if args.Exe != "" && !strings.Contains(strings.ToLower(process.Executable), strings.ToLower(args.Exe)) {
|
||||
continue
|
||||
}
|
||||
if args.Owner != "" && !strings.Contains(strings.ToLower(process.Owner), strings.ToLower(args.Owner)) {
|
||||
continue
|
||||
}
|
||||
processes = append(processes, *processToResult(process))
|
||||
}
|
||||
sort.Slice(processes, func(i, j int) bool {
|
||||
if processes[i].PID == processes[j].PID {
|
||||
return processes[i].PPID < processes[j].PPID
|
||||
}
|
||||
return processes[i].PID < processes[j].PID
|
||||
})
|
||||
|
||||
return marshalToolResult(psResult{
|
||||
Processes: processes,
|
||||
Count: len(processes),
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callEnv(ctx context.Context, args envToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.EnvInfo, error) {
|
||||
return e.backend.GetEnv(callCtx, &sliverpb.EnvReq{
|
||||
Request: req,
|
||||
Name: args.Name,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.EnvInfo { return &sliverpb.EnvInfo{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := envResult{
|
||||
Variables: make([]envVarResult, 0, len(resp.Variables)),
|
||||
}
|
||||
for _, variable := range resp.Variables {
|
||||
if variable == nil {
|
||||
continue
|
||||
}
|
||||
result.Variables = append(result.Variables, envVarResult{Key: variable.Key, Value: variable.Value})
|
||||
}
|
||||
sort.Slice(result.Variables, func(i, j int) bool {
|
||||
return result.Variables[i].Key < result.Variables[j].Key
|
||||
})
|
||||
result.Count = len(result.Variables)
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callGetPID(ctx context.Context, args targetArgs) (string, error) {
|
||||
session, beacon, err := e.lookupTargetMetadata(ctx, args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if session != nil {
|
||||
return marshalToolResult(pidResult{PID: session.PID})
|
||||
}
|
||||
return marshalToolResult(pidResult{PID: beacon.PID})
|
||||
}
|
||||
|
||||
func (e *executor) callGetUID(ctx context.Context, args targetArgs) (string, error) {
|
||||
session, beacon, err := e.lookupTargetMetadata(ctx, args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if session != nil {
|
||||
return marshalToolResult(uidResult{UID: session.UID})
|
||||
}
|
||||
return marshalToolResult(uidResult{UID: beacon.UID})
|
||||
}
|
||||
|
||||
func (e *executor) callInfo(ctx context.Context, args targetArgs) (string, error) {
|
||||
session, beacon, err := e.lookupTargetMetadata(ctx, args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if session != nil {
|
||||
return marshalToolResult(targetInfoResult{
|
||||
TargetType: "session",
|
||||
ID: session.ID,
|
||||
Name: session.Name,
|
||||
Hostname: session.Hostname,
|
||||
UUID: session.UUID,
|
||||
Username: session.Username,
|
||||
UID: session.UID,
|
||||
GID: session.GID,
|
||||
PID: session.PID,
|
||||
OS: session.OS,
|
||||
Arch: session.Arch,
|
||||
Transport: session.Transport,
|
||||
ActiveC2: session.ActiveC2,
|
||||
RemoteAddress: session.RemoteAddress,
|
||||
ProxyURL: session.ProxyURL,
|
||||
Version: session.Version,
|
||||
Locale: session.Locale,
|
||||
Integrity: session.Integrity,
|
||||
IsDead: session.IsDead,
|
||||
FirstContact: session.FirstContact,
|
||||
LastCheckin: session.LastCheckin,
|
||||
ReconnectInterval: session.ReconnectInterval,
|
||||
})
|
||||
}
|
||||
return marshalToolResult(targetInfoResult{
|
||||
TargetType: "beacon",
|
||||
ID: beacon.ID,
|
||||
Name: beacon.Name,
|
||||
Hostname: beacon.Hostname,
|
||||
UUID: beacon.UUID,
|
||||
Username: beacon.Username,
|
||||
UID: beacon.UID,
|
||||
GID: beacon.GID,
|
||||
PID: beacon.PID,
|
||||
OS: beacon.OS,
|
||||
Arch: beacon.Arch,
|
||||
Transport: beacon.Transport,
|
||||
ActiveC2: beacon.ActiveC2,
|
||||
RemoteAddress: beacon.RemoteAddress,
|
||||
ProxyURL: beacon.ProxyURL,
|
||||
Version: beacon.Version,
|
||||
Locale: beacon.Locale,
|
||||
Integrity: beacon.Integrity,
|
||||
IsDead: beacon.IsDead,
|
||||
FirstContact: beacon.FirstContact,
|
||||
LastCheckin: beacon.LastCheckin,
|
||||
ReconnectInterval: beacon.ReconnectInterval,
|
||||
Interval: beacon.Interval,
|
||||
Jitter: beacon.Jitter,
|
||||
NextCheckin: beacon.NextCheckin,
|
||||
TasksCount: beacon.TasksCount,
|
||||
TasksCountCompleted: beacon.TasksCountCompleted,
|
||||
})
|
||||
}
|
||||
|
||||
func (e *executor) callPing(ctx context.Context, args pingToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
nonce := int32(time.Now().UnixNano() % 1_000_000)
|
||||
if args.Nonce != nil {
|
||||
nonce = *args.Nonce
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Ping, error) {
|
||||
return e.backend.Ping(callCtx, &sliverpb.Ping{
|
||||
Request: req,
|
||||
Nonce: nonce,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Ping { return &sliverpb.Ping{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return marshalToolResult(pingResult{Nonce: resp.Nonce})
|
||||
}
|
||||
|
||||
func (e *executor) callScreenshot(ctx context.Context, args screenshotToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Screenshot, error) {
|
||||
return e.backend.Screenshot(callCtx, &sliverpb.ScreenshotReq{Request: req})
|
||||
},
|
||||
func() *sliverpb.Screenshot { return &sliverpb.Screenshot{} },
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result := screenshotResult{
|
||||
ByteLen: len(resp.Data),
|
||||
SHA256: sha256Hex(resp.Data),
|
||||
}
|
||||
if args.IncludeData {
|
||||
result.DataBase64 = base64.StdEncoding.EncodeToString(resp.Data)
|
||||
}
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
func (e *executor) callExec(ctx context.Context, args execToolArgs) (string, error) {
|
||||
target, err := e.resolveTarget(args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if args.Path == "" {
|
||||
return "", fmt.Errorf("path is required")
|
||||
}
|
||||
|
||||
captureOutput := !args.Background
|
||||
if args.CaptureOutput != nil {
|
||||
captureOutput = *args.CaptureOutput && !args.Background
|
||||
}
|
||||
|
||||
var resp *sliverpb.Execute
|
||||
if args.UseToken || args.HideWindow || args.PPID != 0 {
|
||||
session, beacon, err := e.lookupTargetMetadata(ctx, args.SessionID, args.BeaconID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
targetOS := ""
|
||||
if session != nil {
|
||||
targetOS = session.OS
|
||||
} else if beacon != nil {
|
||||
targetOS = beacon.OS
|
||||
}
|
||||
if targetOS != "windows" {
|
||||
return "", fmt.Errorf("use_token, hide_window, and ppid are only supported on windows targets")
|
||||
}
|
||||
if args.EnvInheritance || len(args.Env) > 0 {
|
||||
return "", fmt.Errorf("env and env_inheritance are not supported with use_token, hide_window, or ppid")
|
||||
}
|
||||
|
||||
resp, err = callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Execute, error) {
|
||||
return e.backend.ExecuteWindows(callCtx, &sliverpb.ExecuteWindowsReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
Args: args.Args,
|
||||
Output: captureOutput,
|
||||
Background: args.Background,
|
||||
Stdout: args.Stdout,
|
||||
Stderr: args.Stderr,
|
||||
UseToken: args.UseToken,
|
||||
HideWindow: args.HideWindow,
|
||||
PPid: args.PPID,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Execute { return &sliverpb.Execute{} },
|
||||
)
|
||||
} else {
|
||||
resp, err = callTargetRPC(
|
||||
ctx,
|
||||
target,
|
||||
func(callCtx context.Context, req *commonpb.Request) (*sliverpb.Execute, error) {
|
||||
return e.backend.Execute(callCtx, &sliverpb.ExecuteReq{
|
||||
Request: req,
|
||||
Path: args.Path,
|
||||
Args: args.Args,
|
||||
Output: captureOutput,
|
||||
Background: args.Background,
|
||||
Stdout: args.Stdout,
|
||||
Stderr: args.Stderr,
|
||||
EnvInheritance: args.EnvInheritance,
|
||||
Env: args.Env,
|
||||
})
|
||||
},
|
||||
func() *sliverpb.Execute { return &sliverpb.Execute{} },
|
||||
)
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
stdoutText, stdoutBase64 := bytesToTextAndBase64(resp.Stdout)
|
||||
stderrText, stderrBase64 := bytesToTextAndBase64(resp.Stderr)
|
||||
|
||||
return marshalToolResult(execResult{
|
||||
PID: resp.Pid,
|
||||
Status: resp.Status,
|
||||
StdoutLen: len(resp.Stdout),
|
||||
StderrLen: len(resp.Stderr),
|
||||
StdoutText: stdoutText,
|
||||
StdoutBase64: stdoutBase64,
|
||||
StderrText: stderrText,
|
||||
StderrBase64: stderrBase64,
|
||||
})
|
||||
}
|
||||
|
||||
func processToResult(process *commonpb.Process) *processResult {
|
||||
if process == nil {
|
||||
return nil
|
||||
}
|
||||
return &processResult{
|
||||
PID: process.Pid,
|
||||
PPID: process.Ppid,
|
||||
Executable: process.Executable,
|
||||
Owner: process.Owner,
|
||||
Architecture: process.Architecture,
|
||||
SessionID: process.SessionID,
|
||||
CmdLine: append([]string(nil), process.CmdLine...),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,339 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
consts "github.com/bishopfox/sliver/client/constants"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
func publishAIConversationEvent(event *clientpb.AIConversationEvent) {
|
||||
if event == nil {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(event.GetConversationID()) == "" {
|
||||
switch {
|
||||
case event.GetConversation() != nil:
|
||||
event.ConversationID = event.GetConversation().GetID()
|
||||
case event.GetMessage() != nil:
|
||||
event.ConversationID = event.GetMessage().GetConversationID()
|
||||
}
|
||||
}
|
||||
if event.GetConversation() != nil && len(event.GetConversation().GetMessages()) > 0 {
|
||||
event.Conversation = aiConversationSummary(event.GetConversation())
|
||||
}
|
||||
|
||||
data, err := proto.Marshal(event)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to marshal AI conversation event: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
core.EventBroker.Publish(core.Event{
|
||||
EventType: consts.AIConversationEvent,
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
func aiConversationSummary(conversation *clientpb.AIConversation) *clientpb.AIConversation {
|
||||
if conversation == nil {
|
||||
return nil
|
||||
}
|
||||
summary, ok := proto.Clone(conversation).(*clientpb.AIConversation)
|
||||
if !ok {
|
||||
return conversation
|
||||
}
|
||||
summary.Messages = nil
|
||||
return summary
|
||||
}
|
||||
|
||||
func loadAIConversationSummary(conversationID string) (*clientpb.AIConversation, error) {
|
||||
conversation, err := db.AIConversationByID(conversationID, "", false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return aiConversationSummary(conversation), nil
|
||||
}
|
||||
|
||||
func publishAIConversationSummaryEvent(eventType clientpb.AIConversationEventType, conversation *clientpb.AIConversation, turnID string, errText string) {
|
||||
publishAIConversationEvent(&clientpb.AIConversationEvent{
|
||||
EventType: eventType,
|
||||
ConversationID: selectedAIConversationID(conversation),
|
||||
TurnID: strings.TrimSpace(turnID),
|
||||
ErrorText: strings.TrimSpace(errText),
|
||||
Conversation: aiConversationSummary(conversation),
|
||||
})
|
||||
}
|
||||
|
||||
func publishAIConversationMessageEvent(
|
||||
eventType clientpb.AIConversationEventType,
|
||||
conversation *clientpb.AIConversation,
|
||||
message *clientpb.AIConversationMessage,
|
||||
turnID string,
|
||||
errText string,
|
||||
) {
|
||||
publishAIConversationEvent(&clientpb.AIConversationEvent{
|
||||
EventType: eventType,
|
||||
ConversationID: selectedAIConversationID(conversation),
|
||||
TurnID: strings.TrimSpace(turnID),
|
||||
ErrorText: strings.TrimSpace(errText),
|
||||
Conversation: aiConversationSummary(conversation),
|
||||
Message: message,
|
||||
})
|
||||
}
|
||||
|
||||
func selectedAIConversationID(conversation *clientpb.AIConversation) string {
|
||||
if conversation == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(conversation.GetID())
|
||||
}
|
||||
|
||||
type aiConversationEventSink struct {
|
||||
conversationID string
|
||||
operatorName string
|
||||
runtime *serverai.RuntimeConfig
|
||||
turnID string
|
||||
}
|
||||
|
||||
func newAIConversationEventSink(conversationID string, operatorName string, runtime *serverai.RuntimeConfig, turnID string) *aiConversationEventSink {
|
||||
return &aiConversationEventSink{
|
||||
conversationID: strings.TrimSpace(conversationID),
|
||||
operatorName: strings.TrimSpace(operatorName),
|
||||
runtime: runtime,
|
||||
turnID: strings.TrimSpace(turnID),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) TurnStarted() error {
|
||||
conversation, err := s.updateConversationTurnState(s.turnID, clientpb.AIConversationTurnState_AI_TURN_STATE_IN_PROGRESS)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
publishAIConversationSummaryEvent(clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_TURN_STARTED, conversation, s.turnID, "")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) TurnCompleted() error {
|
||||
conversation, err := s.updateConversationTurnState("", clientpb.AIConversationTurnState_AI_TURN_STATE_IDLE)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
publishAIConversationSummaryEvent(clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_TURN_COMPLETED, conversation, s.turnID, "")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) TurnFailed(failure error) error {
|
||||
errText := ""
|
||||
if failure != nil {
|
||||
errText = strings.TrimSpace(failure.Error())
|
||||
}
|
||||
conversation, err := s.updateConversationTurnState("", clientpb.AIConversationTurnState_AI_TURN_STATE_FAILED)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
publishAIConversationSummaryEvent(clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_TURN_FAILED, conversation, s.turnID, errText)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) ChatMessage(ctx context.Context, item serverai.AgenticChatMessage) error {
|
||||
visibility := clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT
|
||||
if item.UIOnly {
|
||||
visibility = clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY
|
||||
}
|
||||
role := strings.ToLower(strings.TrimSpace(item.Role))
|
||||
if role == "" {
|
||||
role = "assistant"
|
||||
}
|
||||
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: s.conversationID,
|
||||
OperatorName: s.operatorName,
|
||||
Provider: s.provider(),
|
||||
Model: s.model(),
|
||||
Role: role,
|
||||
Content: strings.TrimSpace(item.Content),
|
||||
FinishReason: strings.TrimSpace(item.Status),
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: visibility,
|
||||
IncludeInContext: aiOptionalBool(item.IncludeInContext),
|
||||
State: normalizeAIMessageState(item.Status, clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED),
|
||||
TurnID: s.turnID,
|
||||
ItemID: strings.TrimSpace(item.ItemID),
|
||||
}
|
||||
_, err := s.saveMessageAndPublish(message, clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED, "")
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) ReasoningItem(ctx context.Context, item serverai.AgenticReasoningItem) error {
|
||||
content := strings.TrimSpace(item.Content)
|
||||
summary := strings.TrimSpace(item.Summary)
|
||||
if content == "" {
|
||||
content = summary
|
||||
} else if summary != "" && summary != content {
|
||||
content = "Summary:\n" + summary + "\n\n" + content
|
||||
}
|
||||
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: s.conversationID,
|
||||
OperatorName: s.operatorName,
|
||||
Provider: s.provider(),
|
||||
Model: s.model(),
|
||||
Role: "assistant",
|
||||
Content: content,
|
||||
FinishReason: strings.TrimSpace(item.Status),
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_REASONING,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY,
|
||||
IncludeInContext: aiOptionalBool(false),
|
||||
State: normalizeAIMessageState(item.Status, clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED),
|
||||
TurnID: s.turnID,
|
||||
ItemID: strings.TrimSpace(item.ItemID),
|
||||
}
|
||||
_, err := s.saveMessageAndPublish(message, clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED, "")
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) ToolCallStarted(ctx context.Context, item serverai.AgenticToolCall) error {
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: s.conversationID,
|
||||
OperatorName: s.operatorName,
|
||||
Provider: s.provider(),
|
||||
Model: s.model(),
|
||||
Role: "assistant",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_TOOL_CALL,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY,
|
||||
IncludeInContext: aiOptionalBool(false),
|
||||
State: clientpb.AIConversationMessageState_AI_MESSAGE_STATE_IN_PROGRESS,
|
||||
TurnID: s.turnID,
|
||||
ItemID: strings.TrimSpace(item.ItemID),
|
||||
ToolCallID: strings.TrimSpace(item.CallID),
|
||||
ToolName: strings.TrimSpace(item.Name),
|
||||
ToolArguments: strings.TrimSpace(item.Arguments),
|
||||
FinishReason: strings.TrimSpace(item.Status),
|
||||
}
|
||||
_, err := s.saveMessageAndPublish(message, clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_STARTED, "")
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) ToolCallCompleted(ctx context.Context, item serverai.AgenticToolCallResult) error {
|
||||
errText := ""
|
||||
state := clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED
|
||||
if item.Error != nil {
|
||||
errText = strings.TrimSpace(item.Error.Error())
|
||||
state = clientpb.AIConversationMessageState_AI_MESSAGE_STATE_FAILED
|
||||
}
|
||||
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: s.conversationID,
|
||||
OperatorName: s.operatorName,
|
||||
Provider: s.provider(),
|
||||
Model: s.model(),
|
||||
Role: "assistant",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_TOOL_CALL,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY,
|
||||
IncludeInContext: aiOptionalBool(false),
|
||||
State: state,
|
||||
TurnID: s.turnID,
|
||||
ItemID: strings.TrimSpace(item.ItemID),
|
||||
ToolCallID: strings.TrimSpace(item.CallID),
|
||||
ToolName: strings.TrimSpace(item.Name),
|
||||
ToolArguments: strings.TrimSpace(item.Arguments),
|
||||
ToolResult: strings.TrimSpace(item.Output),
|
||||
ErrorText: errText,
|
||||
FinishReason: strings.TrimSpace(item.Status),
|
||||
}
|
||||
_, err := s.saveMessageAndPublish(message, clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED, errText)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) SaveFailureMessage(content string, errText string) error {
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: s.conversationID,
|
||||
OperatorName: s.operatorName,
|
||||
Provider: s.provider(),
|
||||
Model: s.model(),
|
||||
Role: "system",
|
||||
Content: strings.TrimSpace(content),
|
||||
FinishReason: "error",
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_UI_ONLY,
|
||||
IncludeInContext: aiOptionalBool(false),
|
||||
State: clientpb.AIConversationMessageState_AI_MESSAGE_STATE_FAILED,
|
||||
TurnID: s.turnID,
|
||||
ErrorText: strings.TrimSpace(errText),
|
||||
}
|
||||
_, err := s.saveMessageAndPublish(message, clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED, errText)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) saveMessageAndPublish(
|
||||
message *clientpb.AIConversationMessage,
|
||||
eventType clientpb.AIConversationEventType,
|
||||
errText string,
|
||||
) (*clientpb.AIConversationMessage, error) {
|
||||
savedMessage, err := db.SaveAIConversationMessage(message, s.operatorName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conversation, err := loadAIConversationSummary(s.conversationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
publishAIConversationMessageEvent(eventType, conversation, savedMessage, s.turnID, errText)
|
||||
return savedMessage, nil
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) updateConversationTurnState(activeTurnID string, state clientpb.AIConversationTurnState) (*clientpb.AIConversation, error) {
|
||||
conversation, err := db.AIConversationByID(s.conversationID, "", false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conversation.ActiveTurnID = strings.TrimSpace(activeTurnID)
|
||||
conversation.TurnState = state
|
||||
if conversation.GetProvider() == "" {
|
||||
conversation.Provider = s.provider()
|
||||
}
|
||||
if conversation.GetModel() == "" {
|
||||
conversation.Model = s.model()
|
||||
}
|
||||
if _, err := db.SaveAIConversation(conversation, s.operatorName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return loadAIConversationSummary(s.conversationID)
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) provider() string {
|
||||
if s != nil && s.runtime != nil {
|
||||
return strings.TrimSpace(s.runtime.Provider)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *aiConversationEventSink) model() string {
|
||||
if s != nil && s.runtime != nil {
|
||||
return strings.TrimSpace(s.runtime.Model)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeAIMessageState(status string, fallback clientpb.AIConversationMessageState) clientpb.AIConversationMessageState {
|
||||
switch strings.ToLower(strings.TrimSpace(status)) {
|
||||
case "in_progress":
|
||||
return clientpb.AIConversationMessageState_AI_MESSAGE_STATE_IN_PROGRESS
|
||||
case "incomplete", "failed", "error":
|
||||
return clientpb.AIConversationMessageState_AI_MESSAGE_STATE_FAILED
|
||||
case "completed":
|
||||
return clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
func aiOptionalBool(value bool) *bool {
|
||||
return &value
|
||||
}
|
||||
+158
-66
@@ -5,16 +5,15 @@ import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
consts "github.com/bishopfox/sliver/client/constants"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
serverai "github.com/bishopfox/sliver/server/ai"
|
||||
"github.com/bishopfox/sliver/server/configs"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"github.com/bishopfox/sliver/server/rpc/aitools"
|
||||
"github.com/gofrs/uuid"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
func (rpc *Server) aiOperatorName(ctx context.Context, fallback string) string {
|
||||
@@ -35,21 +34,15 @@ func validateAIProvider(provider string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func publishAIConversationEvent(conversation *clientpb.AIConversation) {
|
||||
if conversation == nil {
|
||||
return
|
||||
func validateAIThinkingLevel(thinkingLevel string) error {
|
||||
thinkingLevel = strings.TrimSpace(thinkingLevel)
|
||||
if thinkingLevel == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
data, err := proto.Marshal(conversation)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to marshal AI conversation event: %s", err)
|
||||
return
|
||||
if serverai.NormalizeThinkingLevelValue(thinkingLevel) == "" {
|
||||
return status.Error(codes.InvalidArgument, fmt.Sprintf("unsupported AI thinking level %q", thinkingLevel))
|
||||
}
|
||||
|
||||
core.EventBroker.Publish(core.Event{
|
||||
EventType: consts.AIConversationEvent,
|
||||
Data: data,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (rpc *Server) GetAIProviders(ctx context.Context, _ *commonpb.Empty) (*clientpb.AIProviderConfigs, error) {
|
||||
@@ -86,13 +79,16 @@ func (rpc *Server) SaveAIConversation(ctx context.Context, req *clientpb.AIConve
|
||||
if err := validateAIProvider(req.Provider); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := validateAIThinkingLevel(req.ThinkingLevel); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req.OperatorName = rpc.aiOperatorName(ctx, req.OperatorName)
|
||||
conversation, err := db.SaveAIConversation(req, req.OperatorName)
|
||||
if err != nil {
|
||||
return nil, rpcError(err)
|
||||
}
|
||||
publishAIConversationEvent(conversation)
|
||||
publishAIConversationSummaryEvent(clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_CONVERSATION_UPDATED, conversation, "", "")
|
||||
return conversation, nil
|
||||
}
|
||||
|
||||
@@ -106,9 +102,13 @@ func (rpc *Server) DeleteAIConversation(ctx context.Context, req *clientpb.AICon
|
||||
if err != nil {
|
||||
return nil, rpcError(err)
|
||||
}
|
||||
publishAIConversationEvent(&clientpb.AIConversation{
|
||||
ID: req.ID,
|
||||
OperatorName: operatorName,
|
||||
publishAIConversationEvent(&clientpb.AIConversationEvent{
|
||||
EventType: clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_CONVERSATION_DELETED,
|
||||
ConversationID: req.ID,
|
||||
Conversation: &clientpb.AIConversation{
|
||||
ID: req.ID,
|
||||
OperatorName: operatorName,
|
||||
},
|
||||
})
|
||||
return &commonpb.Empty{}, nil
|
||||
}
|
||||
@@ -147,13 +147,19 @@ func (rpc *Server) SaveAIConversationMessage(ctx context.Context, req *clientpb.
|
||||
return nil, rpcError(err)
|
||||
}
|
||||
|
||||
conversation, err := db.AIConversationByID(message.ConversationID, "", false)
|
||||
conversation, err := loadAIConversationSummary(message.GetConversationID())
|
||||
if err != nil {
|
||||
return nil, rpcError(err)
|
||||
}
|
||||
publishAIConversationEvent(conversation)
|
||||
publishAIConversationMessageEvent(
|
||||
clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED,
|
||||
conversation,
|
||||
message,
|
||||
message.GetTurnID(),
|
||||
message.GetErrorText(),
|
||||
)
|
||||
if shouldTriggerAICompletion(req, message) {
|
||||
go runAIConversationCompletion(message.GetConversationID(), message.GetOperatorName())
|
||||
go rpc.runAIConversationCompletion(message.GetConversationID(), message.GetOperatorName())
|
||||
}
|
||||
|
||||
return message, nil
|
||||
@@ -169,7 +175,7 @@ func shouldTriggerAICompletion(req *clientpb.AIConversationMessage, saved *clien
|
||||
return strings.EqualFold(strings.TrimSpace(saved.GetRole()), "user")
|
||||
}
|
||||
|
||||
func runAIConversationCompletion(conversationID string, operatorName string) {
|
||||
func (rpc *Server) runAIConversationCompletion(conversationID string, operatorName string) {
|
||||
conversation, err := db.AIConversationByID(conversationID, "", true)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to load AI conversation %q for completion: %s", conversationID, err)
|
||||
@@ -184,18 +190,51 @@ func runAIConversationCompletion(conversationID string, operatorName string) {
|
||||
conversation = updatedConversation
|
||||
}
|
||||
}
|
||||
if resolveErr != nil {
|
||||
persistAIConversationFailure(conversationID, operatorName, runtime, resolveErr)
|
||||
return
|
||||
}
|
||||
|
||||
completion, err := serverai.CompleteConversation(context.Background(), runtime, conversation)
|
||||
turnID, err := newAITurnID()
|
||||
if err != nil {
|
||||
persistAIConversationFailure(conversationID, operatorName, runtime, err)
|
||||
rpcLog.Warnf("Failed to create AI turn id for %q: %s", conversationID, err)
|
||||
return
|
||||
}
|
||||
sink := newAIConversationEventSink(conversationID, operatorName, runtime, turnID)
|
||||
if err := sink.TurnStarted(); err != nil {
|
||||
rpcLog.Warnf("Failed to mark AI turn %q as started for %q: %s", turnID, conversationID, err)
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := db.SaveAIConversationMessage(&clientpb.AIConversationMessage{
|
||||
if resolveErr != nil {
|
||||
persistAIConversationFailure(sink, resolveErr)
|
||||
if err := sink.TurnFailed(resolveErr); err != nil {
|
||||
rpcLog.Warnf("Failed to mark AI turn %q as failed for %q: %s", turnID, conversationID, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var completion *serverai.Completion
|
||||
if serverai.SupportsAgenticConversation(runtime) {
|
||||
completion, err = serverai.CompleteConversationAgentic(
|
||||
context.Background(),
|
||||
runtime,
|
||||
conversation,
|
||||
aitools.NewExecutor(rpc, conversation),
|
||||
sink,
|
||||
)
|
||||
} else {
|
||||
completion, err = serverai.CompleteConversation(context.Background(), runtime, conversation)
|
||||
}
|
||||
if err != nil {
|
||||
persistAIConversationFailure(sink, err)
|
||||
if turnErr := sink.TurnFailed(err); turnErr != nil {
|
||||
rpcLog.Warnf("Failed to mark AI turn %q as failed for %q: %s", turnID, conversationID, turnErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
if updatedConversation, err := persistAIConversationUsage(conversation, operatorName, completion); err != nil {
|
||||
rpcLog.Warnf("Failed to persist AI context window usage for %q: %s", conversationID, err)
|
||||
} else if updatedConversation != nil {
|
||||
conversation = updatedConversation
|
||||
}
|
||||
|
||||
message, err := db.SaveAIConversationMessage(&clientpb.AIConversationMessage{
|
||||
ConversationID: conversationID,
|
||||
OperatorName: operatorName,
|
||||
Provider: completion.Provider,
|
||||
@@ -204,12 +243,32 @@ func runAIConversationCompletion(conversationID string, operatorName string) {
|
||||
Content: completion.Content,
|
||||
ProviderMessageID: completion.ProviderMessageID,
|
||||
FinishReason: completion.FinishReason,
|
||||
}, operatorName); err != nil {
|
||||
Kind: clientpb.AIConversationMessageKind_AI_MESSAGE_KIND_CHAT,
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: aiOptionalBool(true),
|
||||
State: clientpb.AIConversationMessageState_AI_MESSAGE_STATE_COMPLETED,
|
||||
TurnID: turnID,
|
||||
}, operatorName)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to save AI assistant reply for %q: %s", conversationID, err)
|
||||
return
|
||||
}
|
||||
|
||||
publishAIConversationByID(conversationID)
|
||||
summary, err := loadAIConversationSummary(conversationID)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to load AI conversation summary for %q: %s", conversationID, err)
|
||||
} else {
|
||||
publishAIConversationMessageEvent(
|
||||
clientpb.AIConversationEventType_AI_CONVERSATION_EVENT_TYPE_MESSAGE_COMPLETED,
|
||||
summary,
|
||||
message,
|
||||
turnID,
|
||||
"",
|
||||
)
|
||||
}
|
||||
if err := sink.TurnCompleted(); err != nil {
|
||||
rpcLog.Warnf("Failed to mark AI turn %q as completed for %q: %s", turnID, conversationID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func persistAIConversationRuntime(conversation *clientpb.AIConversation, operatorName string, runtime *serverai.RuntimeConfig) (*clientpb.AIConversation, error) {
|
||||
@@ -227,49 +286,82 @@ func persistAIConversationRuntime(conversation *clientpb.AIConversation, operato
|
||||
}
|
||||
|
||||
return db.SaveAIConversation(&clientpb.AIConversation{
|
||||
ID: conversation.GetID(),
|
||||
OperatorName: conversation.GetOperatorName(),
|
||||
Provider: provider,
|
||||
Model: model,
|
||||
Title: conversation.GetTitle(),
|
||||
Summary: conversation.GetSummary(),
|
||||
SystemPrompt: conversation.GetSystemPrompt(),
|
||||
ID: conversation.GetID(),
|
||||
OperatorName: conversation.GetOperatorName(),
|
||||
Provider: provider,
|
||||
Model: model,
|
||||
ThinkingLevel: conversation.GetThinkingLevel(),
|
||||
Title: conversation.GetTitle(),
|
||||
Summary: conversation.GetSummary(),
|
||||
SystemPrompt: conversation.GetSystemPrompt(),
|
||||
ActiveTurnID: conversation.GetActiveTurnID(),
|
||||
TurnState: conversation.GetTurnState(),
|
||||
TargetSessionID: conversation.GetTargetSessionID(),
|
||||
TargetBeaconID: conversation.GetTargetBeaconID(),
|
||||
ContextWindowUsage: conversation.GetContextWindowUsage(),
|
||||
}, operatorName)
|
||||
}
|
||||
|
||||
func persistAIConversationFailure(conversationID string, operatorName string, runtime *serverai.RuntimeConfig, failure error) {
|
||||
content := "AI request failed."
|
||||
if failure != nil && strings.TrimSpace(failure.Error()) != "" {
|
||||
content = "AI request failed: " + strings.TrimSpace(failure.Error())
|
||||
func persistAIConversationUsage(conversation *clientpb.AIConversation, operatorName string, completion *serverai.Completion) (*clientpb.AIConversation, error) {
|
||||
if conversation == nil || completion == nil || completion.ContextWindowUsage == nil {
|
||||
return conversation, nil
|
||||
}
|
||||
|
||||
message := &clientpb.AIConversationMessage{
|
||||
ConversationID: conversationID,
|
||||
OperatorName: operatorName,
|
||||
Role: "system",
|
||||
Content: content,
|
||||
FinishReason: "error",
|
||||
}
|
||||
if runtime != nil {
|
||||
message.Provider = runtime.Provider
|
||||
message.Model = runtime.Model
|
||||
usage := completion.ContextWindowUsage
|
||||
current := conversation.GetContextWindowUsage()
|
||||
if current != nil &&
|
||||
current.GetInputTokens() == usage.InputTokens &&
|
||||
current.GetOutputTokens() == usage.OutputTokens &&
|
||||
current.GetTotalTokens() == usage.TotalTokens &&
|
||||
current.GetContextWindowTokens() == usage.ContextWindowTokens &&
|
||||
current.GetContextWindowTokensEstimated() == usage.ContextWindowTokensEstimated {
|
||||
return conversation, nil
|
||||
}
|
||||
|
||||
if _, err := db.SaveAIConversationMessage(message, operatorName); err != nil {
|
||||
rpcLog.Warnf("Failed to save AI failure message for %q: %s", conversationID, err)
|
||||
return
|
||||
}
|
||||
|
||||
publishAIConversationByID(conversationID)
|
||||
return db.SaveAIConversation(&clientpb.AIConversation{
|
||||
ID: conversation.GetID(),
|
||||
OperatorName: conversation.GetOperatorName(),
|
||||
Provider: fallbackAIString(conversation.GetProvider(), completion.Provider),
|
||||
Model: fallbackAIString(conversation.GetModel(), completion.Model),
|
||||
ThinkingLevel: conversation.GetThinkingLevel(),
|
||||
Title: conversation.GetTitle(),
|
||||
Summary: conversation.GetSummary(),
|
||||
SystemPrompt: conversation.GetSystemPrompt(),
|
||||
ActiveTurnID: conversation.GetActiveTurnID(),
|
||||
TurnState: conversation.GetTurnState(),
|
||||
TargetSessionID: conversation.GetTargetSessionID(),
|
||||
TargetBeaconID: conversation.GetTargetBeaconID(),
|
||||
ContextWindowUsage: &clientpb.AIContextWindowUsage{
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
TotalTokens: usage.TotalTokens,
|
||||
ContextWindowTokens: usage.ContextWindowTokens,
|
||||
ContextWindowTokensEstimated: usage.ContextWindowTokensEstimated,
|
||||
},
|
||||
}, operatorName)
|
||||
}
|
||||
|
||||
func publishAIConversationByID(conversationID string) {
|
||||
conversation, err := db.AIConversationByID(conversationID, "", false)
|
||||
if err != nil {
|
||||
rpcLog.Warnf("Failed to publish AI conversation update for %q: %s", conversationID, err)
|
||||
func persistAIConversationFailure(sink *aiConversationEventSink, failure error) {
|
||||
content := "AI request failed."
|
||||
errText := ""
|
||||
if failure != nil && strings.TrimSpace(failure.Error()) != "" {
|
||||
errText = strings.TrimSpace(failure.Error())
|
||||
content = "AI request failed: " + errText
|
||||
}
|
||||
if sink == nil {
|
||||
return
|
||||
}
|
||||
publishAIConversationEvent(conversation)
|
||||
if err := sink.SaveFailureMessage(content, errText); err != nil {
|
||||
rpcLog.Warnf("Failed to save AI failure message for %q: %s", sink.conversationID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func newAITurnID() (string, error) {
|
||||
id, err := uuid.NewV4()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return id.String(), nil
|
||||
}
|
||||
|
||||
func fallbackAIString(value string, fallback string) string {
|
||||
|
||||
+918
-15
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,15 @@
|
||||
package clientpbutil
|
||||
|
||||
import "github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
|
||||
// AIConversationMessageIncludesContext reports whether a message should be
|
||||
// included in the model-visible context window.
|
||||
func AIConversationMessageIncludesContext(message *clientpb.AIConversationMessage) bool {
|
||||
if message == nil {
|
||||
return false
|
||||
}
|
||||
if message.IncludeInContext != nil {
|
||||
return message.GetIncludeInContext()
|
||||
}
|
||||
return message.GetVisibility() == clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package clientpbutil
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
)
|
||||
|
||||
func TestAIConversationMessageIncludesContextHonorsExplicitFlag(t *testing.T) {
|
||||
message := &clientpb.AIConversationMessage{
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
IncludeInContext: boolPtr(false),
|
||||
}
|
||||
|
||||
if AIConversationMessageIncludesContext(message) {
|
||||
t.Fatalf("expected explicit include-in-context=false to exclude the message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIConversationMessageIncludesContextFallsBackToVisibility(t *testing.T) {
|
||||
message := &clientpb.AIConversationMessage{
|
||||
Visibility: clientpb.AIConversationMessageVisibility_AI_MESSAGE_VISIBILITY_CONTEXT,
|
||||
}
|
||||
|
||||
if !AIConversationMessageIncludesContext(message) {
|
||||
t.Fatalf("expected context visibility to include the message when the flag is unset")
|
||||
}
|
||||
}
|
||||
|
||||
func boolPtr(value bool) *bool {
|
||||
return &value
|
||||
}
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
package progress
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"image/color"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
"charm.land/lipgloss/v2"
|
||||
"github.com/charmbracelet/x/ansi"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultFullCharHalfBlock = '▌'
|
||||
DefaultFullCharFullBlock = '█'
|
||||
DefaultEmptyCharBlock = '░'
|
||||
|
||||
defaultWidth = 40
|
||||
)
|
||||
|
||||
var (
|
||||
defaultFullColor = lipgloss.Color("#7571F9")
|
||||
defaultEmptyColor = lipgloss.Color("#606060")
|
||||
)
|
||||
|
||||
type Option func(*Model)
|
||||
|
||||
func WithColors(colors ...color.Color) Option {
|
||||
if len(colors) == 0 {
|
||||
return func(m *Model) {
|
||||
m.FullColor = defaultFullColor
|
||||
m.blend = nil
|
||||
}
|
||||
}
|
||||
if len(colors) == 1 {
|
||||
return func(m *Model) {
|
||||
m.FullColor = colors[0]
|
||||
m.blend = nil
|
||||
}
|
||||
}
|
||||
return func(m *Model) {
|
||||
m.blend = colors
|
||||
}
|
||||
}
|
||||
|
||||
func WithFillCharacters(full, empty rune) Option {
|
||||
return func(m *Model) {
|
||||
m.Full = full
|
||||
m.Empty = empty
|
||||
}
|
||||
}
|
||||
|
||||
func WithoutPercentage() Option {
|
||||
return func(m *Model) {
|
||||
m.ShowPercentage = false
|
||||
}
|
||||
}
|
||||
|
||||
func WithWidth(w int) Option {
|
||||
return func(m *Model) {
|
||||
m.SetWidth(w)
|
||||
}
|
||||
}
|
||||
|
||||
type Model struct {
|
||||
width int
|
||||
Full rune
|
||||
FullColor color.Color
|
||||
Empty rune
|
||||
EmptyColor color.Color
|
||||
ShowPercentage bool
|
||||
PercentFormat string
|
||||
PercentageStyle lipgloss.Style
|
||||
blend []color.Color
|
||||
}
|
||||
|
||||
func New(opts ...Option) Model {
|
||||
m := Model{
|
||||
width: defaultWidth,
|
||||
Full: DefaultFullCharHalfBlock,
|
||||
FullColor: defaultFullColor,
|
||||
Empty: DefaultEmptyCharBlock,
|
||||
EmptyColor: defaultEmptyColor,
|
||||
ShowPercentage: true,
|
||||
PercentFormat: " %3.0f%%",
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(&m)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *Model) SetWidth(w int) {
|
||||
m.width = w
|
||||
}
|
||||
|
||||
func (m Model) Width() int {
|
||||
return m.width
|
||||
}
|
||||
|
||||
func (m Model) View() string {
|
||||
return m.ViewAs(0)
|
||||
}
|
||||
|
||||
func (m Model) ViewAs(percent float64) string {
|
||||
b := strings.Builder{}
|
||||
percentView := m.percentageView(percent)
|
||||
m.barView(&b, percent, ansi.StringWidth(percentView))
|
||||
b.WriteString(percentView)
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func (m Model) percentageView(percent float64) string {
|
||||
if !m.ShowPercentage {
|
||||
return ""
|
||||
}
|
||||
percent = math.Max(0, math.Min(1, percent))
|
||||
percentage := fmt.Sprintf(m.PercentFormat, percent*100)
|
||||
return m.PercentageStyle.Inline(true).Render(percentage)
|
||||
}
|
||||
|
||||
func (m Model) barView(b *strings.Builder, percent float64, textWidth int) {
|
||||
totalWidth := maxInt(0, m.width-textWidth)
|
||||
filledWidth := int(math.Round(float64(totalWidth) * math.Max(0, math.Min(1, percent))))
|
||||
filledWidth = clampInt(filledWidth, 0, totalWidth)
|
||||
|
||||
if len(m.blend) > 1 {
|
||||
blend := lipgloss.Blend1D(totalWidth, m.blend...)
|
||||
for i := 0; i < filledWidth && i < len(blend); i++ {
|
||||
b.WriteString(lipgloss.NewStyle().
|
||||
Foreground(blend[i]).
|
||||
Render(string(m.Full)))
|
||||
}
|
||||
} else {
|
||||
b.WriteString(lipgloss.NewStyle().
|
||||
Foreground(m.FullColor).
|
||||
Render(strings.Repeat(string(m.Full), filledWidth)))
|
||||
}
|
||||
|
||||
emptyWidth := maxInt(0, totalWidth-filledWidth)
|
||||
if emptyWidth == 0 {
|
||||
return
|
||||
}
|
||||
b.WriteString(lipgloss.NewStyle().
|
||||
Foreground(m.EmptyColor).
|
||||
Render(strings.Repeat(string(m.Empty), emptyWidth)))
|
||||
}
|
||||
|
||||
func maxInt(a, b int) int {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func clampInt(value, low, high int) int {
|
||||
if value < low {
|
||||
return low
|
||||
}
|
||||
if value > high {
|
||||
return high
|
||||
}
|
||||
return value
|
||||
}
|
||||
+16
-6
@@ -199,7 +199,7 @@ func (c *Client) Initialize(
|
||||
}
|
||||
// Add elicitation capability if handler is configured
|
||||
if c.elicitationHandler != nil {
|
||||
capabilities.Elicitation = &struct{}{}
|
||||
capabilities.Elicitation = &mcp.ElicitationCapability{}
|
||||
}
|
||||
|
||||
// Ensure we send a params object with all required fields
|
||||
@@ -213,6 +213,11 @@ func (c *Client) Initialize(
|
||||
Capabilities: capabilities,
|
||||
}
|
||||
|
||||
// By default, use client supported latest protocol version if version not specified
|
||||
if params.ProtocolVersion == "" {
|
||||
params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
|
||||
}
|
||||
|
||||
response, err := c.sendRequest(ctx, "initialize", params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -267,7 +272,7 @@ func (c *Client) ListResourcesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourcesRequest,
|
||||
) (*mcp.ListResourcesResult, error) {
|
||||
result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, "resources/list")
|
||||
result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, request.Header, "resources/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -303,7 +308,7 @@ func (c *Client) ListResourceTemplatesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourceTemplatesRequest,
|
||||
) (*mcp.ListResourceTemplatesResult, error) {
|
||||
result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, "resources/templates/list")
|
||||
result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, request.Header, "resources/templates/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -367,7 +372,7 @@ func (c *Client) ListPromptsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListPromptsRequest,
|
||||
) (*mcp.ListPromptsResult, error) {
|
||||
result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, "prompts/list")
|
||||
result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, request.Header, "prompts/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -415,7 +420,7 @@ func (c *Client) ListToolsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListToolsRequest,
|
||||
) (*mcp.ListToolsResult, error) {
|
||||
result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, "tools/list")
|
||||
result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, request.Header, "tools/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -629,6 +634,10 @@ func (c *Client) handleElicitationRequestTransport(ctx context.Context, request
|
||||
}
|
||||
}
|
||||
|
||||
if err := params.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("invalid elicitation params: %w", err)
|
||||
}
|
||||
|
||||
// Create the MCP request
|
||||
mcpRequest := mcp.ElicitationRequest{
|
||||
Request: mcp.Request{
|
||||
@@ -664,9 +673,10 @@ func listByPage[T any](
|
||||
ctx context.Context,
|
||||
client *Client,
|
||||
request mcp.PaginatedRequest,
|
||||
header http.Header,
|
||||
method string,
|
||||
) (*T, error) {
|
||||
response, err := client.sendRequest(ctx, method, request.Params, nil)
|
||||
response, err := client.sendRequest(ctx, method, request.Params, header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
+2
-2
@@ -11,8 +11,8 @@ import (
|
||||
type ElicitationHandler interface {
|
||||
// Elicit handles an elicitation request from the server and returns the user's response.
|
||||
// The implementation should:
|
||||
// 1. Present the request message to the user
|
||||
// 2. Validate input against the requested schema
|
||||
// 1. Present the request message to the user (and URL if in URL mode)
|
||||
// 2. Validate input against the requested schema (for form mode)
|
||||
// 3. Allow the user to accept, decline, or cancel
|
||||
// 4. Return the appropriate response
|
||||
Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
|
||||
|
||||
+7
@@ -20,6 +20,13 @@ func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
|
||||
return transport.WithHTTPClient(httpClient)
|
||||
}
|
||||
|
||||
// WithHTTPHost sets a custom Host header for the SSE client, enabling manual DNS resolution.
|
||||
// This allows connecting to an IP address while sending a specific Host header to the server.
|
||||
// For example, connecting to "http://192.168.1.100:8080/sse" but sending Host: "api.example.com"
|
||||
func WithHTTPHost(host string) transport.ClientOption {
|
||||
return transport.WithHTTPHost(host)
|
||||
}
|
||||
|
||||
// NewSSEMCPClient creates a new SSE-based MCP client with the given base URL.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewSSEMCPClient(baseURL string, options ...transport.ClientOption) (*Client, error) {
|
||||
|
||||
+38
-6
@@ -21,6 +21,8 @@ var ErrNoToken = errors.New("no token available")
|
||||
type OAuthConfig struct {
|
||||
// ClientID is the OAuth client ID
|
||||
ClientID string
|
||||
// ClientURI is the URI of the client
|
||||
ClientURI string
|
||||
// ClientSecret is the OAuth client secret (for confidential clients)
|
||||
ClientSecret string
|
||||
// RedirectURI is the redirect URI for the OAuth flow
|
||||
@@ -242,8 +244,21 @@ func (h *OAuthHandler) refreshToken(ctx context.Context, refreshToken string) (*
|
||||
return nil, extractOAuthError(body, resp.StatusCode, "refresh token request failed")
|
||||
}
|
||||
|
||||
// Read the response body for parsing
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read token response body: %w", err)
|
||||
}
|
||||
|
||||
// GitHub returns HTTP 200 even for errors, with error details in the JSON body
|
||||
// Check if the response contains an error field before parsing as Token
|
||||
var oauthErr OAuthError
|
||||
if err := json.Unmarshal(body, &oauthErr); err == nil && oauthErr.ErrorCode != "" {
|
||||
return nil, fmt.Errorf("refresh token request failed: %w", oauthErr)
|
||||
}
|
||||
|
||||
var tokenResp Token
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
if err := json.Unmarshal(body, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode token response: %w", err)
|
||||
}
|
||||
|
||||
@@ -414,14 +429,14 @@ func (h *OAuthHandler) getServerMetadata(ctx context.Context) (*AuthServerMetada
|
||||
// Use the first authorization server
|
||||
authServerURL := protectedResource.AuthorizationServers[0]
|
||||
|
||||
// Try OpenID Connect discovery first
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/openid-configuration")
|
||||
// Try OAuth Authorization Server Metadata first
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/oauth-authorization-server")
|
||||
if h.serverMetadata != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// If OpenID Connect discovery fails, try OAuth Authorization Server Metadata
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/oauth-authorization-server")
|
||||
// If OAuth Authorization Server Metadata discovery fails, try OpenID Connect discovery
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/openid-configuration")
|
||||
if h.serverMetadata != nil {
|
||||
return
|
||||
}
|
||||
@@ -548,6 +563,10 @@ func (h *OAuthHandler) RegisterClient(ctx context.Context, clientName string) er
|
||||
"scope": strings.Join(h.config.Scopes, " "),
|
||||
}
|
||||
|
||||
if h.config.ClientURI != "" {
|
||||
regRequest["client_uri"] = h.config.ClientURI
|
||||
}
|
||||
|
||||
// Add client_secret if this is a confidential client
|
||||
if h.config.ClientSecret != "" {
|
||||
regRequest["token_endpoint_auth_method"] = "client_secret_basic"
|
||||
@@ -665,8 +684,21 @@ func (h *OAuthHandler) ProcessAuthorizationResponse(ctx context.Context, code, s
|
||||
return extractOAuthError(body, resp.StatusCode, "token request failed")
|
||||
}
|
||||
|
||||
// Read the response body for parsing
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read token response body: %w", err)
|
||||
}
|
||||
|
||||
// GitHub returns HTTP 200 even for errors, with error details in the JSON body
|
||||
// Check if the response contains an error field before parsing as Token
|
||||
var oauthErr OAuthError
|
||||
if err := json.Unmarshal(body, &oauthErr); err == nil && oauthErr.ErrorCode != "" {
|
||||
return fmt.Errorf("token request failed: %w", oauthErr)
|
||||
}
|
||||
|
||||
var tokenResp Token
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
if err := json.Unmarshal(body, &tokenResp); err != nil {
|
||||
return fmt.Errorf("failed to decode token response: %w", err)
|
||||
}
|
||||
|
||||
|
||||
+100
-36
@@ -34,6 +34,7 @@ type SSE struct {
|
||||
endpointChan chan struct{}
|
||||
headers map[string]string
|
||||
headerFunc HTTPHeaderFunc
|
||||
host string
|
||||
logger util.Logger
|
||||
|
||||
started atomic.Bool
|
||||
@@ -80,6 +81,15 @@ func WithOAuth(config OAuthConfig) ClientOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPHost sets a custom Host header for the SSE client, enabling manual DNS resolution.
|
||||
// This allows connecting to an IP address while sending a specific Host header to the server.
|
||||
// For example, connecting to "http://192.168.1.100:8080/sse" but sending Host: "api.example.com"
|
||||
func WithHTTPHost(host string) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.host = host
|
||||
}
|
||||
}
|
||||
|
||||
// NewSSE creates a new SSE-based MCP client with the given base URL.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
|
||||
@@ -126,6 +136,11 @@ func (c *SSE) Start(ctx context.Context) error {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Set custom Host header if provided
|
||||
if c.host != "" {
|
||||
req.Host = c.host
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("Cache-Control", "no-cache")
|
||||
req.Header.Set("Connection", "keep-alive")
|
||||
@@ -162,11 +177,14 @@ func (c *SSE) Start(ctx context.Context) error {
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
resp.Body.Close()
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
// Handle unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
if c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return ErrUnauthorized
|
||||
}
|
||||
return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
|
||||
}
|
||||
@@ -174,16 +192,31 @@ func (c *SSE) Start(ctx context.Context) error {
|
||||
go c.readSSE(resp.Body)
|
||||
|
||||
// Wait for the endpoint to be received
|
||||
timeout := time.NewTimer(30 * time.Second)
|
||||
defer timeout.Stop()
|
||||
endpointTimeout := 30 * time.Second
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
remaining := time.Until(deadline)
|
||||
// If context deadline has already passed, return immediately
|
||||
if remaining <= 0 {
|
||||
cancel()
|
||||
return ctx.Err()
|
||||
}
|
||||
// Use the shorter of remaining time or default timeout
|
||||
if remaining < endpointTimeout {
|
||||
endpointTimeout = remaining
|
||||
}
|
||||
}
|
||||
|
||||
timer := time.NewTimer(endpointTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-c.endpointChan:
|
||||
// Endpoint received, proceed
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("context cancelled while waiting for endpoint")
|
||||
case <-timeout.C: // Add a timeout
|
||||
return fmt.Errorf("context cancelled while waiting for endpoint: %w", ctx.Err())
|
||||
case <-timer.C:
|
||||
cancel()
|
||||
return fmt.Errorf("timeout waiting for endpoint")
|
||||
return fmt.Errorf("timeout waiting for endpoint after %v", endpointTimeout)
|
||||
}
|
||||
|
||||
c.started.Store(true)
|
||||
@@ -212,22 +245,14 @@ func (c *SSE) readSSE(reader io.ReadCloser) {
|
||||
}
|
||||
c.handleSSEEvent(event, data)
|
||||
}
|
||||
break
|
||||
}
|
||||
// Checking whether the connection was terminated due to NO_ERROR in HTTP2 based on RFC9113
|
||||
// Only handle NO_ERROR specially if onConnectionLost handler is set to maintain backward compatibility
|
||||
if strings.Contains(err.Error(), "NO_ERROR") {
|
||||
c.connectionLostMu.RLock()
|
||||
handler := c.onConnectionLost
|
||||
c.connectionLostMu.RUnlock()
|
||||
|
||||
if handler != nil {
|
||||
// This is not actually an error - HTTP2 idle timeout disconnection
|
||||
handler(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if !c.closed.Load() {
|
||||
c.connectionLostMu.RLock()
|
||||
handler := c.onConnectionLost
|
||||
c.connectionLostMu.RUnlock()
|
||||
if handler != nil {
|
||||
// Notify that the connection will be closed due to an error
|
||||
handler(err)
|
||||
} else if err == io.EOF && !c.closed.Load() {
|
||||
c.logger.Errorf("SSE stream error: %v", err)
|
||||
}
|
||||
return
|
||||
@@ -249,10 +274,10 @@ func (c *SSE) readSSE(reader io.ReadCloser) {
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
data = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if after, ok := strings.CutPrefix(line, "event:"); ok {
|
||||
event = strings.TrimSpace(after)
|
||||
} else if after, ok := strings.CutPrefix(line, "data:"); ok {
|
||||
data = strings.TrimSpace(after)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -369,6 +394,11 @@ func (c *SSE) SendRequest(
|
||||
}
|
||||
}
|
||||
|
||||
// Set custom Host header if provided
|
||||
if c.host != "" {
|
||||
req.Host = c.host
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
@@ -416,6 +446,7 @@ func (c *SSE) SendRequest(
|
||||
resp.Body.Close()
|
||||
|
||||
if err != nil {
|
||||
deleteResponseChan()
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
@@ -423,20 +454,45 @@ func (c *SSE) SendRequest(
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
deleteResponseChan()
|
||||
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
// Handle unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
if c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
// Calculate response timeout
|
||||
responseTimeout := 60 * time.Second
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
remaining := time.Until(deadline)
|
||||
// Check if context deadline has already passed
|
||||
if remaining <= 0 {
|
||||
deleteResponseChan()
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
// Use the shorter of remaining time or default timeout
|
||||
if remaining < responseTimeout {
|
||||
responseTimeout = remaining
|
||||
}
|
||||
}
|
||||
|
||||
timer := time.NewTimer(responseTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
deleteResponseChan()
|
||||
return nil, ctx.Err()
|
||||
case <-timer.C:
|
||||
// Timeout handling
|
||||
deleteResponseChan()
|
||||
return nil, fmt.Errorf("timeout waiting for SSE response after %v", responseTimeout)
|
||||
case response, ok := <-responseChan:
|
||||
if ok {
|
||||
return response, nil
|
||||
@@ -534,6 +590,11 @@ func (c *SSE) SendNotification(ctx context.Context, notification mcp.JSONRPCNoti
|
||||
}
|
||||
}
|
||||
|
||||
// Set custom Host header if provided
|
||||
if c.host != "" {
|
||||
req.Host = c.host
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send notification: %w", err)
|
||||
@@ -541,11 +602,14 @@ func (c *SSE) SendNotification(ctx context.Context, notification mcp.JSONRPCNoti
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
// Handle unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
if c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return ErrUnauthorized
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
|
||||
+62
-24
@@ -16,6 +16,10 @@ import (
|
||||
"github.com/mark3labs/mcp-go/util"
|
||||
)
|
||||
|
||||
// ErrTransportClosed is returned when attempting to send a request or notification
|
||||
// to a transport that has already been closed.
|
||||
var ErrTransportClosed = errors.New("transport closed")
|
||||
|
||||
// Stdio implements the transport layer of the MCP protocol using stdio communication.
|
||||
// It launches a subprocess and communicates with it via standard input/output streams
|
||||
// using JSON-RPC messages. The client handles message routing between requests and
|
||||
@@ -32,8 +36,10 @@ type Stdio struct {
|
||||
stderr io.ReadCloser
|
||||
responses map[string]chan *JSONRPCResponse
|
||||
mu sync.RWMutex
|
||||
done chan struct{}
|
||||
onNotification func(mcp.JSONRPCNotification)
|
||||
done chan struct{}
|
||||
closeOnce sync.Once
|
||||
closeCleanupOnce sync.Once
|
||||
onNotification func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
onRequest RequestHandler
|
||||
requestMu sync.RWMutex
|
||||
@@ -203,33 +209,42 @@ func (c *Stdio) spawnCommand(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// closeDone safely closes the done channel exactly once, unblocking all
|
||||
// in-flight SendRequest calls. Safe to call from multiple goroutines.
|
||||
func (c *Stdio) closeDone() {
|
||||
c.closeOnce.Do(func() { close(c.done) })
|
||||
}
|
||||
|
||||
// Close shuts down the stdio client, closing the stdin pipe and waiting for the subprocess to exit.
|
||||
// Returns an error if there are issues closing stdin or waiting for the subprocess to terminate.
|
||||
// Safe to call multiple times and concurrently with readResponses calling closeDone().
|
||||
func (c *Stdio) Close() error {
|
||||
select {
|
||||
case <-c.done:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
// cancel all in-flight request
|
||||
close(c.done)
|
||||
// Signal all in-flight requests to unblock.
|
||||
c.closeDone()
|
||||
|
||||
if c.stdin != nil {
|
||||
if err := c.stdin.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close stdin: %w", err)
|
||||
// Perform resource cleanup exactly once, even if readResponses already
|
||||
// called closeDone() (e.g. server died). Without this, the old early-return
|
||||
// guard would skip stdin/stderr cleanup and cmd.Wait(), causing FD leaks
|
||||
// and zombie processes.
|
||||
var closeErr error
|
||||
c.closeCleanupOnce.Do(func() {
|
||||
if c.stdin != nil {
|
||||
if err := c.stdin.Close(); err != nil {
|
||||
closeErr = fmt.Errorf("failed to close stdin: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
if c.stderr != nil {
|
||||
if err := c.stderr.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close stderr: %w", err)
|
||||
if c.stderr != nil {
|
||||
if err := c.stderr.Close(); err != nil && closeErr == nil {
|
||||
closeErr = fmt.Errorf("failed to close stderr: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if c.cmd != nil {
|
||||
return c.cmd.Wait()
|
||||
}
|
||||
|
||||
return nil
|
||||
if c.cmd != nil {
|
||||
if err := c.cmd.Wait(); err != nil && closeErr == nil {
|
||||
closeErr = err
|
||||
}
|
||||
}
|
||||
})
|
||||
return closeErr
|
||||
}
|
||||
|
||||
// GetSessionId returns the session ID of the transport.
|
||||
@@ -270,6 +285,9 @@ func (c *Stdio) readResponses() {
|
||||
if err != io.EOF && !errors.Is(err, context.Canceled) {
|
||||
c.logger.Errorf("Error reading from stdout: %v", err)
|
||||
}
|
||||
// Signal done so in-flight SendRequest calls unblock
|
||||
// instead of hanging forever when the server dies.
|
||||
c.closeDone()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -338,8 +356,10 @@ func (c *Stdio) SendRequest(
|
||||
ctx context.Context,
|
||||
request JSONRPCRequest,
|
||||
) (*JSONRPCResponse, error) {
|
||||
// Check if context is already canceled before doing any work
|
||||
// Check if transport is closed or context is already canceled before doing any work
|
||||
select {
|
||||
case <-c.done:
|
||||
return nil, ErrTransportClosed
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
@@ -377,6 +397,16 @@ func (c *Stdio) SendRequest(
|
||||
}
|
||||
|
||||
select {
|
||||
case <-c.done:
|
||||
// Drain responseChan first: a valid response may have been delivered
|
||||
// just before readResponses closed the done channel on EOF.
|
||||
select {
|
||||
case response := <-responseChan:
|
||||
return response, nil
|
||||
default:
|
||||
}
|
||||
deleteResponseChan()
|
||||
return nil, ErrTransportClosed
|
||||
case <-ctx.Done():
|
||||
deleteResponseChan()
|
||||
return nil, ctx.Err()
|
||||
@@ -390,6 +420,14 @@ func (c *Stdio) SendNotification(
|
||||
ctx context.Context,
|
||||
notification mcp.JSONRPCNotification,
|
||||
) error {
|
||||
select {
|
||||
case <-c.done:
|
||||
return ErrTransportClosed
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
if c.stdin == nil {
|
||||
return fmt.Errorf("stdio client not started")
|
||||
}
|
||||
|
||||
+62
-40
@@ -87,6 +87,15 @@ func WithSession(sessionID string) StreamableHTTPCOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithStreamableHTTPHost sets a custom Host header for the StreamableHTTP client, enabling manual DNS resolution.
|
||||
// This allows connecting to an IP address while sending a specific Host header to the server.
|
||||
// For example, connecting to "http://192.168.1.100:8080/mcp" but sending Host: "api.example.com"
|
||||
func WithStreamableHTTPHost(host string) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.host = host
|
||||
}
|
||||
}
|
||||
|
||||
// StreamableHTTP implements Streamable HTTP transport.
|
||||
//
|
||||
// It transmits JSON-RPC messages over individual HTTP requests. One message per request.
|
||||
@@ -103,6 +112,7 @@ type StreamableHTTP struct {
|
||||
httpClient *http.Client
|
||||
headers map[string]string
|
||||
headerFunc HTTPHeaderFunc
|
||||
host string
|
||||
logger util.Logger
|
||||
getListeningEnabled bool
|
||||
|
||||
@@ -119,11 +129,11 @@ type StreamableHTTP struct {
|
||||
requestHandler RequestHandler
|
||||
requestMu sync.RWMutex
|
||||
|
||||
closed chan struct{}
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
|
||||
// OAuth support
|
||||
oauthHandler *OAuthHandler
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewStreamableHTTP creates a new Streamable HTTP transport with the given server URL.
|
||||
@@ -188,21 +198,14 @@ func (c *StreamableHTTP) Start(ctx context.Context) error {
|
||||
|
||||
// Close closes the all the HTTP connections to the server.
|
||||
func (c *StreamableHTTP) Close() error {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
// Cancel all in-flight requests
|
||||
close(c.closed)
|
||||
c.closeOnce.Do(func() {
|
||||
// Cancel all in-flight requests
|
||||
close(c.closed)
|
||||
|
||||
sessionId := c.sessionID.Load().(string)
|
||||
if sessionId != "" {
|
||||
c.sessionID.Store("")
|
||||
c.wg.Add(1)
|
||||
// notify server session closed
|
||||
go func() {
|
||||
defer c.wg.Done()
|
||||
sessionId := c.sessionID.Load().(string)
|
||||
if sessionId != "" {
|
||||
c.sessionID.Store("")
|
||||
// notify server session closed
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, c.serverURL.String(), nil)
|
||||
@@ -217,15 +220,19 @@ func (c *StreamableHTTP) Close() error {
|
||||
req.Header.Set(HeaderKeyProtocolVersion, version)
|
||||
}
|
||||
}
|
||||
|
||||
// Set custom Host header if provided
|
||||
if c.host != "" {
|
||||
req.Host = c.host
|
||||
}
|
||||
res, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to send close request: %v", err)
|
||||
return
|
||||
}
|
||||
res.Body.Close()
|
||||
}()
|
||||
}
|
||||
c.wg.Wait()
|
||||
}
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -268,17 +275,14 @@ func (c *StreamableHTTP) SendRequest(
|
||||
resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream", request.Header)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrSessionTerminated) && request.Method == string(mcp.MethodInitialize) {
|
||||
// If the request is initialize, should not return a SessionTerminated error
|
||||
// It should be a genuine endpoint-routing issue.
|
||||
// ( Fall through to return StatusCode checking. )
|
||||
} else {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
// Per the MCP spec's backwards compatibility section: a 404 on an
|
||||
// initialize POST means the server likely only supports legacy SSE.
|
||||
return nil, ErrLegacySSEServer
|
||||
}
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
|
||||
// Only proceed if we have a valid response.
|
||||
// When sendHTTP fails and resp is nil but method is mcp.MethodInitialize
|
||||
// defer resp.Body.Close() fails with nil pointer dereference.
|
||||
if resp == nil {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
@@ -287,11 +291,21 @@ func (c *StreamableHTTP) SendRequest(
|
||||
// Check if we got an error response
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
// Handle unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
if c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
|
||||
// Per the MCP spec's backwards compatibility section: if an initialize
|
||||
// POST receives an HTTP 4xx (e.g. 405 Method Not Allowed, 404 Not Found),
|
||||
// the server likely only supports the legacy HTTP+SSE transport.
|
||||
if request.Method == string(mcp.MethodInitialize) && resp.StatusCode >= 400 && resp.StatusCode < 500 {
|
||||
return nil, ErrLegacySSEServer
|
||||
}
|
||||
|
||||
// handle error response
|
||||
@@ -376,6 +390,11 @@ func (c *StreamableHTTP) sendHTTP(
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
// Set custom Host header if provided
|
||||
if c.host != "" {
|
||||
req.Host = c.host
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
@@ -532,10 +551,10 @@ func (c *StreamableHTTP) readSSE(ctx context.Context, reader io.ReadCloser, hand
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
data = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if eventStr, ok := strings.CutPrefix(line, "event:"); ok {
|
||||
event = strings.TrimSpace(eventStr)
|
||||
} else if dataStr, ok := strings.CutPrefix(line, "data:"); ok {
|
||||
data = strings.TrimSpace(dataStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -558,14 +577,17 @@ func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
switch resp.StatusCode {
|
||||
case http.StatusOK, http.StatusAccepted, http.StatusNoContent:
|
||||
return nil
|
||||
case http.StatusUnauthorized:
|
||||
if c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
|
||||
return ErrUnauthorized
|
||||
default:
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf(
|
||||
"notification failed with status %d: %s",
|
||||
@@ -573,8 +595,6 @@ func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) SetNotificationHandler(handler func(mcp.JSONRPCNotification)) {
|
||||
@@ -643,6 +663,8 @@ func (c *StreamableHTTP) listenForever(ctx context.Context) {
|
||||
var (
|
||||
ErrSessionTerminated = fmt.Errorf("session terminated (404). need to re-initialize")
|
||||
ErrGetMethodNotAllowed = fmt.Errorf("GET method not allowed")
|
||||
ErrUnauthorized = fmt.Errorf("unauthorized (401)")
|
||||
ErrLegacySSEServer = fmt.Errorf("server returned 4xx for initialize POST, likely a legacy SSE server")
|
||||
|
||||
retryInterval = 1 * time.Second // a variable is convenient for testing
|
||||
)
|
||||
|
||||
+3
@@ -6,4 +6,7 @@ const (
|
||||
ContentTypeAudio = "audio"
|
||||
ContentTypeLink = "resource_link"
|
||||
ContentTypeResource = "resource"
|
||||
|
||||
ElicitationModeForm = "form"
|
||||
ElicitationModeURL = "url"
|
||||
)
|
||||
+47
@@ -1,6 +1,7 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
@@ -29,6 +30,28 @@ var (
|
||||
ErrResourceNotFound = errors.New("resource not found")
|
||||
)
|
||||
|
||||
// URLElicitationRequiredError is returned when the server requires URL elicitation to proceed.
|
||||
type URLElicitationRequiredError struct {
|
||||
Elicitations []ElicitationParams `json:"elicitations"`
|
||||
}
|
||||
|
||||
func (e URLElicitationRequiredError) Error() string {
|
||||
return fmt.Sprintf("URL elicitation required: %d elicitation(s) needed", len(e.Elicitations))
|
||||
}
|
||||
|
||||
func (e URLElicitationRequiredError) JSONRPCError() JSONRPCError {
|
||||
return JSONRPCError{
|
||||
JSONRPC: JSONRPC_VERSION,
|
||||
Error: JSONRPCErrorDetails{
|
||||
Code: URL_ELICITATION_REQUIRED,
|
||||
Message: e.Error(),
|
||||
Data: map[string]any{
|
||||
"elicitations": e.Elicitations,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// UnsupportedProtocolVersionError is returned when the server responds with
|
||||
// a protocol version that the client doesn't support.
|
||||
type UnsupportedProtocolVersionError struct {
|
||||
@@ -39,6 +62,12 @@ func (e UnsupportedProtocolVersionError) Error() string {
|
||||
return fmt.Sprintf("unsupported protocol version: %q", e.Version)
|
||||
}
|
||||
|
||||
// Is implements the errors.Is interface for better error handling
|
||||
func (e URLElicitationRequiredError) Is(target error) bool {
|
||||
_, ok := target.(URLElicitationRequiredError)
|
||||
return ok
|
||||
}
|
||||
|
||||
// Is implements the errors.Is interface for better error handling
|
||||
func (e UnsupportedProtocolVersionError) Is(target error) bool {
|
||||
_, ok := target.(UnsupportedProtocolVersionError)
|
||||
@@ -72,6 +101,24 @@ func (e *JSONRPCErrorDetails) AsError() error {
|
||||
err = ErrRequestInterrupted
|
||||
case RESOURCE_NOT_FOUND:
|
||||
err = ErrResourceNotFound
|
||||
case URL_ELICITATION_REQUIRED:
|
||||
// Attempt to reconstruct URLElicitationRequiredError from Data
|
||||
if e.Data != nil {
|
||||
// Round-trip through JSON to parse into struct
|
||||
// This handles both map[string]any (from unmarshal) and other forms
|
||||
if dataBytes, marshalErr := json.Marshal(e.Data); marshalErr == nil {
|
||||
var data struct {
|
||||
Elicitations []ElicitationParams `json:"elicitations"`
|
||||
}
|
||||
if unmarshalErr := json.Unmarshal(dataBytes, &data); unmarshalErr == nil {
|
||||
return URLElicitationRequiredError{
|
||||
Elicitations: data.Elicitations,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Fallback if data is missing or invalid
|
||||
return URLElicitationRequiredError{}
|
||||
default:
|
||||
return errors.New(e.Message)
|
||||
}
|
||||
|
||||
+10
@@ -56,6 +56,8 @@ type Prompt struct {
|
||||
// A list of arguments to use for templating the prompt.
|
||||
// The presence of arguments indicates this is a template prompt.
|
||||
Arguments []PromptArgument `json:"arguments,omitempty"`
|
||||
// Icons provides visual identifiers for the prompt
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the name of the prompt.
|
||||
@@ -136,6 +138,14 @@ func WithPromptDescription(description string) PromptOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithPromptIcons adds icons to the Prompt.
|
||||
// Icons provide visual identifiers for the prompt.
|
||||
func WithPromptIcons(icons ...Icon) PromptOption {
|
||||
return func(p *Prompt) {
|
||||
p.Icons = icons
|
||||
}
|
||||
}
|
||||
|
||||
// WithArgument adds an argument to the prompt's argument list.
|
||||
// The argument will be configured based on the provided options.
|
||||
func WithArgument(name string, opts ...ArgumentOption) PromptOption {
|
||||
|
||||
+56
-9
@@ -1,6 +1,10 @@
|
||||
package mcp
|
||||
|
||||
import "github.com/yosida95/uritemplate/v3"
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/yosida95/uritemplate/v3"
|
||||
)
|
||||
|
||||
// ResourceOption is a function that configures a Resource.
|
||||
// It provides a flexible way to set various properties of a Resource using the functional options pattern.
|
||||
@@ -38,15 +42,29 @@ func WithMIMEType(mimeType string) ResourceOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithAnnotations adds annotations to the Resource.
|
||||
// Annotations can provide additional metadata about the resource's intended use.
|
||||
func WithAnnotations(audience []Role, priority float64) ResourceOption {
|
||||
// WithAnnotations returns a ResourceOption that sets the resource's Annotations fields.
|
||||
// It initializes Annotations if nil, sets Audience to the provided slice,
|
||||
// stores Priority as a pointer to the provided value, and sets LastModified to the provided timestamp.
|
||||
func WithAnnotations(audience []Role, priority float64, lastModified string) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
if r.Annotations == nil {
|
||||
r.Annotations = &Annotations{}
|
||||
}
|
||||
r.Annotations.Audience = audience
|
||||
r.Annotations.Priority = priority
|
||||
r.Annotations.Priority = &priority
|
||||
r.Annotations.LastModified = lastModified
|
||||
}
|
||||
}
|
||||
|
||||
// WithLastModified returns a ResourceOption that sets the resource's Annotations.LastModified
|
||||
// to the provided timestamp. If the resource's Annotations is nil, it will be initialized.
|
||||
// The timestamp is expected to be an ISO 8601 (RFC3339) formatted string (e.g., "2025-01-12T15:00:58Z").
|
||||
func WithLastModified(timestamp string) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
if r.Annotations == nil {
|
||||
r.Annotations = &Annotations{}
|
||||
}
|
||||
r.Annotations.LastModified = timestamp
|
||||
}
|
||||
}
|
||||
|
||||
@@ -86,14 +104,43 @@ func WithTemplateMIMEType(mimeType string) ResourceTemplateOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithTemplateAnnotations adds annotations to the ResourceTemplate.
|
||||
// Annotations can provide additional metadata about the template's intended use.
|
||||
func WithTemplateAnnotations(audience []Role, priority float64) ResourceTemplateOption {
|
||||
// WithTemplateAnnotations returns a ResourceTemplateOption that sets the template's
|
||||
// Annotations field, initializing it if nil, and setting Audience, Priority, and LastModified.
|
||||
func WithTemplateAnnotations(audience []Role, priority float64, lastModified string) ResourceTemplateOption {
|
||||
return func(t *ResourceTemplate) {
|
||||
if t.Annotations == nil {
|
||||
t.Annotations = &Annotations{}
|
||||
}
|
||||
t.Annotations.Audience = audience
|
||||
t.Annotations.Priority = priority
|
||||
t.Annotations.Priority = &priority
|
||||
t.Annotations.LastModified = lastModified
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateISO8601Timestamp verifies that timestamp is a valid ISO 8601 timestamp
|
||||
// using the RFC3339 layout. An empty string is considered valid. It returns nil
|
||||
// when the timestamp is valid, or the parsing error when it is not.
|
||||
func ValidateISO8601Timestamp(timestamp string) error {
|
||||
if timestamp == "" {
|
||||
return nil // Empty is valid (optional field)
|
||||
}
|
||||
// Use time.RFC3339 for ISO 8601 compatibility
|
||||
_, err := time.Parse(time.RFC3339, timestamp)
|
||||
return err
|
||||
}
|
||||
|
||||
// WithResourceIcons adds icons to the Resource.
|
||||
// Icons provide visual identifiers for the resource.
|
||||
func WithResourceIcons(icons ...Icon) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
r.Icons = icons
|
||||
}
|
||||
}
|
||||
|
||||
// WithTemplateIcons adds icons to the ResourceTemplate.
|
||||
// Icons provide visual identifiers for the resource template.
|
||||
func WithTemplateIcons(icons ...Icon) ResourceTemplateOption {
|
||||
return func(rt *ResourceTemplate) {
|
||||
rt.Icons = icons
|
||||
}
|
||||
}
|
||||
+208
@@ -0,0 +1,208 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// TaskOption is a function that configures a Task.
|
||||
// It provides a flexible way to set various properties of a Task using the functional options pattern.
|
||||
type TaskOption func(*Task)
|
||||
|
||||
//
|
||||
// Core Task Functions
|
||||
//
|
||||
|
||||
// NewTask creates a new Task with the given ID and options.
|
||||
// The task will be configured based on the provided options.
|
||||
// Options are applied in order, allowing for flexible task configuration.
|
||||
func NewTask(taskId string, opts ...TaskOption) Task {
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
task := Task{
|
||||
TaskId: taskId,
|
||||
Status: TaskStatusWorking,
|
||||
CreatedAt: now,
|
||||
LastUpdatedAt: now,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&task)
|
||||
}
|
||||
|
||||
return task
|
||||
}
|
||||
|
||||
// WithTaskStatus sets the status of the task.
|
||||
func WithTaskStatus(status TaskStatus) TaskOption {
|
||||
return func(t *Task) {
|
||||
t.Status = status
|
||||
}
|
||||
}
|
||||
|
||||
// WithTaskStatusMessage sets a human-readable status message for the task.
|
||||
func WithTaskStatusMessage(message string) TaskOption {
|
||||
return func(t *Task) {
|
||||
t.StatusMessage = message
|
||||
}
|
||||
}
|
||||
|
||||
// WithTaskTTL sets the time-to-live for the task in milliseconds.
|
||||
// After this duration from creation, the task may be deleted.
|
||||
func WithTaskTTL(ttlMs int64) TaskOption {
|
||||
return func(t *Task) {
|
||||
t.TTL = &ttlMs
|
||||
}
|
||||
}
|
||||
|
||||
// WithTaskPollInterval sets the suggested polling interval in milliseconds.
|
||||
func WithTaskPollInterval(intervalMs int64) TaskOption {
|
||||
return func(t *Task) {
|
||||
t.PollInterval = &intervalMs
|
||||
}
|
||||
}
|
||||
|
||||
// WithTaskCreatedAt sets a specific creation timestamp for the task.
|
||||
// By default, NewTask uses the current time.
|
||||
func WithTaskCreatedAt(createdAt string) TaskOption {
|
||||
return func(t *Task) {
|
||||
t.CreatedAt = createdAt
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Task Helper Functions
|
||||
//
|
||||
|
||||
// NewTaskParams creates TaskParams with the given TTL.
|
||||
func NewTaskParams(ttlMs *int64) TaskParams {
|
||||
return TaskParams{
|
||||
TTL: ttlMs,
|
||||
}
|
||||
}
|
||||
|
||||
// NewCreateTaskResult creates a CreateTaskResult with the given task.
|
||||
func NewCreateTaskResult(task Task) CreateTaskResult {
|
||||
return CreateTaskResult{
|
||||
Task: task,
|
||||
}
|
||||
}
|
||||
|
||||
// NewGetTaskResult creates a GetTaskResult from a Task.
|
||||
func NewGetTaskResult(task Task) GetTaskResult {
|
||||
return GetTaskResult{
|
||||
Task: task,
|
||||
}
|
||||
}
|
||||
|
||||
// NewListTasksResult creates a ListTasksResult with the given tasks.
|
||||
func NewListTasksResult(tasks []Task) ListTasksResult {
|
||||
return ListTasksResult{
|
||||
Tasks: tasks,
|
||||
}
|
||||
}
|
||||
|
||||
// NewCancelTaskResult creates a CancelTaskResult from a Task.
|
||||
func NewCancelTaskResult(task Task) CancelTaskResult {
|
||||
return CancelTaskResult{
|
||||
Task: task,
|
||||
}
|
||||
}
|
||||
|
||||
// NewTaskStatusNotification creates a notification for a task status change.
|
||||
func NewTaskStatusNotification(task Task) TaskStatusNotification {
|
||||
return TaskStatusNotification{
|
||||
Notification: Notification{
|
||||
Method: string(MethodNotificationTasksStatus),
|
||||
},
|
||||
Params: TaskStatusNotificationParams{
|
||||
Task: task,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Task Capability Helper Functions
|
||||
//
|
||||
|
||||
// NewTasksCapability creates a TasksCapability with all operations enabled.
|
||||
func NewTasksCapability() *TasksCapability {
|
||||
return &TasksCapability{
|
||||
List: &struct{}{},
|
||||
Cancel: &struct{}{},
|
||||
Requests: &TaskRequestsCapability{
|
||||
Tools: &struct {
|
||||
Call *struct{} `json:"call,omitempty"`
|
||||
}{
|
||||
Call: &struct{}{},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewTasksCapabilityWithToolsOnly creates a TasksCapability with only tool call support.
|
||||
// List and Cancel operations are not enabled with this capability.
|
||||
func NewTasksCapabilityWithToolsOnly() *TasksCapability {
|
||||
return &TasksCapability{
|
||||
Requests: &TaskRequestsCapability{
|
||||
Tools: &struct {
|
||||
Call *struct{} `json:"call,omitempty"`
|
||||
}{
|
||||
Call: &struct{}{},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Related Task Metadata Functions
|
||||
//
|
||||
|
||||
// RelatedTaskMetaKey is the metadata key for associating a message with a task.
|
||||
const RelatedTaskMetaKey = "io.modelcontextprotocol/related-task"
|
||||
|
||||
// RelatedTaskMeta creates the metadata for associating a message with a task.
|
||||
// The returned map contains a "taskId" field with the provided task ID.
|
||||
func RelatedTaskMeta(taskID string) map[string]any {
|
||||
return map[string]any{
|
||||
"taskId": taskID,
|
||||
}
|
||||
}
|
||||
|
||||
// WithRelatedTask returns a Meta with the related task ID set.
|
||||
// This is useful for associating task results with their originating task.
|
||||
func WithRelatedTask(taskID string) *Meta {
|
||||
return &Meta{
|
||||
AdditionalFields: map[string]any{
|
||||
RelatedTaskMetaKey: RelatedTaskMeta(taskID),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Model Immediate Response Metadata Functions
|
||||
//
|
||||
|
||||
// ModelImmediateResponseMetaKey is the metadata key for providing an immediate response to the model.
|
||||
// Servers can use this optional key in the _meta field of CreateTaskResult to provide
|
||||
// a string that should be passed as an immediate tool result to the model while the task
|
||||
// continues executing asynchronously in the background.
|
||||
const ModelImmediateResponseMetaKey = "io.modelcontextprotocol/model-immediate-response"
|
||||
|
||||
// WithModelImmediateResponse creates Meta with an immediate response message for the model.
|
||||
// This allows the model to continue processing while the task executes asynchronously.
|
||||
// The message parameter is a human-readable string that will be shown to the model.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// return &mcp.CreateTaskResult{
|
||||
// Task: task,
|
||||
// Result: mcp.Result{
|
||||
// Meta: mcp.WithModelImmediateResponse("Processing your request. This may take a few minutes."),
|
||||
// },
|
||||
// }
|
||||
func WithModelImmediateResponse(message string) *Meta {
|
||||
return &Meta{
|
||||
AdditionalFields: map[string]any{
|
||||
ModelImmediateResponseMetaKey: message,
|
||||
},
|
||||
}
|
||||
}
|
||||
+122
-9
@@ -58,9 +58,10 @@ type CallToolRequest struct {
|
||||
}
|
||||
|
||||
type CallToolParams struct {
|
||||
Name string `json:"name"`
|
||||
Arguments any `json:"arguments,omitempty"`
|
||||
Meta *Meta `json:"_meta,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Arguments any `json:"arguments,omitempty"`
|
||||
Meta *Meta `json:"_meta,omitempty"`
|
||||
Task *TaskParams `json:"task,omitempty"`
|
||||
}
|
||||
|
||||
// GetArguments returns the Arguments as map[string]any for backward compatibility
|
||||
@@ -553,6 +554,24 @@ type ToolListChangedNotification struct {
|
||||
Notification
|
||||
}
|
||||
|
||||
// TaskSupport indicates how a tool supports task augmentation.
|
||||
type TaskSupport string
|
||||
|
||||
const (
|
||||
// TaskSupportForbidden means the tool cannot be invoked as a task (default).
|
||||
TaskSupportForbidden TaskSupport = "forbidden"
|
||||
// TaskSupportOptional means the tool can be invoked as a task or normally.
|
||||
TaskSupportOptional TaskSupport = "optional"
|
||||
// TaskSupportRequired means the tool must be invoked as a task.
|
||||
TaskSupportRequired TaskSupport = "required"
|
||||
)
|
||||
|
||||
// ToolExecution describes execution behavior for a tool.
|
||||
type ToolExecution struct {
|
||||
// TaskSupport indicates whether the tool supports task augmentation.
|
||||
TaskSupport TaskSupport `json:"taskSupport,omitempty"`
|
||||
}
|
||||
|
||||
// Tool represents the definition for a tool the client can call.
|
||||
type Tool struct {
|
||||
// Meta is a metadata object that is reserved by MCP for storing additional information.
|
||||
@@ -571,6 +590,12 @@ type Tool struct {
|
||||
RawOutputSchema json.RawMessage `json:"-"` // Hide this from JSON marshaling
|
||||
// Optional properties describing tool behavior
|
||||
Annotations ToolAnnotation `json:"annotations"`
|
||||
// Support for deferred loading
|
||||
DeferLoading bool `json:"defer_loading,omitempty"`
|
||||
// Icons provides visual identifiers for the tool
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
// Execution describes execution behavior for the tool
|
||||
Execution *ToolExecution `json:"execution,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the name of the tool.
|
||||
@@ -613,27 +638,70 @@ func (t Tool) MarshalJSON() ([]byte, error) {
|
||||
|
||||
m["annotations"] = t.Annotations
|
||||
|
||||
if t.DeferLoading {
|
||||
m["defer_loading"] = t.DeferLoading
|
||||
}
|
||||
|
||||
// Marshal Meta if present
|
||||
if t.Meta != nil {
|
||||
m["_meta"] = t.Meta
|
||||
}
|
||||
|
||||
if t.Icons != nil {
|
||||
m["icons"] = t.Icons
|
||||
}
|
||||
|
||||
if t.Execution != nil {
|
||||
m["execution"] = t.Execution
|
||||
}
|
||||
|
||||
return json.Marshal(m)
|
||||
}
|
||||
|
||||
// ToolArgumentsSchema represents a JSON Schema for tool arguments.
|
||||
type ToolArgumentsSchema struct {
|
||||
Defs map[string]any `json:"$defs,omitempty"`
|
||||
Type string `json:"type"`
|
||||
Properties map[string]any `json:"properties,omitempty"`
|
||||
Required []string `json:"required,omitempty"`
|
||||
Defs map[string]any `json:"$defs,omitempty"`
|
||||
Type string `json:"type"`
|
||||
Properties map[string]any `json:"properties"`
|
||||
Required []string `json:"required,omitempty"`
|
||||
AdditionalProperties any `json:"additionalProperties,omitempty"`
|
||||
}
|
||||
|
||||
type ToolInputSchema ToolArgumentsSchema // For retro-compatibility
|
||||
type ToolOutputSchema ToolArgumentsSchema
|
||||
|
||||
// MarshalJSON implements the json.Marshaler interface for ToolInputSchema.
|
||||
func (tis ToolInputSchema) MarshalJSON() ([]byte, error) {
|
||||
return toolArgumentsSchemaMarshalJSON(ToolArgumentsSchema(tis))
|
||||
}
|
||||
|
||||
// MarshalJSON implements the json.Marshaler interface for ToolOutputSchema.
|
||||
func (tis ToolOutputSchema) MarshalJSON() ([]byte, error) {
|
||||
return toolArgumentsSchemaMarshalJSON(ToolArgumentsSchema(tis))
|
||||
}
|
||||
|
||||
// MarshalJSON implements the json.Marshaler interface for ToolArgumentsSchema.
|
||||
func (tis ToolArgumentsSchema) MarshalJSON() ([]byte, error) {
|
||||
return toolArgumentsSchemaMarshalJSON(tis)
|
||||
}
|
||||
|
||||
// UnmarshalJSON implements the json.Unmarshaler interface for ToolInputSchema.
|
||||
func (tis *ToolInputSchema) UnmarshalJSON(data []byte) error {
|
||||
return toolArgumentsSchemaUnmarshalJSON(data, (*ToolArgumentsSchema)(tis))
|
||||
}
|
||||
|
||||
// UnmarshalJSON implements the json.Unmarshaler interface for ToolOutputSchema.
|
||||
func (tis *ToolOutputSchema) UnmarshalJSON(data []byte) error {
|
||||
return toolArgumentsSchemaUnmarshalJSON(data, (*ToolArgumentsSchema)(tis))
|
||||
}
|
||||
|
||||
// UnmarshalJSON implements the json.Unmarshaler interface for ToolArgumentsSchema.
|
||||
func (tis *ToolArgumentsSchema) UnmarshalJSON(data []byte) error {
|
||||
return toolArgumentsSchemaUnmarshalJSON(data, tis)
|
||||
}
|
||||
|
||||
// toolArgumentsSchemaMarshalJSON handles the fields stored in ToolArgumentsSchema when json.Marshaler is called
|
||||
func toolArgumentsSchemaMarshalJSON(tis ToolArgumentsSchema) ([]byte, error) {
|
||||
m := make(map[string]any)
|
||||
m["type"] = tis.Type
|
||||
|
||||
@@ -644,19 +712,27 @@ func (tis ToolArgumentsSchema) MarshalJSON() ([]byte, error) {
|
||||
// Marshal Properties to '{}' rather than `nil` when its length equals zero
|
||||
if tis.Properties != nil {
|
||||
m["properties"] = tis.Properties
|
||||
} else {
|
||||
m["properties"] = map[string]any{}
|
||||
}
|
||||
|
||||
// Marshal Required to '[]' rather than `nil` when its length equals zero
|
||||
if len(tis.Required) > 0 {
|
||||
m["required"] = tis.Required
|
||||
} else {
|
||||
m["required"] = []string{}
|
||||
}
|
||||
|
||||
if tis.AdditionalProperties != nil {
|
||||
m["additionalProperties"] = tis.AdditionalProperties
|
||||
}
|
||||
|
||||
return json.Marshal(m)
|
||||
}
|
||||
|
||||
// UnmarshalJSON implements the json.Unmarshaler interface for ToolArgumentsSchema.
|
||||
// It handles both "$defs" (JSON Schema 2019-09+) and "definitions" (JSON Schema draft-07)
|
||||
// by reading either field and storing it in the Defs field.
|
||||
func (tis *ToolArgumentsSchema) UnmarshalJSON(data []byte) error {
|
||||
func toolArgumentsSchemaUnmarshalJSON(data []byte, tis *ToolArgumentsSchema) error {
|
||||
// Use a temporary type to avoid infinite recursion
|
||||
type Alias ToolArgumentsSchema
|
||||
aux := &struct {
|
||||
@@ -755,6 +831,14 @@ func WithDescription(description string) ToolOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithDeferLoading sets the defer_loading flag for the tool.
|
||||
// This is used to implement dynamic tool loading/searching patterns.
|
||||
func WithDeferLoading(deferLoading bool) ToolOption {
|
||||
return func(t *Tool) {
|
||||
t.DeferLoading = deferLoading
|
||||
}
|
||||
}
|
||||
|
||||
// WithInputSchema creates a ToolOption that sets the input schema for a tool.
|
||||
// It accepts any Go type, usually a struct, and automatically generates a JSON schema from it.
|
||||
func WithInputSchema[T any]() ToolOption {
|
||||
@@ -785,6 +869,26 @@ func WithInputSchema[T any]() ToolOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithToolIcons adds icons to the Tool.
|
||||
// Icons provide visual identifiers for the tool.
|
||||
func WithToolIcons(icons ...Icon) ToolOption {
|
||||
return func(t *Tool) {
|
||||
t.Icons = icons
|
||||
}
|
||||
}
|
||||
|
||||
// WithTaskSupport sets the task support mode for the tool.
|
||||
// It configures whether the tool can be invoked as a task (asynchronously).
|
||||
// Valid values are TaskSupportForbidden (default), TaskSupportOptional, or TaskSupportRequired.
|
||||
func WithTaskSupport(support TaskSupport) ToolOption {
|
||||
return func(t *Tool) {
|
||||
if t.Execution == nil {
|
||||
t.Execution = &ToolExecution{}
|
||||
}
|
||||
t.Execution.TaskSupport = support
|
||||
}
|
||||
}
|
||||
|
||||
// WithRawInputSchema sets a raw JSON schema for the tool's input.
|
||||
// Use this when you need full control over the schema or when working with
|
||||
// complex schemas that can't be generated from Go types. The jsonschema library
|
||||
@@ -891,6 +995,15 @@ func WithOpenWorldHintAnnotation(value bool) ToolOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithSchemaAdditionalProperties sets the additionalProperties field on the tool's input schema.
|
||||
// It accepts false (disallow extra properties), true (allow any), or a schema map
|
||||
// to validate additional properties against.
|
||||
func WithSchemaAdditionalProperties(schema any) ToolOption {
|
||||
return func(t *Tool) {
|
||||
t.InputSchema.AdditionalProperties = schema
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Common Property Options
|
||||
//
|
||||
|
||||
+375
-26
@@ -59,10 +59,29 @@ const (
|
||||
// https://modelcontextprotocol.io/docs/concepts/elicitation
|
||||
MethodElicitationCreate MCPMethod = "elicitation/create"
|
||||
|
||||
// MethodNotificationElicitationComplete notifies when a URL mode elicitation completes.
|
||||
MethodNotificationElicitationComplete MCPMethod = "notifications/elicitation/complete"
|
||||
|
||||
// MethodListRoots requests roots list from the client during interactions.
|
||||
// https://modelcontextprotocol.io/specification/2025-06-18/client/roots
|
||||
MethodListRoots MCPMethod = "roots/list"
|
||||
|
||||
// MethodTasksGet retrieves the current status of a task.
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks
|
||||
MethodTasksGet MCPMethod = "tasks/get"
|
||||
|
||||
// MethodTasksList lists all tasks for the current session.
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks
|
||||
MethodTasksList MCPMethod = "tasks/list"
|
||||
|
||||
// MethodTasksResult retrieves the result of a completed task.
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks
|
||||
MethodTasksResult MCPMethod = "tasks/result"
|
||||
|
||||
// MethodTasksCancel cancels an in-progress task.
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks
|
||||
MethodTasksCancel MCPMethod = "tasks/cancel"
|
||||
|
||||
// MethodNotificationResourcesListChanged notifies when the list of available resources changes.
|
||||
// https://modelcontextprotocol.io/specification/2025-03-26/server/resources#list-changed-notification
|
||||
MethodNotificationResourcesListChanged = "notifications/resources/list_changed"
|
||||
@@ -80,6 +99,14 @@ const (
|
||||
// MethodNotificationRootsListChanged notifies when the list of available roots changes.
|
||||
// https://modelcontextprotocol.io/specification/2025-06-18/client/roots#root-list-changes
|
||||
MethodNotificationRootsListChanged = "notifications/roots/list_changed"
|
||||
|
||||
// MethodNotificationTasksStatus notifies when a task's status changes.
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks
|
||||
MethodNotificationTasksStatus = "notifications/tasks/status"
|
||||
|
||||
// MethodCompletionComplete returns completion suggestions for a given argument
|
||||
// https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion
|
||||
MethodCompletionComplete MCPMethod = "completion/complete"
|
||||
)
|
||||
|
||||
type URITemplate struct {
|
||||
@@ -109,11 +136,12 @@ func (t *URITemplate) UnmarshalJSON(data []byte) error {
|
||||
type JSONRPCMessage any
|
||||
|
||||
// LATEST_PROTOCOL_VERSION is the most recent version of the MCP protocol.
|
||||
const LATEST_PROTOCOL_VERSION = "2025-06-18"
|
||||
const LATEST_PROTOCOL_VERSION = "2025-11-25"
|
||||
|
||||
// ValidProtocolVersions lists all known valid MCP protocol versions.
|
||||
var ValidProtocolVersions = []string{
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
"2025-06-18",
|
||||
"2025-03-26",
|
||||
"2024-11-05",
|
||||
}
|
||||
@@ -395,8 +423,11 @@ const (
|
||||
|
||||
// MCP error codes
|
||||
const (
|
||||
// RESOURCE_NOT_FOUND indicates a requested resource was not found.
|
||||
// RESOURCE_NOT_FOUND indicates that the requested resource was not found.
|
||||
RESOURCE_NOT_FOUND = -32002
|
||||
|
||||
// URL_ELICITATION_REQUIRED is the error code for when URL elicitation is required.
|
||||
URL_ELICITATION_REQUIRED = -32042
|
||||
)
|
||||
|
||||
/* Empty result */
|
||||
@@ -490,7 +521,9 @@ type ClientCapabilities struct {
|
||||
// Present if the client supports sampling from an LLM.
|
||||
Sampling *struct{} `json:"sampling,omitempty"`
|
||||
// Present if the client supports elicitation requests from the server.
|
||||
Elicitation *struct{} `json:"elicitation,omitempty"`
|
||||
Elicitation *ElicitationCapability `json:"elicitation,omitempty"`
|
||||
// Present if the client supports task-based execution.
|
||||
Tasks *TasksCapability `json:"tasks,omitempty"`
|
||||
}
|
||||
|
||||
// ServerCapabilities represents capabilities that a server may support. Known
|
||||
@@ -522,16 +555,43 @@ type ServerCapabilities struct {
|
||||
ListChanged bool `json:"listChanged,omitempty"`
|
||||
} `json:"tools,omitempty"`
|
||||
// Present if the server supports elicitation requests to the client.
|
||||
Elicitation *struct{} `json:"elicitation,omitempty"`
|
||||
Elicitation *ElicitationCapability `json:"elicitation,omitempty"`
|
||||
// Present if the server supports roots requests to the client.
|
||||
Roots *struct{} `json:"roots,omitempty"`
|
||||
// Present if the server supports task-based execution.
|
||||
Tasks *TasksCapability `json:"tasks,omitempty"`
|
||||
// Present if the server supports completions requests to the client.
|
||||
Completions *struct{} `json:"completions,omitempty"`
|
||||
}
|
||||
|
||||
// Icon represents a visual identifier for MCP entities.
|
||||
//
|
||||
// Security considerations:
|
||||
// - Clients MUST support at least image/png and image/jpeg MIME types
|
||||
// - Clients SHOULD support image/svg+xml and image/webp
|
||||
// - Icons should be treated as untrusted input
|
||||
// - URI scheme validation (HTTPS or data URI only)
|
||||
// - Size/dimension limits to prevent resource exhaustion
|
||||
type Icon struct {
|
||||
// URI pointing to the icon resource (HTTPS URL or data URI)
|
||||
Src string `json:"src"`
|
||||
|
||||
// Optional MIME type (e.g., "image/png", "image/svg+xml")
|
||||
MIMEType string `json:"mimeType,omitempty"`
|
||||
|
||||
// Optional size specifications (e.g., ["48x48"], ["any"] for SVG)
|
||||
Sizes []string `json:"sizes,omitempty"`
|
||||
}
|
||||
|
||||
// Implementation describes the name and version of an MCP implementation.
|
||||
type Implementation struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
WebsiteURL string `json:"websiteUrl,omitempty"`
|
||||
// Icons provides visual identifiers for the implementation
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
}
|
||||
|
||||
/* Ping */
|
||||
@@ -708,6 +768,8 @@ type Resource struct {
|
||||
Description string `json:"description,omitempty"`
|
||||
// The MIME type of this resource, if known.
|
||||
MIMEType string `json:"mimeType,omitempty"`
|
||||
// Icons provides visual identifiers for the resource
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the name of the resource.
|
||||
@@ -736,6 +798,8 @@ type ResourceTemplate struct {
|
||||
// The MIME type for all resources that match this template. This should only
|
||||
// be included if all resources matching this template have the same type.
|
||||
MIMEType string `json:"mimeType,omitempty"`
|
||||
// Icons provides visual identifiers for the resource template
|
||||
Icons []Icon `json:"icons,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the name of the resourceTemplate.
|
||||
@@ -861,10 +925,49 @@ type ElicitationRequest struct {
|
||||
|
||||
// ElicitationParams contains the parameters for an elicitation request.
|
||||
type ElicitationParams struct {
|
||||
Meta *Meta `json:"_meta,omitempty"`
|
||||
// Mode specifies the type of elicitation: "form" or "url". Defaults to "form".
|
||||
Mode string `json:"mode,omitempty"`
|
||||
// A human-readable message explaining what information is being requested and why.
|
||||
Message string `json:"message"`
|
||||
|
||||
// Form mode fields
|
||||
|
||||
// A JSON Schema defining the expected structure of the user's response.
|
||||
RequestedSchema any `json:"requestedSchema"`
|
||||
RequestedSchema any `json:"requestedSchema,omitempty"`
|
||||
|
||||
// URL mode fields
|
||||
|
||||
// ElicitationID is a unique identifier for the elicitation request.
|
||||
ElicitationID string `json:"elicitationId,omitempty"`
|
||||
// URL is the URL to be opened by the user.
|
||||
URL string `json:"url,omitempty"`
|
||||
}
|
||||
|
||||
// Validate checks if the elicitation parameters are valid.
|
||||
func (p ElicitationParams) Validate() error {
|
||||
mode := p.Mode
|
||||
if mode == "" {
|
||||
mode = ElicitationModeForm
|
||||
}
|
||||
|
||||
switch mode {
|
||||
case ElicitationModeForm:
|
||||
if p.RequestedSchema == nil {
|
||||
return fmt.Errorf("requestedSchema is required for form elicitation")
|
||||
}
|
||||
case ElicitationModeURL:
|
||||
if p.ElicitationID == "" {
|
||||
return fmt.Errorf("elicitationId is required for url elicitation")
|
||||
}
|
||||
if p.URL == "" {
|
||||
return fmt.Errorf("url is required for url elicitation")
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("invalid elicitation mode: %s", mode)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ElicitationResult represents the result of an elicitation request.
|
||||
@@ -952,7 +1055,10 @@ type Annotations struct {
|
||||
// A value of 1 means "most important," and indicates that the data is
|
||||
// effectively required, while 0 means "least important," and indicates that
|
||||
// the data is entirely optional.
|
||||
Priority float64 `json:"priority,omitempty"`
|
||||
// Priority ranges from 0.0 to 1.0 (1 = most important, 0 = least important).
|
||||
Priority *float64 `json:"priority,omitempty"`
|
||||
// ISO 8601 formatted timestamp (e.g., "2025-01-12T15:00:58Z")
|
||||
LastModified string `json:"lastModified,omitempty"`
|
||||
}
|
||||
|
||||
// Annotated is the base for objects that include optional annotations for the
|
||||
@@ -1104,29 +1210,81 @@ type CompleteRequest struct {
|
||||
Header http.Header `json:"-"`
|
||||
}
|
||||
|
||||
// CompleteParams are the parameters for a completion/complete request
|
||||
type CompleteParams struct {
|
||||
Ref any `json:"ref"` // Can be PromptReference or ResourceReference
|
||||
Argument struct {
|
||||
// The name of the argument
|
||||
Name string `json:"name"`
|
||||
// The value of the argument to use for completion matching.
|
||||
Value string `json:"value"`
|
||||
} `json:"argument"`
|
||||
Ref any `json:"ref"` // Can be PromptReference or ResourceReference
|
||||
Argument CompleteArgument `json:"argument"`
|
||||
Context CompleteContext `json:"context"`
|
||||
}
|
||||
|
||||
func (p *CompleteParams) UnmarshalJSON(data []byte) error {
|
||||
// Use a temporary type to avoid infinite recursion on UnmarshalJSON
|
||||
type Alias CompleteParams
|
||||
aux := &struct {
|
||||
// Use RawMessage to delay unmarshalling until after the type is known
|
||||
Ref json.RawMessage `json:"ref"`
|
||||
*Alias
|
||||
}{
|
||||
Alias: (*Alias)(p),
|
||||
}
|
||||
if err := json.Unmarshal(data, aux); err != nil {
|
||||
return err
|
||||
}
|
||||
// Use a temporary "type peek" struct to determine the type
|
||||
var typePeek struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if err := json.Unmarshal(aux.Ref, &typePeek); err != nil {
|
||||
return err
|
||||
}
|
||||
switch typePeek.Type {
|
||||
case "ref/prompt":
|
||||
var prompt PromptReference
|
||||
if err := json.Unmarshal(aux.Ref, &prompt); err != nil {
|
||||
return err
|
||||
}
|
||||
p.Ref = prompt
|
||||
case "ref/resource":
|
||||
var resource ResourceReference
|
||||
if err := json.Unmarshal(aux.Ref, &resource); err != nil {
|
||||
return err
|
||||
}
|
||||
p.Ref = resource
|
||||
default:
|
||||
return fmt.Errorf("unknown reference type: %s", typePeek.Type)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CompleteResult is the server's response to a completion/complete request
|
||||
type CompleteResult struct {
|
||||
Result
|
||||
Completion struct {
|
||||
// An array of completion values. Must not exceed 100 items.
|
||||
Values []string `json:"values"`
|
||||
// The total number of completion options available. This can exceed the
|
||||
// number of values actually sent in the response.
|
||||
Total int `json:"total,omitempty"`
|
||||
// Indicates whether there are additional completion options beyond those
|
||||
// provided in the current response, even if the exact total is unknown.
|
||||
HasMore bool `json:"hasMore,omitempty"`
|
||||
} `json:"completion"`
|
||||
Completion Completion `json:"completion"`
|
||||
}
|
||||
|
||||
// CompleteArgument is an argument to a completion request
|
||||
type CompleteArgument struct {
|
||||
// The name of the argument
|
||||
Name string `json:"name"`
|
||||
// The value of the argument to use for completion matching.
|
||||
Value string `json:"value"`
|
||||
}
|
||||
|
||||
// CompleteContext is the context about already-resolved arguments
|
||||
type CompleteContext struct {
|
||||
Arguments map[string]string `json:"arguments"`
|
||||
}
|
||||
|
||||
// Completion is the server's response to a completion/complete request
|
||||
type Completion struct {
|
||||
// An array of completion values. Must not exceed 100 items.
|
||||
Values []string `json:"values"`
|
||||
// The total number of completion options available. This can exceed the
|
||||
// number of values actually sent in the response.
|
||||
Total int `json:"total,omitempty"`
|
||||
// Indicates whether there are additional completion options beyond those
|
||||
// provided in the current response, even if the exact total is unknown.
|
||||
HasMore bool `json:"hasMore,omitempty"`
|
||||
}
|
||||
|
||||
// ResourceReference is a reference to a resource or resource template definition.
|
||||
@@ -1186,6 +1344,176 @@ type RootsListChangedNotification struct {
|
||||
Notification
|
||||
}
|
||||
|
||||
/* Tasks */
|
||||
|
||||
// TasksCapability represents the task capabilities that a client or server may support.
|
||||
// Tasks enable long-running, asynchronous operations with status polling.
|
||||
type TasksCapability struct {
|
||||
// Whether the party supports the tasks/list operation.
|
||||
List *struct{} `json:"list,omitempty"`
|
||||
// Whether the party supports the tasks/cancel operation.
|
||||
Cancel *struct{} `json:"cancel,omitempty"`
|
||||
// Requests that can be augmented with task metadata.
|
||||
Requests *TaskRequestsCapability `json:"requests,omitempty"`
|
||||
}
|
||||
|
||||
// TaskRequestsCapability indicates which request types support task augmentation.
|
||||
type TaskRequestsCapability struct {
|
||||
// Tool-related capabilities.
|
||||
Tools *struct {
|
||||
// Whether tools/call can be augmented with task metadata.
|
||||
Call *struct{} `json:"call,omitempty"`
|
||||
} `json:"tools,omitempty"`
|
||||
// Sampling-related capabilities.
|
||||
Sampling *struct {
|
||||
// Whether sampling/createMessage can be augmented with task metadata.
|
||||
CreateMessage *struct{} `json:"createMessage,omitempty"`
|
||||
} `json:"sampling,omitempty"`
|
||||
// Elicitation-related capabilities.
|
||||
Elicitation *struct {
|
||||
// Whether elicitation/create can be augmented with task metadata.
|
||||
Create *struct{} `json:"create,omitempty"`
|
||||
} `json:"elicitation,omitempty"`
|
||||
}
|
||||
|
||||
// TaskStatus represents the execution state of a task.
|
||||
type TaskStatus string
|
||||
|
||||
const (
|
||||
// TaskStatusWorking indicates the request is currently being processed.
|
||||
TaskStatusWorking TaskStatus = "working"
|
||||
// TaskStatusInputRequired indicates the receiver needs input from the requestor.
|
||||
// NOTE: This status is defined by the spec but not yet implemented in this SDK.
|
||||
// The input_required flow requires integration with elicitation which is planned
|
||||
// for a future release.
|
||||
TaskStatusInputRequired TaskStatus = "input_required"
|
||||
// TaskStatusCompleted indicates the request completed successfully.
|
||||
TaskStatusCompleted TaskStatus = "completed"
|
||||
// TaskStatusFailed indicates the request did not complete successfully.
|
||||
TaskStatusFailed TaskStatus = "failed"
|
||||
// TaskStatusCancelled indicates the request was cancelled before completion.
|
||||
TaskStatusCancelled TaskStatus = "cancelled"
|
||||
)
|
||||
|
||||
// IsTerminal returns true if the task status is terminal (completed, failed, or cancelled).
|
||||
func (s TaskStatus) IsTerminal() bool {
|
||||
return s == TaskStatusCompleted || s == TaskStatusFailed || s == TaskStatusCancelled
|
||||
}
|
||||
|
||||
// Task represents the execution state of a request.
|
||||
type Task struct {
|
||||
// Unique identifier for the task.
|
||||
TaskId string `json:"taskId"`
|
||||
// Current state of the task execution.
|
||||
Status TaskStatus `json:"status"`
|
||||
// Optional human-readable message describing the current state.
|
||||
StatusMessage string `json:"statusMessage,omitempty"`
|
||||
// ISO 8601 timestamp when the task was created.
|
||||
CreatedAt string `json:"createdAt"`
|
||||
// ISO 8601 timestamp when the task was last updated.
|
||||
LastUpdatedAt string `json:"lastUpdatedAt"`
|
||||
// Time in milliseconds from creation before task may be deleted.
|
||||
// If null, the task has no expiration.
|
||||
TTL *int64 `json:"ttl"`
|
||||
// Suggested time in milliseconds between status checks.
|
||||
PollInterval *int64 `json:"pollInterval,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the task ID, implementing the Named interface for pagination.
|
||||
func (t Task) GetName() string {
|
||||
return t.TaskId
|
||||
}
|
||||
|
||||
// TaskParams represents the task metadata included when augmenting a request.
|
||||
type TaskParams struct {
|
||||
// Requested duration in milliseconds to retain task from creation.
|
||||
TTL *int64 `json:"ttl,omitempty"`
|
||||
}
|
||||
|
||||
// CreateTaskResult is returned immediately when a task-augmented request is accepted.
|
||||
// It contains task metadata rather than the actual operation result.
|
||||
type CreateTaskResult struct {
|
||||
Result
|
||||
Task Task `json:"task"`
|
||||
}
|
||||
|
||||
// GetTaskRequest retrieves the current status of a task.
|
||||
type GetTaskRequest struct {
|
||||
Request
|
||||
Header http.Header `json:"-"`
|
||||
Params GetTaskParams `json:"params"`
|
||||
}
|
||||
|
||||
type GetTaskParams struct {
|
||||
TaskId string `json:"taskId"`
|
||||
}
|
||||
|
||||
// GetTaskResult returns the current state of a task.
|
||||
type GetTaskResult struct {
|
||||
Result
|
||||
Task
|
||||
}
|
||||
|
||||
// ListTasksRequest retrieves a paginated list of tasks.
|
||||
type ListTasksRequest struct {
|
||||
PaginatedRequest
|
||||
Header http.Header `json:"-"`
|
||||
}
|
||||
|
||||
// ListTasksResult returns a list of tasks.
|
||||
type ListTasksResult struct {
|
||||
PaginatedResult
|
||||
Tasks []Task `json:"tasks"`
|
||||
}
|
||||
|
||||
// TaskResultRequest retrieves the result of a completed task.
|
||||
type TaskResultRequest struct {
|
||||
Request
|
||||
Header http.Header `json:"-"`
|
||||
Params TaskResultParams `json:"params"`
|
||||
}
|
||||
|
||||
type TaskResultParams struct {
|
||||
TaskId string `json:"taskId"`
|
||||
}
|
||||
|
||||
// TaskResultResult contains the actual operation result.
|
||||
// For task-augmented tool calls, this embeds the CallToolResult fields.
|
||||
type TaskResultResult struct {
|
||||
Result
|
||||
// Tool call result fields (for task-augmented tool calls)
|
||||
Content []Content `json:"content,omitempty"`
|
||||
StructuredContent any `json:"structuredContent,omitempty"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
}
|
||||
|
||||
// CancelTaskRequest cancels an in-progress task.
|
||||
type CancelTaskRequest struct {
|
||||
Request
|
||||
Header http.Header `json:"-"`
|
||||
Params CancelTaskParams `json:"params"`
|
||||
}
|
||||
|
||||
type CancelTaskParams struct {
|
||||
TaskId string `json:"taskId"`
|
||||
}
|
||||
|
||||
// CancelTaskResult returns the cancelled task state.
|
||||
type CancelTaskResult struct {
|
||||
Result
|
||||
Task
|
||||
}
|
||||
|
||||
// TaskStatusNotification is sent when a task's status changes.
|
||||
type TaskStatusNotification struct {
|
||||
Notification
|
||||
Params TaskStatusNotificationParams `json:"params"`
|
||||
}
|
||||
|
||||
type TaskStatusNotificationParams struct {
|
||||
Task
|
||||
}
|
||||
|
||||
// ClientRequest represents any request that can be sent from client to server.
|
||||
type ClientRequest any
|
||||
|
||||
@@ -1250,3 +1578,24 @@ func UnmarshalContent(data []byte) (Content, error) {
|
||||
return nil, fmt.Errorf("unknown content type: %s", contentType)
|
||||
}
|
||||
}
|
||||
|
||||
// ElicitationCapability represents the elicitation capabilities of a client or server.
|
||||
type ElicitationCapability struct {
|
||||
Form *struct{} `json:"form,omitempty"` // Supports form mode
|
||||
URL *struct{} `json:"url,omitempty"` // Supports URL mode
|
||||
}
|
||||
|
||||
// NewElicitationCompleteNotification creates a new elicitation complete notification.
|
||||
func NewElicitationCompleteNotification(elicitationID string) JSONRPCNotification {
|
||||
return JSONRPCNotification{
|
||||
JSONRPC: JSONRPC_VERSION,
|
||||
Notification: Notification{
|
||||
Method: string(MethodNotificationElicitationComplete),
|
||||
Params: NotificationParams{
|
||||
AdditionalFields: map[string]any{
|
||||
"elicitationId": elicitationID,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
+23
-3
@@ -353,7 +353,7 @@ func NewToolResultImage(text, imageData, mimeType string) *CallToolResult {
|
||||
}
|
||||
|
||||
// NewToolResultAudio creates a new CallToolResult with both text and audio content
|
||||
func NewToolResultAudio(text, imageData, mimeType string) *CallToolResult {
|
||||
func NewToolResultAudio(text, audioData, mimeType string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
@@ -362,7 +362,7 @@ func NewToolResultAudio(text, imageData, mimeType string) *CallToolResult {
|
||||
},
|
||||
AudioContent{
|
||||
Type: ContentTypeAudio,
|
||||
Data: imageData,
|
||||
Data: audioData,
|
||||
MIMEType: mimeType,
|
||||
},
|
||||
},
|
||||
@@ -536,13 +536,22 @@ func ExtractString(data map[string]any, key string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// ParseAnnotations parses priority, audience, and lastModified fields from the provided map
|
||||
// and returns an Annotations struct populated with any valid values found.
|
||||
// If data is nil, ParseAnnotations returns nil. Priority is set when a numeric value can be
|
||||
// parsed and is stored as a *float64. Audience is populated from string values and includes
|
||||
// only RoleUser and RoleAssistant entries. LastModified is set when the value is a string.
|
||||
func ParseAnnotations(data map[string]any) *Annotations {
|
||||
if data == nil {
|
||||
return nil
|
||||
}
|
||||
annotations := &Annotations{}
|
||||
if value, ok := data["priority"]; ok {
|
||||
annotations.Priority = cast.ToFloat64(value)
|
||||
if value != nil {
|
||||
if priority, err := cast.ToFloat64E(value); err == nil {
|
||||
annotations.Priority = &priority
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if value, ok := data["audience"]; ok {
|
||||
@@ -553,6 +562,12 @@ func ParseAnnotations(data map[string]any) *Annotations {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if value, ok := data["lastModified"]; ok {
|
||||
if str, ok := value.(string); ok {
|
||||
annotations.LastModified = str
|
||||
}
|
||||
}
|
||||
return annotations
|
||||
|
||||
}
|
||||
@@ -950,6 +965,11 @@ func ToBoolPtr(b bool) *bool {
|
||||
return &b
|
||||
}
|
||||
|
||||
// ToInt64Ptr returns a pointer to the given int64 value
|
||||
func ToInt64Ptr(i int64) *int64 {
|
||||
return &i
|
||||
}
|
||||
|
||||
// GetTextFromContent extracts text from a Content interface that might be a TextContent struct
|
||||
// or a map[string]any that was unmarshaled from JSON. This is useful when dealing with content
|
||||
// that comes from different transport layers that may handle JSON differently.
|
||||
|
||||
Loaded 100 of 127 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user