Skip to content

Commit

Permalink
chore: make cred stores be tools
Browse files Browse the repository at this point in the history
  • Loading branch information
ibuildthecloud committed Nov 5, 2024
1 parent 4ce687f commit 9aeb1cd
Show file tree
Hide file tree
Showing 20 changed files with 268 additions and 422 deletions.
27 changes: 5 additions & 22 deletions pkg/cli/credential.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,7 @@ import (
"time"

cmd2 "github.com/gptscript-ai/cmd"
"github.com/gptscript-ai/gptscript/pkg/config"
"github.com/gptscript-ai/gptscript/pkg/credentials"
"github.com/gptscript-ai/gptscript/pkg/gptscript"
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
"github.com/spf13/cobra"
)

Expand All @@ -37,33 +34,19 @@ func (c *Credential) Customize(cmd *cobra.Command) {
}

func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
if err != nil {
return fmt.Errorf("failed to read CLI config: %w", err)
}

opts, err := c.root.NewGPTScriptOpts()
if err != nil {
return err
}
opts = gptscript.Complete(opts)
if opts.Runner.RuntimeManager == nil {
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir, opts.SystemToolsDir)
}

ctxs := opts.CredentialContexts
if c.AllContexts {
ctxs = []string{credentials.AllCredentialContexts}
}

if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg); err != nil {
gptScript, err := gptscript.New(cmd.Context(), opts)
if err != nil {
return err
}
defer gptScript.Close(true)

// Initialize the credential store and get all the credentials.
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctxs, opts.Cache.CacheDir)
store, err := gptScript.CredentialStoreFactory.NewStore(gptScript.DefaultCredentialContexts)
if err != nil {
return fmt.Errorf("failed to get credentials store: %w", err)
return err
}

creds, err := store.List(cmd.Context())
Expand Down
19 changes: 4 additions & 15 deletions pkg/cli/credential_delete.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,7 @@ package cli
import (
"fmt"

"github.com/gptscript-ai/gptscript/pkg/config"
"github.com/gptscript-ai/gptscript/pkg/credentials"
"github.com/gptscript-ai/gptscript/pkg/gptscript"
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
"github.com/spf13/cobra"
)

Expand All @@ -28,23 +25,15 @@ func (c *Delete) Run(cmd *cobra.Command, args []string) error {
return err
}

cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
gptScript, err := gptscript.New(cmd.Context(), opts)
if err != nil {
return fmt.Errorf("failed to read CLI config: %w", err)
}

opts = gptscript.Complete(opts)
if opts.Runner.RuntimeManager == nil {
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir, opts.SystemToolsDir)
}

if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg); err != nil {
return err
}
defer gptScript.Close(true)

store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, opts.CredentialContexts, opts.Cache.CacheDir)
store, err := gptScript.CredentialStoreFactory.NewStore(gptScript.DefaultCredentialContexts)
if err != nil {
return fmt.Errorf("failed to get credentials store: %w", err)
return err
}

if err = store.Remove(cmd.Context(), args[0]); err != nil {
Expand Down
19 changes: 4 additions & 15 deletions pkg/cli/credential_show.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,7 @@ import (
"os"
"text/tabwriter"

"github.com/gptscript-ai/gptscript/pkg/config"
"github.com/gptscript-ai/gptscript/pkg/credentials"
"github.com/gptscript-ai/gptscript/pkg/gptscript"
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
"github.com/spf13/cobra"
)

Expand All @@ -30,23 +27,15 @@ func (c *Show) Run(cmd *cobra.Command, args []string) error {
return err
}

cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
gptScript, err := gptscript.New(cmd.Context(), opts)
if err != nil {
return fmt.Errorf("failed to read CLI config: %w", err)
}

opts = gptscript.Complete(opts)
if opts.Runner.RuntimeManager == nil {
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir, opts.SystemToolsDir)
}

if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg); err != nil {
return err
}
defer gptScript.Close(true)

store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, opts.CredentialContexts, opts.Cache.CacheDir)
store, err := gptScript.CredentialStoreFactory.NewStore(gptScript.DefaultCredentialContexts)
if err != nil {
return fmt.Errorf("failed to get credentials store: %w", err)
return err
}

