Antonio Mika
·
2026-10-25
1package pico
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "fmt"
8 "log/slog"
9 "path/filepath"
10 "strings"
11 "time"
12
13 "github.com/picosh/pico/pkg/db"
14 "github.com/picosh/pico/pkg/pssh"
15 "github.com/picosh/pico/pkg/shared"
16
17 pipeLogger "github.com/picosh/utils/pipe/log"
18)
19
20func getUser(s *pssh.SSHServerConnSession, dbpool db.DB) (*db.User, error) {
21 if s.PublicKey() == nil {
22 return nil, fmt.Errorf("key not found")
23 }
24
25 key := shared.KeyForKeyText(s.PublicKey())
26
27 user, err := dbpool.FindUserByKey(s.User(), key)
28 if err != nil {
29 return nil, err
30 }
31
32 if user.Name == "" {
33 return nil, fmt.Errorf("must have username set")
34 }
35
36 return user, nil
37}
38
39type Cmd struct {
40 User *db.User
41 SshSession *pssh.SSHServerConnSession
42 Session shared.CmdSession
43 Log *slog.Logger
44 Dbpool db.DB
45 Write bool
46
47 // DomainTargets are the hosts dns-check expects a custom domain to point at.
48 DomainTargets []string
49
50 lookupTXT func(host, space string) ([]string, error)
51 lookupHost func(host string) (cname string, addrs []string, err error)
52}
53
54func (c *Cmd) output(out string) {
55 _, _ = c.Session.Write([]byte(out + "\r\n"))
56}
57
58func (c *Cmd) help() {
59 helpStr := "Commands: [help, user, logs, access-logs, chat, not-found, dns-check]\n"
60 helpStr += "help - this message\n"
61 helpStr += "user - display user information (returns name, id, account created, pico+ expiration)\n"
62 helpStr += "logs - stream user logs\n"
63 helpStr += "access-logs - fetch access logs from the last 30 days\n"
64 helpStr += "chat - IRC chat (must enable pty with `-t` to the SSH command)\n"
65 helpStr += "not-found - return all status 404 requests for a host (hostname.com [year|month])\n"
66 helpStr += "dns-check - check a domain's CNAME/A and TXT records for pgs, prose and tuns custom domains (hostname.com or *.hostname.com)\n"
67 c.output(helpStr)
68}
69
70func (c *Cmd) user() {
71 plus := ""
72 ff, _ := c.Dbpool.FindFeature(c.User.ID, "plus")
73 if ff != nil {
74 plus = ff.ExpiresAt.Format(time.RFC3339)
75 }
76 helpStr := fmt.Sprintf(
77 "%s\n%s\n%s\n%s",
78 c.User.Name,
79 c.User.ID,
80 c.User.CreatedAt.Format(time.RFC3339),
81 plus,
82 )
83 c.output(helpStr)
84}
85
86func (c *Cmd) notFound(host, interval string) error {
87 origin := shared.StartOfYear()
88 if interval == "month" {
89 origin = shared.StartOfMonth()
90 }
91 c.output(fmt.Sprintf("starting from: %s\n", origin.Format(time.RFC3339)))
92 urls, err := c.Dbpool.VisitUrlNotFound(&db.SummaryOpts{
93 Host: host,
94 UserID: c.User.ID,
95 Limit: 100,
96 Origin: origin,
97 })
98 if err != nil {
99 return err
100 }
101 for _, url := range urls {
102 c.output(fmt.Sprintf("%d %s", url.Count, url.Url))
103 }
104 return nil
105}
106
107func (c *Cmd) accessLogs(_ context.Context) error {
108 fromDate := time.Now().AddDate(0, 0, -30)
109 logs, err := c.Dbpool.FindAccessLogs(c.User.ID, &fromDate)
110 if err != nil {
111 return err
112 }
113
114 for _, log := range logs {
115 jsonData, err := json.Marshal(log)
116 if err != nil {
117 c.Log.Error("json marshal", "err", err)
118 continue
119 }
120 c.output(string(jsonData))
121 }
122 return nil
123}
124
125func (c *Cmd) logs(ctx context.Context) error {
126 conn := shared.NewPicoPipeClient()
127 stdoutPipe, err := pipeLogger.ReadLogs(ctx, c.Log, conn)
128
129 if err != nil {
130 return err
131 }
132
133 logChan := make(chan string)
134 defer close(logChan)
135
136 go func() {
137 for {
138 select {
139 case <-ctx.Done():
140 return
141 case log, ok := <-logChan:
142 if log == "" {
143 continue
144 }
145 if !ok {
146 return
147 }
148 _, _ = fmt.Fprintln(c.SshSession, log)
149 }
150 }
151 }()
152
153 scanner := bufio.NewScanner(stdoutPipe)
154 scanner.Buffer(make([]byte, 32*1024), 32*1024)
155 for scanner.Scan() {
156 line := scanner.Text()
157 parsedData := map[string]any{}
158
159 err := json.Unmarshal([]byte(line), &parsedData)
160 if err != nil {
161 c.Log.Error("json unmarshal", "err", err, "line", line, "hidden", true)
162 continue
163 }
164
165 user := shared.AnyToStr(parsedData, "user")
166 userId := shared.AnyToStr(parsedData, "userId")
167
168 hidden := shared.AnyToBool(parsedData, "hidden")
169
170 if !hidden && (user == c.User.Name || userId == c.User.ID) {
171 select {
172 case logChan <- line:
173 case <-ctx.Done():
174 return nil
175 default:
176 c.Log.Error("logChan is full, dropping log", "log", line)
177 continue
178 }
179 }
180 }
181 return scanner.Err()
182}
183
184type CliHandler struct {
185 DBPool db.DB
186 Logger *slog.Logger
187 DomainTargets []string
188}
189
190func Middleware(handler *CliHandler) pssh.SSHServerMiddleware {
191 dbpool := handler.DBPool
192 log := handler.Logger
193
194 return func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
195 return func(sesh *pssh.SSHServerConnSession) error {
196 args := sesh.Command()
197 if len(args) == 0 {
198 // Without a pty there is no TUI, and the upload middleware would wait on stdin.
199 if _, _, ok := sesh.Pty(); !ok {
200 opts := Cmd{Session: sesh}
201 opts.help()
202 return nil
203 }
204 return next(sesh)
205 }
206
207 user, err := getUser(sesh, dbpool)
208 if err != nil {
209 _, _ = fmt.Fprintf(sesh.Stderr(), "detected ssh command: %s\n", args)
210 s := fmt.Errorf("ERROR: you need to create an account before using the remote cli: %w", err)
211 sesh.Fatal(s)
212 return s
213 }
214
215 if len(args) > 0 && args[0] == "chat" {
216 _, _, hasPty := sesh.Pty()
217 if !hasPty {
218 err := fmt.Errorf(
219 "in order to render chat you need to enable PTY with the `ssh -t` flag",
220 )
221
222 sesh.Fatal(err)
223 return err
224 }
225
226 ff, err := dbpool.FindFeature(user.ID, "plus")
227 if err != nil {
228 handler.Logger.Error("Unable to find plus feature flag", "err", err, "user", user, "command", args)
229 ff, err = dbpool.FindFeature(user.ID, "bouncer")
230 if err != nil {
231 handler.Logger.Error("Unable to find bouncer feature flag", "err", err, "user", user, "command", args)
232 sesh.Fatal(err)
233 return err
234 }
235 }
236
237 if ff == nil {
238 err = fmt.Errorf("unable to find plus or bouncer feature flag")
239 sesh.Fatal(err)
240 return err
241 }
242
243 pass, err := dbpool.UpsertToken(user.ID, "pico-chat")
244 if err != nil {
245 sesh.Fatal(err)
246 return err
247 }
248 app, err := shared.NewSenpaiApp(sesh, user.Name, pass)
249 if err != nil {
250 sesh.Fatal(err)
251 return err
252 }
253 app.Run()
254 app.Close()
255 return err
256 }
257
258 opts := Cmd{
259 Session: sesh,
260 SshSession: sesh,
261 User: user,
262 Log: log,
263 Dbpool: dbpool,
264 Write: false,
265
266 DomainTargets: handler.DomainTargets,
267 }
268
269 cmd := strings.TrimSpace(args[0])
270 if len(args) == 1 {
271 switch cmd {
272 case "help":
273 opts.help()
274 return nil
275 case "user":
276 opts.user()
277 return nil
278 case "logs":
279 err = opts.logs(sesh.Context())
280 if err != nil {
281 sesh.Fatal(err)
282 }
283 return nil
284 case "dns-check":
285 sesh.Fatal(fmt.Errorf("must provide a domain name (`dns-check example.com`)"))
286 return nil
287 // access_logs is the command's old name.
288 case "access-logs", "access_logs":
289 err = opts.accessLogs(sesh.Context())
290 if err != nil {
291 sesh.Fatal(err)
292 }
293 return nil
294 }
295 }
296
297 if cmd == "dns-check" {
298 err = opts.dnsCheck(args[1])
299 if err != nil {
300 sesh.Fatal(err)
301 }
302 return nil
303 }
304
305 if cmd == "not-found" {
306 if len(args) < 3 {
307 sesh.Fatal(fmt.Errorf("must provide host name and interval (`month` or `year`)"))
308 return nil
309 }
310 err = opts.notFound(args[1], args[2])
311 if err != nil {
312 sesh.Fatal(err)
313 }
314 return nil
315 }
316
317 if passthrough(cmd) {
318 return next(sesh)
319 }
320
321 opts.help()
322 sesh.Fatal(fmt.Errorf("unknown command: %s", cmd))
323 return nil
324 }
325 }
326}
327
328// passthrough reports whether a later middleware handles cmd. Anything else
329// would reach the upload middleware, which waits on stdin.
330func passthrough(cmd string) bool {
331 switch cmd {
332 case "scp", "rsync", "command":
333 return true
334 }
335 return filepath.Base(cmd) == "authorized_keys"
336}