repos / pico

pico services - prose.sh, pastes.sh, imgs.sh, feeds.sh, pgs.sh
git clone https://github.com/picosh/pico.git

pico / imgs
Antonio Mika · 08 Oct 24

cli.go

  1package imgs
  2
  3import (
  4	"encoding/json"
  5	"flag"
  6	"fmt"
  7	"io"
  8	"log/slog"
  9	"net/http"
 10	"path/filepath"
 11	"strings"
 12
 13	"github.com/charmbracelet/ssh"
 14	"github.com/charmbracelet/wish"
 15	"github.com/google/uuid"
 16	"github.com/picosh/pico/db"
 17	"github.com/picosh/pico/shared/storage"
 18	"github.com/picosh/pico/tui/common"
 19	sst "github.com/picosh/pobj/storage"
 20	psub "github.com/picosh/pubsub"
 21	sendutils "github.com/picosh/send/utils"
 22	"github.com/picosh/utils"
 23)
 24
 25func getUser(s ssh.Session, dbpool db.DB) (*db.User, error) {
 26	if s.PublicKey() == nil {
 27		return nil, fmt.Errorf("key not found")
 28	}
 29
 30	key := utils.KeyForKeyText(s.PublicKey())
 31
 32	user, err := dbpool.FindUserForKey(s.User(), key)
 33	if err != nil {
 34		return nil, err
 35	}
 36
 37	if user.Name == "" {
 38		return nil, fmt.Errorf("must have username set")
 39	}
 40
 41	return user, nil
 42}
 43
 44func flagSet(cmdName string, sesh ssh.Session) (*flag.FlagSet, *bool) {
 45	cmd := flag.NewFlagSet(cmdName, flag.ContinueOnError)
 46	cmd.SetOutput(sesh)
 47	write := cmd.Bool("write", false, "apply changes")
 48	return cmd, write
 49}
 50
 51func flagCheck(cmd *flag.FlagSet, posArg string, cmdArgs []string) bool {
 52	_ = cmd.Parse(cmdArgs)
 53
 54	if posArg == "-h" || posArg == "--help" || posArg == "-help" {
 55		cmd.Usage()
 56		return false
 57	}
 58	return true
 59}
 60
 61type Cmd struct {
 62	User        *db.User
 63	Session     utils.CmdSession
 64	Log         *slog.Logger
 65	Dbpool      db.DB
 66	Write       bool
 67	Styles      common.Styles
 68	Storage     sst.ObjectStorage
 69	RegistryUrl string
 70}
 71
 72func (c *Cmd) output(out string) {
 73	_, _ = c.Session.Write([]byte(out + "\r\n"))
 74}
 75
 76func (c *Cmd) error(err error) {
 77	_, _ = fmt.Fprint(c.Session.Stderr(), err, "\r\n")
 78	_ = c.Session.Exit(1)
 79	_ = c.Session.Close()
 80}
 81
 82func (c *Cmd) bail(err error) {
 83	if err == nil {
 84		return
 85	}
 86	c.Log.Error(err.Error())
 87	c.error(err)
 88}
 89
 90func (c *Cmd) notice() {
 91	if !c.Write {
 92		c.output("\nNOTICE: changes not commited, use `--write` to save operation")
 93	}
 94}
 95
 96func (c *Cmd) help() {
 97	helpStr := "Commands: [help, ls, rm]\n"
 98	helpStr += "NOTICE: *must* append with `--write` for the changes to persist.\n"
 99	c.output(helpStr)
100}
101
102func (c *Cmd) rm(repo string) error {
103	bucket, err := c.Storage.GetBucket("imgs")
104	if err != nil {
105		return err
106	}
107
108	fp := filepath.Join("docker/registry/v2/repositories", c.User.Name, repo)
109
110	fileList, err := c.Storage.ListObjects(bucket, fp, true)
111	if err != nil {
112		return err
113	}
114
115	if len(fileList) == 0 {
116		c.output(fmt.Sprintf("repo not found (%s)", repo))
117		return nil
118	}
119	c.output(fmt.Sprintf("found (%d) objects for repo (%s), removing", len(fileList), repo))
120
121	for _, obj := range fileList {
122		fname := filepath.Join(fp, obj.Name())
123		intent := fmt.Sprintf("deleted (%s)", obj.Name())
124		c.Log.Info(
125			"attempting to delete file",
126			"user", c.User.Name,
127			"bucket", bucket.Name,
128			"repo", repo,
129			"filename", fname,
130		)
131		if c.Write {
132			err := c.Storage.DeleteObject(bucket, fname)
133			if err != nil {
134				return err
135			}
136		}
137		c.output(intent)
138	}
139
140	return nil
141}
142
143type RegistryCatalog struct {
144	Repos []string `json:"repositories"`
145}
146
147func (c *Cmd) ls() error {
148	res, err := http.Get(
149		fmt.Sprintf("http://%s/v2/_catalog", c.RegistryUrl),
150	)
151	if err != nil {
152		return err
153	}
154
155	body, err := io.ReadAll(res.Body)
156	if err != nil {
157		return err
158	}
159
160	var data RegistryCatalog
161	err = json.Unmarshal(body, &data)
162
163	if err != nil {
164		return err
165	}
166
167	if len(data.Repos) == 0 {
168		c.output("You don't have any repos on imgs.sh")
169		return nil
170	}
171
172	user := c.User.Name
173	out := "repos\n"
174	out += "-----\n"
175	for _, repo := range data.Repos {
176		if !strings.HasPrefix(repo, user+"/") {
177			continue
178		}
179		rr := strings.TrimPrefix(repo, user+"/")
180		out += fmt.Sprintf("%s\n", rr)
181	}
182	c.output(out)
183	return nil
184}
185
186type CliHandler struct {
187	DBPool      db.DB
188	Logger      *slog.Logger
189	Storage     storage.StorageServe
190	RegistryUrl string
191	PubSub      psub.PubSub
192}
193
194func WishMiddleware(handler *CliHandler) wish.Middleware {
195	dbpool := handler.DBPool
196	log := handler.Logger
197	st := handler.Storage
198	pubsub := handler.PubSub
199
200	return func(next ssh.Handler) ssh.Handler {
201		return func(sesh ssh.Session) {
202			user, err := getUser(sesh, dbpool)
203			if err != nil {
204				sendutils.ErrorHandler(sesh, err)
205				return
206			}
207
208			args := sesh.Command()
209
210			opts := Cmd{
211				Session:     sesh,
212				User:        user,
213				Log:         log,
214				Dbpool:      dbpool,
215				Write:       false,
216				Storage:     st,
217				RegistryUrl: handler.RegistryUrl,
218			}
219
220			if len(args) == 0 {
221				next(sesh)
222				return
223			}
224
225			cmd := strings.TrimSpace(args[0])
226			if len(args) == 1 {
227				if cmd == "help" {
228					opts.help()
229					return
230				} else if cmd == "ls" {
231					err := opts.ls()
232					opts.bail(err)
233					return
234				} else {
235					next(sesh)
236					return
237				}
238			}
239
240			repoName := strings.TrimSpace(args[1])
241			cmdArgs := args[2:]
242			log.Info(
243				"imgs middleware detected command",
244				"args", args,
245				"cmd", cmd,
246				"repoName", repoName,
247				"cmdArgs", cmdArgs,
248			)
249
250			if cmd == "rm" {
251				rmCmd, write := flagSet("rm", sesh)
252				if !flagCheck(rmCmd, repoName, cmdArgs) {
253					return
254				}
255				opts.Write = *write
256
257				err := opts.rm(repoName)
258				opts.notice()
259				opts.bail(err)
260				return
261			} else if cmd == "sub" {
262				err = pubsub.Sub(sesh.Context(), uuid.NewString(), sesh, []*psub.Channel{
263					psub.NewChannel(fmt.Sprintf("%s/%s", user.Name, repoName)),
264				}, false)
265
266				if err != nil {
267					wish.Errorln(sesh, err)
268				}
269			} else {
270				next(sesh)
271				return
272			}
273		}
274	}
275}