cred, exists, err := store.Get(cmd.Context(), args[0])
Expand Down
72 changes: 21 additions & 51 deletions pkg/config/cliconfig.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,9 @@ package config
import (
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"os"
"runtime"
"slices"
"strings"
"sync"

Expand All @@ -21,28 +19,13 @@ const (
SecretserviceCredHelper = "secretservice"
PassCredHelper = "pass"
FileCredHelper = "file"
SqliteCredHelper = "sqlite"
PostgresCredHelper = "postgres"

GPTScriptHelperPrefix = "gptscript-credential-"
)

var (
darwinHelpers = []string{OsxkeychainCredHelper, FileCredHelper, SqliteCredHelper, PostgresCredHelper}
windowsHelpers = []string{WincredCredHelper, FileCredHelper}
linuxHelpers = []string{SecretserviceCredHelper, PassCredHelper, FileCredHelper, SqliteCredHelper, PostgresCredHelper}
// Helpers is a list of all supported credential helpers from github.com/gptscript-ai/gptscript-credential-helpers
Helpers = []string{WincredCredHelper, OsxkeychainCredHelper, SecretserviceCredHelper, PassCredHelper}
)

func listAsString(helpers []string) string {
if len(helpers) == 0 {
return ""
} else if len(helpers) == 1 {
return helpers[0]
}

return strings.Join(helpers[:len(helpers)-1], ", ") + " or " + helpers[len(helpers)-1]
}

type AuthConfig types.AuthConfig

func (a AuthConfig) MarshalJSON() ([]byte, error) {
Expand Down Expand Up @@ -74,8 +57,8 @@ func (a *AuthConfig) UnmarshalJSON(data []byte) error {
type CLIConfig struct {
Auths map[string]AuthConfig `json:"auths,omitempty"`
CredentialsStore string `json:"credsStore,omitempty"`
Integrations map[string]string `json:"integrations,omitempty"`

raw []byte
auths map[string]types.AuthConfig
authsLock *sync.Mutex
location string
Expand Down Expand Up @@ -108,7 +91,19 @@ func (c *CLIConfig) Save() error {
}
c.auths = nil
}
data, err := json.Marshal(c)

// This is to not overwrite additional fields that might be the config file
out := map[string]any{}
if len(c.raw) > 0 {
err := json.Unmarshal(c.raw, &out)
if err != nil {
return err
}
}
out["auths"] = c.Auths
out["credsStore"] = c.CredentialsStore

data, err := json.Marshal(out)
if err != nil {
return err
}
Expand Down Expand Up @@ -154,34 +149,22 @@ func ReadCLIConfig(gptscriptConfigFile string) (*CLIConfig, error) {
result := &CLIConfig{
authsLock: &sync.Mutex{},
location: gptscriptConfigFile,
raw: data,
}
if err := json.Unmarshal(data, result); err != nil {
return nil, fmt.Errorf("failed to unmarshal %s: %v", gptscriptConfigFile, err)
}

if store := os.Getenv("GPTSCRIPT_CREDENTIAL_STORE"); store != "" {
result.CredentialsStore = store
}

if result.CredentialsStore == "" {
if err := result.setDefaultCredentialsStore(); err != nil {
return nil, err
}
}

if !isValidCredentialHelper(result.CredentialsStore) {
errMsg := fmt.Sprintf("invalid credential store '%s'", result.CredentialsStore)
switch runtime.GOOS {
case "darwin":
errMsg += fmt.Sprintf(" (use %s)", listAsString(darwinHelpers))
case "windows":
errMsg += fmt.Sprintf(" (use %s)", listAsString(windowsHelpers))
case "linux":
errMsg += fmt.Sprintf(" (use %s)", listAsString(linuxHelpers))
default:
errMsg += " (use file)"
}
errMsg += fmt.Sprintf("\nPlease edit your config file at %s to fix this.", result.location)

return nil, errors.New(errMsg)
}

return result, nil
}

Expand All @@ -197,19 +180,6 @@ func (c *CLIConfig) setDefaultCredentialsStore() error {
return c.Save()
}

func isValidCredentialHelper(helper string) bool {
switch runtime.GOOS {
case "darwin":
return slices.Contains(darwinHelpers, helper)
case "windows":
return slices.Contains(windowsHelpers, helper)
case "linux":
return slices.Contains(linuxHelpers, helper)
default:
return helper == FileCredHelper
}
}

func readFile(path string) ([]byte, error) {
data, err := os.ReadFile(path)
if os.IsNotExist(err) {
Expand Down
77 changes: 77 additions & 0 deletions pkg/credentials/factory.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
package credentials

import (
"context"

"github.com/docker/docker-credential-helpers/client"
"github.com/gptscript-ai/gptscript/pkg/config"
"github.com/gptscript-ai/gptscript/pkg/types"
)

type ProgramLoaderRunner interface {
Load(ctx context.Context, toolName string) (prg types.Program, err error)
Run(ctx context.Context, prg types.Program, input string) (output string, err error)
}

func NewFactory(ctx context.Context, cfg *config.CLIConfig, plr ProgramLoaderRunner) (StoreFactory, error) {
toolName := translateToolName(cfg.CredentialsStore)
if toolName == config.FileCredHelper {
return StoreFactory{
file: true,
cfg: cfg,
}, nil
}

prg, err := plr.Load(ctx, toolName)
if err != nil {
return StoreFactory{}, err
}

return StoreFactory{
ctx: ctx,
prg: prg,
runner: plr,
cfg: cfg,
}, nil
}

type StoreFactory struct {
ctx context.Context
prg types.Program
file bool
runner ProgramLoaderRunner
cfg *config.CLIConfig
}

func (s *StoreFactory) NewStore(credCtxs []string) (CredentialStore, error) {
if err := validateCredentialCtx(credCtxs); err != nil {
return nil, err
}
if s.file {
return Store{
credCtxs: credCtxs,
cfg: s.cfg,
}, nil
}
return Store{
credCtxs: credCtxs,
cfg: s.cfg,
program: s.program,
}, nil
}

func (s *StoreFactory) program(args ...string) client.Program {
return &runnerProgram{
factory: s,
action: args[0],
}
}

func translateToolName(toolName string) string {
for _, helper := range config.Helpers {
if helper == toolName {
return "github.com/gptscript-ai/gptscript-credential-helpers/" + toolName + "/cmd"
}
}
return toolName
}
4 changes: 3 additions & 1 deletion pkg/credentials/noop.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package credentials

import "context"
import (
"context"
)

type NoopStore struct{}

Expand Down
29 changes: 29 additions & 0 deletions pkg/credentials/runnerprogram.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
package credentials

import (
"io"
)

type runnerProgram struct {
factory *StoreFactory
action string
output string
err error
}

func (r *runnerProgram) Output() ([]byte, error) {
return []byte(r.output), r.err
}

func (r *runnerProgram) Input(in io.Reader) {
input, err := io.ReadAll(in)
if err != nil {
r.err = err
return
}

prg := r.factory.prg
prg.EntryToolID = prg.ToolSet[prg.EntryToolID].LocalTools[r.action]

r.output, r.err = r.factory.runner.Run(r.factory.ctx, prg, string(input))
}
Loading

0 comments on commit 9aeb1cd

Please sign in to comment.