Antonio Mika
·
10 Sep 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"
18 "github.com/picosh/pico/shared/storage"
19 "github.com/picosh/pico/tui/common"
20 sst "github.com/picosh/pobj/storage"
21 psub "github.com/picosh/pubsub"
22 "github.com/picosh/send/send/utils"
23)
24
25func getUser(s ssh.Session, dbpool db.DB) (*db.User, error) {
26 var err error
27 key, err := shared.KeyText(s)
28 if err != nil {
29 return nil, fmt.Errorf("key not found")
30 }
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 shared.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.Cfg
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 utils.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.PubSub.Sub(fmt.Sprintf("%s/%s", user.Name, repoName), &psub.Sub{
263 ID: uuid.NewString(),
264 Writer: sesh,
265 Done: make(chan struct{}),
266 Data: make(chan []byte),
267 })
268
269 if err != nil {
270 wish.Errorln(sesh, err)
271 }
272 } else {
273 next(sesh)
274 return
275 }
276 }
277 }
278}