ollama-proxy/internal/cli/accounts.go

package cli

import (
	"bufio"
	"fmt"
	"os"
	"strings"

	"github.com/Atte149/ollama-proxy/internal/config"
	"golang.org/x/term"
)

// RunAccounts dispatches to the accounts subcommand.
func RunAccounts(args []string) int {
	if len(args) < 1 {
		printAccountsUsage()
		return 1
	}
	switch args[0] {
	case "add":
		return accountsAdd(args[1:])
	case "list":
		return accountsList(args[1:])
	case "remove", "rm":
		return accountsRemove(args[1:])
	case "set-base-url":
		return accountsSetBaseURL(args[1:])
	case "-h", "--help", "help":
		printAccountsUsage()
		return 0
	default:
		fmt.Fprintf(os.Stderr, "unknown accounts subcommand: %s\n\n", args[0])
		printAccountsUsage()
		return 2
	}
}

func printAccountsUsage() {
	fmt.Print(`Usage:
  ollama-proxy accounts add [API_KEY] [--name <alias>]
  ollama-proxy accounts list
  ollama-proxy accounts remove <id|name>
  ollama-proxy accounts set-base-url <url>
`)
}

// accountsAdd implements `accounts add [API_KEY] [--name <alias>]`.
func accountsAdd(args []string) int {
	var apiKey, name string
	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--name":
			if i+1 >= len(args) {
				fmt.Fprintln(os.Stderr, "--name requires an argument")
				return 2
			}
			name = args[i+1]
			i++
		case "-h", "--help":
			fmt.Println("usage: accounts add [API_KEY] [--name <alias>]")
			return 0
		default:
			if apiKey == "" {
				apiKey = args[i]
			} else {
				fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", args[i])
				return 2
			}
		}
	}

	if name == "" {
		name = promptString("Account name (alias)", "acct"+config.NewID()[:4])
	}
	if err := config.ValidateName(name); err != nil {
		fmt.Fprintln(os.Stderr, "invalid name:", err)
		return 2
	}

	if apiKey == "" {
		fmt.Print("Enter Ollama Cloud API key (input hidden): ")
		b, err := term.ReadPassword(int(os.Stdin.Fd()))
		fmt.Println()
		if err != nil {
			fmt.Fprintln(os.Stderr, "read key:", err)
			return 1
		}
		apiKey = strings.TrimSpace(string(b))
	}
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "empty API key, nothing to add")
		return 2
	}

	if err := config.EnsureAccountsDir(); err != nil {
		fmt.Fprintln(os.Stderr, "create config dir:", err)
		return 1
	}
	af, err := config.LoadAccounts()
	if err != nil {
		fmt.Fprintln(os.Stderr, "load accounts:", err)
		return 1
	}
	if af.FindByName(name) != nil {
		fmt.Fprintf(os.Stderr, "an account named %q already exists\n", name)
		return 2
	}
	acct := config.Account{
		ID:      config.NewID(),
		Name:    name,
		APIKey:  apiKey,
		Created: nowFn(),
	}
	af.Accounts = append(af.Accounts, acct)
	if err := af.Save(); err != nil {
		fmt.Fprintln(os.Stderr, "save accounts:", err)
		return 1
	}
	fmt.Printf("added account %s (name=%s key=%s)\n", acct.ID[:8], acct.Name, maskKey(acct.APIKey))
	return 0
}

// accountsList prints a table of accounts.
func accountsList(args []string) int {
	af, err := config.LoadAccounts()
	if err != nil {
		fmt.Fprintln(os.Stderr, "load accounts:", err)
		return 1
	}
	if len(af.Accounts) == 0 {
		fmt.Println("no accounts configured. Run: ollama-proxy accounts add")
		return 0
	}
	w := bufio.NewWriter(os.Stdout)
	defer w.Flush()
	fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", "ID", "NAME", "KEY", "CREATED")
	fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", strings.Repeat("-", 8), strings.Repeat("-", 14), strings.Repeat("-", 16), strings.Repeat("-", 20))
	for _, a := range af.Accounts {
		id := a.ID
		if len(id) > 8 {
			id = id[:8]
		}
		created := a.Created.Format("2006-01-02 15:04 MST")
		fmt.Fprintf(w, "%-10s %-16s %-18s %-22s\n", id, a.Name, maskKey(a.APIKey), created)
	}
	if af.BaseURL != "" {
		fmt.Fprintf(w, "\nDefault upstream: %s\n", af.BaseURL)
	}
	return 0
}

// accountsRemove removes an account by id or name.
func accountsRemove(args []string) int {
	if len(args) < 1 || args[0] == "-h" || args[0] == "--help" {
		fmt.Println("usage: accounts remove <id|name>")
		if len(args) < 1 {
			return 2
		}
		return 0
	}
	ident := args[0]
	af, err := config.LoadAccounts()
	if err != nil {
		fmt.Fprintln(os.Stderr, "load accounts:", err)
		return 1
	}
	if err := af.Remove(ident); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}
	if err := af.Save(); err != nil {
		fmt.Fprintln(os.Stderr, "save:", err)
		return 1
	}
	fmt.Printf("removed account %s\n", ident)
	return 0
}

// accountsSetBaseURL sets the default upstream URL stored in accounts.json.
func accountsSetBaseURL(args []string) int {
	if len(args) < 1 {
		fmt.Println("usage: accounts set-base-url <url>")
		return 2
	}
	url := args[0]
	if err := config.EnsureAccountsDir(); err != nil {
		fmt.Fprintln(os.Stderr, "create config dir:", err)
		return 1
	}
	af, err := config.LoadAccounts()
	if err != nil {
		fmt.Fprintln(os.Stderr, "load accounts:", err)
		return 1
	}
	af.BaseURL = url
	if err := af.Save(); err != nil {
		fmt.Fprintln(os.Stderr, "save:", err)
		return 1
	}
	fmt.Printf("default upstream set to %s\n", url)
	return 0
}

// promptString reads a line from stdin, returning def when the user just hits
// enter. Used for the optional account name prompt.
func promptString(label, def string) string {
	fmt.Printf("%s [%s]: ", label, def)
	r := bufio.NewReader(os.Stdin)
	line, _ := r.ReadString('\n')
	line = strings.TrimSpace(line)
	if line == "" {
		return def
	}
	return line
}

// maskKey mirrors config.accounts.maskKey but is exposed here for display.
func maskKey(key string) string {
	if len(key) <= 4 {
		return "****"
	}
	return "****" + key[len(key)-4:]
}