llmgw/internal/cli/users.go

package cli

import (
	"errors"
	"fmt"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

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

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

	cmd := &cobra.Command{
		Use:   "users",
		Short: "Approve and manage the people who sign up",
		Long: "Signing up creates an account that cannot do anything. Approval is what\n" +
			"makes it work, and it happens here rather than in a browser: approving is\n" +
			"the step that decides who gets to spend the cluster, and it needs a shell\n" +
			"on this machine rather than a session cookie.",
	}
	sf.bind(cmd)
	open := func(cmd *cobra.Command) (*store.DB, error) {
		return sf.open(cmd)
	}

	cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open),
		newUserStateCmd(open, "suspend", store.UserSuspended,
			"Stop an account working, keeping it and its history"),
		newUserStateCmd(open, "reinstate", store.UserActive,
			"Let a suspended account work again"),
		newUserCreateCmd(open), newUserResetPasswordCmd(open),
		newUserJoinCmd(open), newUserLeaveCmd(open), newUserDeleteCmd(open),
		newUserAdminCmd(open, "promote", true, "Let an account see the cluster dashboard at /admin"),
		newUserAdminCmd(open, "demote", false, "Take away access to the cluster dashboard"),
		newUserRoleCmd(open))
	return cmd
}

type userOpener func(*cobra.Command) (*store.DB, error)

func newUserListCmd(open userOpener) *cobra.Command {
	var pending bool
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List accounts",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			state := store.UserState("")
			if pending {
				state = store.UserPending
			}
			users, err := s.Users(cmd.Context(), state)
			if err != nil {
				return err
			}
			if len(users) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "no accounts")
				return nil
			}

			all, err := s.AllMemberships(cmd.Context())
			if err != nil {
				return err
			}
			where := map[string][]string{}
			for _, m := range all {
				where[m.UserID] = append(where[m.UserID], m.Path)
			}
			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "ID\tEMAIL\tSTATE\tADMIN\tMEMBER OF\tSIGNED UP\tLAST LOGIN")
			now := time.Now()
			for _, u := range users {
				admin := "-"
				if u.IsAdmin {
					admin = "yes"
				}
				member := "on their own"
				if ms := where[u.ID]; len(ms) > 0 {
					member = strings.Join(ms, "; ")
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
					u.ID, u.Email, u.State, admin, member,
					ago(u.CreatedAt, now), ago(u.LastLoginAt, now))
			}
			return w.Flush()
		},
	}
	cmd.Flags().BoolVar(&pending, "pending", false, "only accounts waiting for approval")
	return cmd
}

func newUserApproveCmd(open userOpener) *cobra.Command {
	var project, by string
	cmd := &cobra.Command{
		Use:   "approve <user>",
		Short: "Let an account sign in and issue its own keys",
		Long: "A project is optional. Approved into nothing, a person uses the gateway on\n" +
			"their own — personal keys and chat within their personal limits — and can\n" +
			"be added to organizations, teams and projects with 'llmgw users join'.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			var p store.ProjectPath
			if project != "" {
				if p, err = resolveProject(cmd.Context(), s, project); err != nil {
					return err
				}
			}
			if err := s.ApproveUser(cmd.Context(), u.ID, p.ID, by, time.Now()); err != nil {
				return fmt.Errorf("approving %s: %w", u.Email, err)
			}
			if p.ID == "" {
				fmt.Fprintf(cmd.OutOrStdout(), "%s approved, on their own for now\n", u.Email)
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s approved into %s\n", u.Email, p.Path)
			}
			cmd.PrintErrln("They can sign in and create their own keys from /profile.")
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "",
		"a project to put them in (optional): id, path (Acme/platform/gateway) or unique name")
	cmd.Flags().StringVar(&by, "by", "", "who approved, for the record")
	return cmd
}

func newUserStateCmd(open userOpener, verb string, state store.UserState, short string) *cobra.Command {
	return &cobra.Command{
		Use:   verb + " <user>",
		Short: short,
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			if err := s.SetUserState(cmd.Context(), u.ID, state); err != nil {
				return err
			}
			if state == store.UserSuspended {
				// Otherwise the person stays signed in and suspension is
				// advisory until their session happens to expire.
				if err := s.DeleteUserSessions(cmd.Context(), u.ID); err != nil {
					return err
				}
				cmd.PrintErrln("Their sessions were closed. Their API keys still work — revoke those separately.")
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s is now %s\n", u.Email, state)
			return nil
		},
	}
}

