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