package app
import (
"cmp"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"os"
"slices"
"strings"
"time"
"github.com/sahilm/fuzzy"
"github.com/samber/lo"
"golang.org/x/term"
"tobi/internal/remote"
"tobi/internal/session"
"tobi/internal/target"
"tobi/internal/termcolor"
)
func Find(ctx context.Context, query string) []session.Session {
return findSessions(ctx, query, remote.DiscoverPeers(ctx))
}
func findSessions(ctx context.Context, query string, peers []remote.Peer) []session.Session {
raw := strings.TrimSpace(query)
q := strings.ToLower(raw)
t := target.Parse(raw)
id := strings.ToLower(t.ID)
matchID := func(list []session.Session) []session.Session {
return lo.Filter(list, func(s session.Session, _ int) bool { return strings.HasPrefix(strings.ToLower(s.ID), id) })
}
if t.ID != "" || t.User != "" {
if session.IsLocal(t.Host) {
return matchID(session.ListLocal())
}
peer := lo.ContainsBy(peers, func(p remote.Peer) bool { return strings.EqualFold(p.Name, t.Host) }) || len(remote.KnownDests(t.Host)) > 1
if peer {
if session.ValidateID(t.ID) {
return []session.Session{{Host: t.Dest(), ID: t.ID}}
}
return matchID(remote.QueryHost(ctx, t.Dest()))
}
}
all := session.ListLocal()
if match, ok := lo.Find(all, func(s session.Session) bool {
return q != "" && (strings.EqualFold(s.ID, q) || strings.EqualFold(s.Canonical(), q))
}); ok {
return []session.Session{match}
}
online := lo.Filter(peers, func(p remote.Peer, _ int) bool { return p.Online })
if q != "" {
if named := lo.Filter(online, func(p remote.Peer, _ int) bool { return strings.HasPrefix(strings.ToLower(p.Name), q) }); len(named) > 0 {
online = named
}
}
lists := make(chan []session.Session, len(online))
for _, p := range online {
go func() {
pCtx, cancel := context.WithTimeout(ctx, 1200*time.Millisecond)
defer cancel()
lists <- remote.QueryHost(pCtx, p.Name)
}()
}
for range online {
all = append(all, <-lists...)
}
byNewest(all)
return rankSessions(all, q, true)
}
func rankSessions(list []session.Session, query string, directFirst bool) []session.Session {
q := strings.ToLower(strings.TrimSpace(query))
if q == "" {
return list
}
if directFirst {
direct := lo.Filter(list, func(s session.Session, _ int) bool {
return strings.HasPrefix(strings.ToLower(s.Canonical()), q) || strings.HasPrefix(strings.ToLower(s.ID), q) ||
strings.HasPrefix(strings.ToLower(s.Command), q) || strings.Contains(strings.ToLower(s.Cwd), q)
})
if len(direct) > 0 {
return lo.UniqBy(direct, (session.Session).Canonical)
}
}
matches := fuzzy.Find(q, lo.Map(list, func(s session.Session, _ int) string {
return fmt.Sprintf("%s %s %s %s", s.Host, s.Canonical(), s.Command, s.Cwd)
}))
return lo.UniqBy(lo.Map(matches, func(m fuzzy.Match, _ int) session.Session { return list[m.Index] }), (session.Session).Canonical)
}
func byNewest(s []session.Session) {
slices.SortFunc(s, func(a, b session.Session) int {
if c := b.StartedAt.Compare(a.StartedAt); c != 0 {
return c
}
return strings.Compare(a.ID, b.ID)
})
}
func connectOrHop(dst session.Session) error {
if cs := os.Getenv("TOBI_CONTROL"); cs != "" {
return session.SendHop(cs, dst.Canonical())
}
return RunHopLoop(dst)
}
func Interactive(query string) error {
if q := strings.TrimSpace(query); !isTTY() || strings.Contains(q, ".") {
if sessions := Find(context.Background(), q); len(sessions) > 0 {
return connectOrHop(sessions[0])
}
return errors.New("no sessions found")
}
return runPicker(session.ListLocal(), query)
}
func suggest(err error, q string) error {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
alt := lo.Filter(Find(ctx, q), func(s session.Session, _ int) bool { return !strings.EqualFold(s.Canonical(), q) })
if len(alt) > 0 {
return fmt.Errorf("%w (did you mean %s?)", err, alt[0].Canonical())
}
return err
}
func RunHopLoop(dst session.Session) error {
if rawState, err := term.MakeRaw(int(os.Stdin.Fd())); err == nil {
defer term.Restore(int(os.Stdin.Fd()), rawState)
}
inputCh := make(chan []byte, 32)
probed := make(chan map[int]termcolor.Color, 1)
go func() {
defer close(inputCh)
reads := make(chan stdinRead)
go func() {
buf := make([]byte, 4096)
for {
n, err := os.Stdin.Read(buf)
reads <- stdinRead{slices.Clone(buf[:n]), err}
if err != nil {
return
}
}
}()
colors := probeTerminal(reads, func(b []byte) { inputCh <- b })
probed <- colors
for r := range reads {
if len(r.b) > 0 {
inputCh <- r.b
}
if r.err != nil {
return
}
}
}()
colors := <-probed
dial := func(s session.Session) (io.ReadWriteCloser, error) {
if session.IsLocal(s.Host) {
return net.Dial("unix", session.SocketPath(s.ID))
}
return remote.AttachRemote(s.Host, s.ID)
}
current := dst
for {
stream, err := dial(current)
if err != nil {
return suggest(fmt.Errorf("connect %s: %w", current.Canonical(), err), current.Canonical())
}
if err = session.Attach(stream, inputCh, colors); err == nil || err == session.ErrDetached {
return nil
}
hop, ok := err.(session.HopRequest)
if !ok {
return err
}
t := target.Parse(hop.Target)
if t.Host == "" || !session.ValidateID(t.ID) {
return suggest(fmt.Errorf("invalid hop target: %s", hop.Target), hop.Target)
}
current = session.Session{Host: t.Dest(), ID: t.ID}
}
}
type stdinRead struct {
b []byte
err error
}
// probeTerminal asks the real terminal for its colors and collects replies
// until the da1 sentinel arrives or the window closes. it runs inside the
// stdin reader goroutine before any attach, so replies can never reach a
// session as input. keys typed during the probe are queued for the input
// loop. it returns nil and skips the probe when stdin or stdout is not a
// terminal.
func probeTerminal(reads <-chan stdinRead, push func([]byte)) map[int]termcolor.Color {
if !term.IsTerminal(int(os.Stdin.Fd())) || !term.IsTerminal(int(os.Stdout.Fd())) {
return nil
}
return runProbe(reads, os.Stdout, push, probeWindow)
}
const probeWindow = 300 * time.Millisecond
func runProbe(reads <-chan stdinRead, w io.Writer, push func([]byte), window time.Duration) map[int]termcolor.Color {
if _, err := w.Write([]byte(termcolor.ProbeQuery)); err != nil {
return nil
}
var acc []byte
timer := time.NewTimer(window)
defer timer.Stop()
for {
select {
case r := <-reads:
acc = append(acc, r.b...)
if end := termcolor.DA1End(acc); end >= 0 {
if end < len(acc) {
push(acc[end:])
}
return termcolor.ParseReplies(acc[:end])
}
if r.err != nil {
return termcolor.ParseReplies(acc)
}
case <-timer.C:
return termcolor.ParseReplies(acc)
}
}
}
func Kill(ctx context.Context, s session.Session) error {
if session.IsLocal(s.Host) {
return session.Kill(s.ID)
}
return remote.KillRemote(ctx, s.Host, s.ID)
}
func Spawn(addr string, cmdArgs []string) error {
if len(cmdArgs) == 0 {
return errors.New("usage: tobi [host.[id]] -- <command> [args...]")
}
t := target.Parse(addr)
t.Host = cmp.Or(t.Host, strings.ToLower(session.ShortHostname()))
id := cmp.Or(t.ID, session.GenerateID())
if !session.ValidateID(id) {
return fmt.Errorf("invalid session id: %q (must be 1-64 alphanumeric, - or _)", id)
}
if t.User == "" && session.IsLocal(t.Host) {
if err := session.EnsureOwner(id, cmdArgs); err != nil {
return err
}
return connectOrHop(session.Session{ID: id, Host: session.ShortHostname()})
}
if err := remote.SpawnRemote(t.Dest(), id, cmdArgs); err != nil {
return err
}
return connectOrHop(session.Session{ID: id, Host: t.Dest()})
}
func Complete(rawPrefix string) error {
ctx, cancel := context.WithTimeout(context.Background(), 800*time.Millisecond)
defer cancel()
peers := remote.DiscoverPeers(ctx)
candidates := append(findSessions(ctx, rawPrefix, peers), hostCandidates(rawPrefix, peers)...)
return json.NewEncoder(os.Stdout).Encode(lo.Map(candidates, func(s session.Session, _ int) map[string]string {
desc := fmt.Sprintf("%-12s %s", s.Command, s.Cwd)
if s.ID == "" {
desc = "host"
}
return map[string]string{"value": s.Canonical(), "description": desc}
}))
}
func hostCandidates(prefix string, peers []remote.Peer) []session.Session {
t := target.Parse(prefix)
q, user := strings.ToLower(t.Host), lo.Ternary(t.User != "", t.User+"@", "")
return lo.FilterMap(
append([]remote.Peer{{Name: session.ShortHostname()}}, peers...),
func(p remote.Peer, _ int) (session.Session, bool) {
return session.Session{Host: user + p.Name}, strings.HasPrefix(strings.ToLower(p.Name), q)
},
)
}