func newUserRoleCmd(open userOpener) *cobra.Command {
	return &cobra.Command{
		Use:   "role <user> <viewer|operator|billing|admin|none>",
		Short: "Give an account a part in running the gateway, from the pages under /admin",
		Long: "viewer looks at every administration page and changes nothing. operator\n" +
			"also runs the cluster from them: models through the mirror, a node's\n" +
			"engines. billing sees usage and quotas. admin may do anything. none takes\n" +
			"the role away. Every change made from those pages is in `llmgw ctl audit`.",
		Args: cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			role := args[1]
			if role == "none" {
				role = ""
			}
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()
			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			if err := s.SetStaffRole(cmd.Context(), u.ID, role); err != nil {
				return err
			}
			if role == "" {
				fmt.Fprintf(cmd.OutOrStdout(), "%s has no part in running the gateway now\n", u.Email)
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s is now %s\n", u.Email, role)
			}
			return nil
		},
	}
}

func newUserAdminCmd(open userOpener, verb string, admin bool, short string) *cobra.Command {
	return &cobra.Command{
		Use:   verb + " <user>",
		Short: short,
		Long: "The dashboard is read-only, so this grants looking and nothing else.\n" +
			"Everything that changes the cluster or issues a credential stays here on\n" +
			"the CLI, where it needs a shell rather than a session — which keeps the\n" +
			"worst case of a stolen admin cookie to \"saw the cluster\".",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			if err := s.SetAdmin(cmd.Context(), u.ID, admin); err != nil {
				return err
			}
			if admin {
				fmt.Fprintf(cmd.OutOrStdout(), "%s can now see /admin\n", u.Email)
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s can no longer see /admin\n", u.Email)
				cmd.PrintErrln("Their session stays open; the page is read-only. " +
					"Suspend the account to close it.")
			}
			return nil
		},
	}
}

func newUserResetPasswordCmd(open userOpener) *cobra.Command {
	var password string
	cmd := &cobra.Command{
		Use:   "reset-password <user>",
		Short: "Set a new password for someone who lost theirs",
		Long: "There is no self-service reset, because this gateway sends no email and\n" +
			"a reset link with nowhere to send it is not a reset. Set one here, tell\n" +
			"them out of band, and have them change it from their profile.\n\n" +
			"Every session they have open is closed: if they lost the password, the\n" +
			"reason might be that somebody else has it.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if password == "" {
				return fmt.Errorf("--password is required")
			}
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			hash, err := account.HashPassword(password)
			if err != nil {
				return err
			}
			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			if err := s.SetPassword(cmd.Context(), u.ID, hash); err != nil {
				return err
			}
			if err := s.DeleteUserSessions(cmd.Context(), u.ID); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "password set for %s\n", u.Email)
			cmd.PrintErrln("Their sessions were closed. Their API keys still work.")
			return nil
		},
	}
	cmd.Flags().StringVar(&password, "password", "", "the new password")
	return cmd
}

func newUserCreateCmd(open userOpener) *cobra.Command {
	var project, password string
	cmd := &cobra.Command{
		Use:   "create <email>",
		Short: "Create an account directly, already approved",
		Long: "For when signup is closed, or for the first account on a new deployment.\n" +
			"The password is set here and should be changed by its owner.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if password == "" {
				return fmt.Errorf("--password is required")
			}
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()

			var p store.ProjectPath
			if project != "" {
				if p, err = resolveProject(cmd.Context(), s, project); err != nil {
					return err
				}
			}
			hash, err := account.HashPassword(password)
			if err != nil {
				return err
			}
			id := randomID("usr")
			u := store.User{
				ID: id, Email: args[0], Name: args[0],
				PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
			}
			if err := s.CreateUser(cmd.Context(), u); err != nil {
				return err
			}
			if err := s.ApproveUser(cmd.Context(), id, p.ID, "cli", time.Now()); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s\n", id)
			if p.ID == "" {
				cmd.PrintErrf("%s created and approved, on their own for now\n", args[0])
			} else {
				cmd.PrintErrf("%s created and approved into %s\n", args[0], p.Path)
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "",
		"a project to put them in (optional): id, path or unique name")
	cmd.Flags().StringVar(&password, "password", "", "initial password")
	return cmd
}

