llmgw/internal/cli/keys.go

package cli

import (
	"context"
	"crypto/rand"
	"encoding/base32"
	"errors"
	"fmt"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func newTenancyCmd() *cobra.Command {
	var sf storeFlags

	cmd := &cobra.Command{
		Use:   "keys",
		Short: "Manage tenants and API keys",
		Long: "Works directly on the store file rather than through the admin API, so keys\n" +
			"can be issued before the server is running and without exposing an endpoint\n" +
			"that mints credentials.",
	}
	sf.bind(cmd)

	// Takes a context rather than the command, because some of these run
	// before the command tree is entered.
	open := func(ctx context.Context) (*store.DB, error) { return sf.openCtx(ctx) }

	cmd.AddCommand(
		newKeyInitCmd(open),
		newKeyCreateCmd(open),
		newKeyListCmd(open),
		newKeyRevokeCmd(open),
		newLimitsCmd(open),
		newThreadsCmd(open),
		newMemoryCmd(open),
		newMoveStoreCmd(),
	)
	return cmd
}

type opener func(context.Context) (*store.DB, error)

func newKeyInitCmd(open opener) *cobra.Command {
	var org, team, project string
	cmd := &cobra.Command{
		Use:   "init",
		Short: "Create the tenancy tree a key hangs off",
		Long: "Idempotent: running it again renames rather than duplicating, so it can be\n" +
			"kept in whatever provisions the machine.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			ctx := cmd.Context()
			if err := s.PutOrg(ctx, store.Org{ID: slug(org), Name: org}); err != nil {
				return err
			}
			if err := s.PutTeam(ctx, store.Team{ID: slug(team), OrgID: slug(org), Name: team}); err != nil {
				return err
			}
			if err := s.PutProject(ctx, store.Project{
				ID: slug(project), TeamID: slug(team), Name: project,
			}); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s / %s / %s\n", slug(org), slug(team), slug(project))
			return nil
		},
	}
	cmd.Flags().StringVar(&org, "org", "default", "organization name")
	cmd.Flags().StringVar(&team, "team", "default", "team name")
	cmd.Flags().StringVar(&project, "project", "default", "project name")
	return cmd
}

func newKeyCreateCmd(open opener) *cobra.Command {
	var project, band, profile, baseURL string
	var expires, claimTTL time.Duration
	var claim bool

	cmd := &cobra.Command{
		Use:   "create <name>",
		Short: "Issue an API key",
		Long: "The key is printed once and never again: only its hash is stored, so there\n" +
			"is nothing to recover it from. That is the point — a leaked database hands\n" +
			"over no working credential.\n\n" +
			"With --claim it prints a one-time link instead. Send the link rather than\n" +
			"the key: a key pasted into a chat window stays in that chat window forever,\n" +
			"while a link works once and then expires. The key is sealed with the link's\n" +
			"own token, which is not stored, so a pending claim is not a usable\n" +
			"credential sitting in the database either.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			secret, k, err := store.NewKey(randomID("key"), args[0], project)
			if err != nil {
				return err
			}
			k.CreatedAt = time.Now()
			k.Band, k.RoutingProfile = band, profile
			if expires > 0 {
				k.ExpiresAt = k.CreatedAt.Add(expires)
			}
			if err := s.CreateKey(cmd.Context(), k); err != nil {
				return fmt.Errorf("creating the key (does project %q exist? run `llmgw keys init`): %w",
					project, err)
			}

			// The one output that matters goes to stdout and everything else to
			// stderr, so `keys create ... > out.txt` captures it and nothing
			// else. Cobra's Print writes to stderr, which would have made the
			// useful half the half that cannot be piped.
			out := cmd.OutOrStdout()

			if claim {
				if baseURL == "" {
					return errors.New("--claim needs --base-url, the address users reach this gateway at")
				}
				token, err := store.NewClaimToken()
				if err != nil {
					return err
				}
				if err := s.CreateClaim(cmd.Context(), k.ID, token, secret, claimTTL); err != nil {
					return err
				}
				link := strings.TrimRight(baseURL, "/") + "/claim/" + token
				fmt.Fprintln(out, link)
				cmd.PrintErrf("\nid       %s\nproject  %s\nexpires  %s\n",
					k.ID, k.ProjectID, time.Now().Add(orDefault(claimTTL, store.ClaimTTL)).Format(time.RFC3339))
				cmd.PrintErrln("\nSend the link. It works once, and the key is shown only to whoever opens it.")
				return nil
			}

			fmt.Fprintln(out, secret)
			cmd.PrintErrf("\nid       %s\nproject  %s\n", k.ID, k.ProjectID)
			if !k.ExpiresAt.IsZero() {
				cmd.PrintErrf("expires  %s\n", k.ExpiresAt.Format(time.RFC3339))
			}
			cmd.PrintErrln("\nThis is the only time the key is shown. Store it now.")
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "default", "project the key belongs to")
	cmd.Flags().StringVar(&band, "band", "", "admission band; empty inherits")
	cmd.Flags().StringVar(&profile, "profile", "", "routing profile; empty inherits")
	cmd.Flags().DurationVar(&expires, "expires-in", 0, "lifetime; 0 never expires")
	cmd.Flags().BoolVar(&claim, "claim", false, "print a one-time link instead of the key")
	cmd.Flags().StringVar(&baseURL, "base-url", "", "public address of the gateway, for --claim")
	cmd.Flags().DurationVar(&claimTTL, "claim-ttl", store.ClaimTTL, "how long the link stays usable")
	return cmd
}