// scopeFlags names one org, team or project by whichever flag is given.
type scopeFlags struct{ org, team, project string }

func (f *scopeFlags) bind(cmd *cobra.Command) {
	cmd.Flags().StringVar(&f.org, "org", "", "an organization: id or name")
	cmd.Flags().StringVar(&f.team, "team", "", "a team: id, path (Acme/ml) or unique name")
	cmd.Flags().StringVar(&f.project, "project", "", "a project: id, path (Acme/ml/training) or unique name")
}

func (f *scopeFlags) resolve(cmd *cobra.Command, s *store.DB) (store.Scope, string, string, error) {
	n := 0
	for _, v := range []string{f.org, f.team, f.project} {
		if v != "" {
			n++
		}
	}
	if n != 1 {
		return "", "", "", fmt.Errorf("give exactly one of --org, --team or --project")
	}
	ctx := cmd.Context()
	switch {
	case f.org != "":
		o, err := resolveOrg(ctx, s, f.org)
		return store.ScopeOrg, o.ID, o.Name, err
	case f.team != "":
		t, err := resolveTeam(ctx, s, f.team)
		return store.ScopeTeam, t.ID, t.OrgName + " / " + t.Name, err
	default:
		p, err := resolveProject(ctx, s, f.project)
		return store.ScopeProject, p.ID, p.Path, err
	}
}

func newUserJoinCmd(open userOpener) *cobra.Command {
	var where scopeFlags
	var role, by string
	cmd := &cobra.Command{
		Use:   "join <user> --org|--team|--project <which>",
		Short: "Make a person a member of an organization, team or project",
		Long: "A person can be a member of any number of them. A member of a team or\n" +
			"organization may use every project in it.",
		Example: "  llmgw users join ayse@example.com --team Acme/ml\n" +
			"  llmgw users join ayse@example.com --project Scania/research/vision",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()
			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			scope, id, label, err := where.resolve(cmd, s)
			if err != nil {
				return err
			}
			if err := s.AddMembership(cmd.Context(), u.ID, scope, id, role, by, time.Now()); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s is a %s of %s\n", u.Email, role, label)
			return nil
		},
	}
	where.bind(cmd)
	cmd.Flags().StringVar(&role, "role", store.RoleMember, "member, or manager: with --org, may see the organization's page")
	cmd.Flags().StringVar(&by, "by", "", "who added them, for the record")
	return cmd
}

func newUserLeaveCmd(open userOpener) *cobra.Command {
	var where scopeFlags
	cmd := &cobra.Command{
		Use:   "leave <user> --org|--team|--project <which>",
		Short: "Take a person out of an organization, team or project",
		Long: "Keys they made in a project they can no longer use are revoked; a project\n" +
			"they still reach through another membership keeps them. Their personal\n" +
			"keys are not touched.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()
			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			scope, id, label, err := where.resolve(cmd, s)
			if err != nil {
				return err
			}
			revoked, err := s.RemoveMembership(cmd.Context(), u.ID, scope, id, time.Now())
			if errors.Is(err, store.ErrNotFound) {
				return fmt.Errorf("%s is not a member of %s itself", u.Email, label)
			}
			if err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s is no longer in %s (%s revoked)\n",
				u.Email, label, countOf(revoked, "key"))
			return nil
		},
	}
	where.bind(cmd)
	return cmd
}

func newUserDeleteCmd(open userOpener) *cobra.Command {
	var yes bool
	cmd := &cobra.Command{
		Use:   "delete <user>",
		Short: "Remove an account for good",
		Long: "Signs them out, revokes every key they made, and removes the account. The\n" +
			"keys are kept, revoked and ownerless, because usage history refers to them.\n" +
			"To stop an account working and keep it, use suspend.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd)
			if err != nil {
				return err
			}
			defer s.Close()
			u, err := resolveUser(cmd.Context(), s, args[0])
			if err != nil {
				return err
			}
			if !yes {
				return fmt.Errorf("this removes %s and revokes their keys; run again with --yes", u.Email)
			}
			if err := s.DeleteUser(cmd.Context(), u.ID, time.Now()); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", u.Email)
			return nil
		},
	}
	cmd.Flags().BoolVar(&yes, "yes", false, "confirm")
	return cmd
}