func orDefault(d, fallback time.Duration) time.Duration {
	if d <= 0 {
		return fallback
	}
	return d
}

func newKeyListCmd(open opener) *cobra.Command {
	var project string
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List a project's keys",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			keys, err := s.Keys(cmd.Context(), project)
			if err != nil {
				return err
			}
			if len(keys) == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "no keys in project %q\n", project)
				return nil
			}

			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "ID\tNAME\tKEY\tBAND\tSTATE\tLAST USED")
			now := time.Now()
			for _, k := range keys {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
					k.ID, k.Name, k.Prefix, dash(k.Band), keyState(k, now), ago(k.LastUsedAt, now))
			}
			return w.Flush()
		},
	}
	cmd.Flags().StringVar(&project, "project", "default", "project to list")
	return cmd
}

func newKeyRevokeCmd(open opener) *cobra.Command {
	return &cobra.Command{
		Use:   "revoke <key-id>",
		Short: "Stop a key working, without deleting it",
		Long: "A soft delete. Usage rows reference the key, and an audit trail that loses\n" +
			"its subject is not an audit trail.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			if err := s.RevokeKey(cmd.Context(), args[0], time.Now()); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "revoked %s\n", args[0])
			cmd.PrintErrln("A running server keeps accepting it until its cache expires (auth.cache_ttl).")
			return nil
		},
	}
}

func newLimitsCmd(open opener) *cobra.Command {
	var scope, id, period string
	var rpm, tpm int
	var budget int64
	var models []string

	cmd := &cobra.Command{
		Use:   "limits",
		Short: "Set limits at one level of the tenancy tree",
		Long: "Limits compose downward by taking the most restrictive value, so a key\n" +
			"cannot raise what its organization set. Model allowlists intersect: a key\n" +
			"cannot add a model its team withheld.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			l := store.Limits{Scope: store.Scope(scope), ScopeID: id}
			if cmd.Flags().Changed("rpm") {
				l.RPM = &rpm
			}
			if cmd.Flags().Changed("tpm") {
				l.TPM = &tpm
			}
			if cmd.Flags().Changed("budget") {
				l.Budget = &budget
				l.BudgetPeriod = store.Period(period)
			}
			if cmd.Flags().Changed("models") {
				// An explicitly empty list is an allowlist that permits
				// nothing, which is different from setting no list at all.
				if models == nil {
					models = []string{}
				}
				l.AllowedModels = models
			}
			if err := s.PutLimits(cmd.Context(), l); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "limits set on %s %s\n", scope, id)
			return nil
		},
	}
	cmd.Flags().StringVar(&scope, "scope", "project", "org | team | project | key")
	cmd.Flags().StringVar(&id, "id", "default", "id at that scope")
	cmd.Flags().IntVar(&rpm, "rpm", 0, "requests per minute")
	cmd.Flags().IntVar(&tpm, "tpm", 0, "tokens per minute")
	cmd.Flags().Int64Var(&budget, "budget", 0, "budget per period, in micro-units")
	cmd.Flags().StringVar(&period, "period", "month", "day | month")
	cmd.Flags().StringSliceVar(&models, "models", nil, "allowed models; empty allows none")
	return cmd
}

func keyState(k store.Key, now time.Time) string {
	switch {
	case !k.RevokedAt.IsZero():
		return "revoked"
	case !k.ExpiresAt.IsZero() && !now.Before(k.ExpiresAt):
		return "expired"
	default:
		return "active"
	}
}

func ago(t, now time.Time) string {
	if t.IsZero() {
		return "never"
	}
	return now.Sub(t).Round(time.Second).String() + " ago"
}

// slug turns a human name into an id.
func slug(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(strings.TrimSpace(name)) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == ' ' || r == '-' || r == '_':
			b.WriteByte('-')
		}
	}
	if b.Len() == 0 {
		return "default"
	}
	return b.String()
}

func randomID(prefix string) string {
	var buf [10]byte
	_, _ = rand.Read(buf[:])
	return prefix + "-" + strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
}