package main import ( "crypto/hmac" "crypto/rand" "crypto/sha1" "crypto/subtle" "encoding/base32" "encoding/binary" "encoding/hex" "encoding/json" "flag" "fmt" "log" "net" "net/http" "net/url" "os" "path/filepath" "strings" "sync" "time" "github.com/gorilla/websocket" "golang.org/x/crypto/bcrypt" "wxch.org/wxch/catalog" "wxch.org/wxch/deploy" "wxch.org/wxch/internal/modules" "wxch.org/wxch/internal/protocol" "wxch.org/wxch/web" ) type Config struct { Listen string `json:"listen"` PasswordHash string `json:"password_hash"` TOTPSecret string `json:"totp_secret,omitempty"` // optional 2FA: empty = password only AgentTokens map[string]string `json:"agent_tokens"` // agent name -> token PublicURL string `json:"public_url"` // e.g. https://control.example.org (used in printed install commands) CCTVUpstream string `json:"cctv_upstream,omitempty"` // verified loopback tunnel endpoint only } func randHex(n int) string { b := make([]byte, n) rand.Read(b) return hex.EncodeToString(b) } func totp(secret string, t time.Time) string { key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(secret) if err != nil { return "" } m := hmac.New(sha1.New, key) var c [8]byte binary.BigEndian.PutUint64(c[:], uint64(t.Unix()/30)) m.Write(c[:]) s := m.Sum(nil) o := s[len(s)-1] & 0x0f v := (binary.BigEndian.Uint32(s[o:o+4]) & 0x7fffffff) % 1000000 return fmt.Sprintf("%06d", v) } func totpOK(secret, code string) bool { now := time.Now() for _, d := range []time.Duration{-30 * time.Second, 0, 30 * time.Second} { if subtle.ConstantTimeCompare([]byte(totp(secret, now.Add(d))), []byte(code)) == 1 { return true } } return false } // initConfig writes a new config. 2FA is off unless withTOTP is set. func initConfig(path, password string, withTOTP bool) { if _, err := os.Stat(path); err == nil { log.Fatalf("%s already exists; refusing to overwrite", path) } if password == "" { password = randHex(10) fmt.Println("generated password:", password) } if !passwordOK(password) { log.Fatalf("the password must be %d-72 characters", minPasswordLen) } h, _ := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost) cfg := Config{ Listen: "127.0.0.1:8080", PasswordHash: string(h), AgentTokens: map[string]string{"pgx": randHex(24)}, } if withTOTP { sb := make([]byte, 20) rand.Read(sb) cfg.TOTPSecret = base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(sb) } b, _ := json.MarshalIndent(cfg, "", " ") if err := os.WriteFile(path, b, 0600); err != nil { log.Fatal(err) } if cfg.TOTPSecret != "" { fmt.Println("totp secret (add to authenticator app):", cfg.TOTPSecret) fmt.Printf("otpauth://totp/wxch:control?secret=%s&issuer=wxch\n", cfg.TOTPSecret) } fmt.Println("agent token for pgx:", cfg.AgentTokens["pgx"]) fmt.Println("config written to", path) } // ---- hub ---- type conn struct { ws *websocket.Conn mu sync.Mutex sink func(Msg) // set for in-process callers (admin exec) instead of a websocket } func (c *conn) send(v any) { if c.sink != nil { if m, ok := v.(Msg); ok { c.sink(m) } return } c.mu.Lock() defer c.mu.Unlock() c.ws.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.ws.WriteJSON(v); err != nil { c.ws.Close() } } // Msg is the websocket envelope shared by the UI, the server and the agents. type Msg = protocol.Msg type rpcEntry struct { ui *conn uiID string agent string } type session struct { ui *conn uiSid string agent string } type Hub struct { mu sync.Mutex agents map[string]*conn uis map[*conn]bool sessions map[string]*session // global sid -> session rpcs map[string]*rpcEntry // global rpc id -> caller stats map[string]*statBuf // latest telemetry + short history per machine infers map[string]*inferJob // gateway requests in flight, by id inflight map[string]int // gateway requests in flight per machine } func (h *Hub) agentList() []string { l := []string{} for n := range h.agents { l = append(l, n) } return l } func (h *Hub) broadcastAgents() { h.mu.Lock() msg := Msg{T: "agents", List: h.agentList()} var us []*conn for u := range h.uis { us = append(us, u) } h.mu.Unlock() for _, u := range us { u.send(msg) } } // ---- auth ---- type Server struct { chatSess *sessionSet // login sessions: kept in a file, see sessions.go studioSess *sessionSet comfySess *sessionSet cctvSess *sessionSet noteSess *sessionSet talk *talkStore talkTasks *talkTaskStore feedback *feedbackStore // suggestions written on note.* seen *seenStore // when each machine last connected and last dropped, see seen.go tuns tunRegistry // pages of engines carried over the agent connections, see comfy.go jobs *JobStore // AI module jobs, see jobs.go keys *KeyStore recent recentLog npSess *sessionSet cfgPath string cfg Config hub *Hub mu sync.Mutex sessions *sessionSet // control page sessions fails map[string][]time.Time } func clientIP(r *http.Request) string { if f := r.Header.Get("X-Forwarded-For"); f != "" { return strings.TrimSpace(strings.Split(f, ",")[0]) } h, _, _ := net.SplitHostPort(r.RemoteAddr) return h } func (s *Server) locked(ip string) bool { s.mu.Lock() defer s.mu.Unlock() var keep []time.Time for _, t := range s.fails[ip] { if time.Since(t) < 15*time.Minute { keep = append(keep, t) } } s.fails[ip] = keep return len(keep) >= 5 } func (s *Server) fail(ip string) { s.mu.Lock() s.fails[ip] = append(s.fails[ip], time.Now()) s.mu.Unlock() } func (s *Server) authed(r *http.Request) bool { c, err := r.Cookie("wxch_session") if err != nil { return false } return s.sessions.valid(c.Value) } // renewUI extends the caller's session (and its cookie) when it is getting old: every visit renews it, so it only ends // after uiSessionTTL without use. func (s *Server) renewUI(w http.ResponseWriter, r *http.Request) { if c, err := r.Cookie("wxch_session"); err == nil && s.sessions.renew(c.Value, uiSessionTTL) { http.SetCookie(w, &http.Cookie{Name: "wxch_session", Value: c.Value, Path: "/", HttpOnly: true, Secure: true, SameSite: http.SameSiteStrictMode, MaxAge: int(uiSessionTTL.Seconds())}) } } func (s *Server) login(w http.ResponseWriter, r *http.Request) { ip := clientIP(r) if s.locked(ip) { http.Error(w, "too many attempts, try later", http.StatusTooManyRequests) return } var req struct{ Password, Code string } if json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&req) != nil { http.Error(w, "bad request", 400) return } pwOK := bcrypt.CompareHashAndPassword([]byte(s.pwHash()), []byte(req.Password)) == nil codeOK := s.cfg.TOTPSecret == "" || totpOK(s.cfg.TOTPSecret, req.Code) // 2FA only when a secret is configured if !pwOK || !codeOK { s.fail(ip) log.Printf("login failed from %s", ip) time.Sleep(time.Second) // slows guessing on top of the 5-per-15-minutes lockout http.Error(w, "invalid credentials", http.StatusUnauthorized) return } tok := randHex(32) s.sessions.add(tok, uiSessionTTL) http.SetCookie(w, &http.Cookie{Name: "wxch_session", Value: tok, Path: "/", HttpOnly: true, Secure: true, SameSite: http.SameSiteStrictMode, MaxAge: int(uiSessionTTL.Seconds())}) log.Printf("login ok from %s", ip) w.Write([]byte("ok")) } var upgrader = websocket.Upgrader{ ReadBufferSize: 64 * 1024, WriteBufferSize: 64 * 1024, CheckOrigin: func(r *http.Request) bool { o := r.Header.Get("Origin") if o == "" { return true // non-browser (agent) } u, err := url.Parse(o) return err == nil && u.Host == r.Host }, } // ---- websockets ---- func (s *Server) wsAgent(w http.ResponseWriter, r *http.Request) { tok := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") name := "" tokens := s.cfg.AgentTokens if b, err := os.ReadFile(s.cfgPath); err == nil { var c Config if json.Unmarshal(b, &c) == nil { tokens = c.AgentTokens } } for n, t := range tokens { if subtle.ConstantTimeCompare([]byte(t), []byte(tok)) == 1 && tok != "" { name = n } } if name == "" { log.Printf("agent auth failed from %s", clientIP(r)) http.Error(w, "unauthorized", 401) return } ws, err := upgrader.Upgrade(w, r, nil) if err != nil { return } c := &conn{ws: ws} h := s.hub h.mu.Lock() if old := h.agents[name]; old != nil { old.ws.Close() } h.agents[name] = c h.mu.Unlock() log.Printf("agent %s connected", name) s.seen.up(name) h.broadcastAgents() ws.SetReadLimit(48 << 20) // engine page traffic (previews, files) travels in chunks, but a websocket frame is one message ws.SetReadDeadline(time.Now().Add(60 * time.Second)) ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(60 * time.Second)); return nil }) go func() { for range time.Tick(20 * time.Second) { c.mu.Lock() err := ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second)) c.mu.Unlock() if err != nil { return } } }() for { var m Msg if err := ws.ReadJSON(&m); err != nil { break } ws.SetReadDeadline(time.Now().Add(60 * time.Second)) if strings.HasPrefix(m.T, "tun_") { s.tuns.deliver(name, m) continue } if strings.HasPrefix(m.T, "job_") { s.jobs.onAgentMsg(name, m) continue } if strings.HasPrefix(m.T, "infer_") { h.mu.Lock() job := h.infers[m.ID] h.mu.Unlock() if job != nil { select { case job.ch <- m: case <-job.done: } } continue } if m.T == "stats" { h.recordStats(name, m.Data) continue } if m.T == "rpc_result" || m.T == "rpc_event" { h.mu.Lock() e := h.rpcs[m.ID] if e != nil && m.T == "rpc_result" { delete(h.rpcs, m.ID) } h.mu.Unlock() if e != nil { m.ID = strings.TrimPrefix(m.ID, e.uiID+":") e.ui.send(m) } continue } h.mu.Lock() sess := h.sessions[m.Sid] h.mu.Unlock() if sess == nil { continue } uiSid := sess.uiSid m.Sid = uiSid sess.ui.send(m) if m.T == "exit" { h.mu.Lock() for k, v := range h.sessions { if v == sess { delete(h.sessions, k) } } h.mu.Unlock() } } // agent gone h.mu.Lock() wasCurrent := h.agents[name] == c // a connection replaced by a newer one is not the machine dropping if wasCurrent { delete(h.agents, name) } var dead []*session for k, v := range h.sessions { if v.agent == name { dead = append(dead, v) delete(h.sessions, k) } } var deadRPC []*rpcEntry var deadRPCIDs []string for k, v := range h.rpcs { if v.agent == name { deadRPC = append(deadRPC, v) deadRPCIDs = append(deadRPCIDs, k) delete(h.rpcs, k) } } h.mu.Unlock() for _, d := range dead { d.ui.send(Msg{T: "exit", Sid: d.uiSid, Err: "agent_lost"}) // the page uses it to follow an installer across an agent restart } var deadJobs []*inferJob h.mu.Lock() for _, j := range h.infers { if j.agent == name { deadJobs = append(deadJobs, j) } } h.mu.Unlock() for _, j := range deadJobs { select { case j.ch <- Msg{T: "infer_end", Err: "machine disconnected"}: default: } } for i, d := range deadRPC { d.ui.send(Msg{T: "rpc_result", ID: strings.TrimPrefix(deadRPCIDs[i], d.uiID+":"), Err: "agent disconnected"}) } s.jobs.machineGone(name) s.tuns.agentGone(name) ws.Close() log.Printf("agent %s disconnected", name) if wasCurrent { s.seen.down(name) } h.broadcastAgents() } func (s *Server) wsUI(w http.ResponseWriter, r *http.Request) { if !s.authed(r) { http.Error(w, "unauthorized", 401) return } ws, err := upgrader.Upgrade(w, r, nil) if err != nil { return } c := &conn{ws: ws} h := s.hub h.mu.Lock() h.uis[c] = true h.mu.Unlock() ws.SetReadLimit(1 << 20) c.send(Msg{T: "agents", List: func() []string { h.mu.Lock(); defer h.mu.Unlock(); return h.agentList() }()}) h.sendStatsSnapshot(c) uiID := randHex(4) for { var m Msg if err := ws.ReadJSON(&m); err != nil { break } gsid := uiID + ":" + m.Sid switch m.T { case "open": h.mu.Lock() a := h.agents[m.Agent] if a != nil { h.sessions[gsid] = &session{ui: c, uiSid: m.Sid, agent: m.Agent} } h.mu.Unlock() if a == nil { c.send(Msg{T: "exit", Sid: m.Sid}) continue } if len(m.Script) > 1<<18 { c.send(Msg{T: "exit", Sid: m.Sid}) continue } if m.Script != "" { log.Printf("script run on %s (%s): %s", m.Agent, gsid, m.Script) } else { log.Printf("terminal opened on %s (%s)", m.Agent, gsid) } a.send(Msg{T: "open", Sid: gsid, Cols: m.Cols, Rows: m.Rows, Script: m.Script}) case "rpc", "update": h.mu.Lock() a := h.agents[m.Agent] gid := uiID + ":" + m.ID if a != nil { h.rpcs[gid] = &rpcEntry{ui: c, uiID: uiID, agent: m.Agent} } h.mu.Unlock() if a == nil { c.send(Msg{T: "rpc_result", ID: m.ID, Err: "machine offline"}) continue } log.Printf("%s %s on %s", m.T, m.Method, m.Agent) a.send(Msg{T: m.T, ID: gid, Method: m.Method, Params: m.Params}) case "in", "resize", "close": h.mu.Lock() sess := h.sessions[gsid] var a *conn if sess != nil { a = h.agents[sess.agent] } h.mu.Unlock() if a != nil { m.Sid = gsid a.send(m) } } } // ui gone: close its sessions h.mu.Lock() delete(h.uis, c) type pair struct { a *conn sid string } var ps []pair for k, v := range h.sessions { if v.ui == c { ps = append(ps, pair{h.agents[v.agent], k}) delete(h.sessions, k) } } for k, v := range h.rpcs { if v.ui == c { delete(h.rpcs, k) } } h.mu.Unlock() for _, p := range ps { if p.a != nil { p.a.send(Msg{T: "close", Sid: p.sid}) } } ws.Close() } func main() { cfgPath := flag.String("config", "config.json", "config file") doInit := flag.Bool("init", false, "create config with random secrets") execAgent := flag.String("exec", "", "run -script on this machine via the local admin socket and print the output") execScript := flag.String("script", "", "script for -exec (reads stdin if empty)") newKey := flag.String("new-key", "", "create a gateway API key with this name (prints it once)") showStats := flag.Bool("stats", false, "print the latest telemetry of every machine (via the local admin socket)") rpcMethod := flag.String("rpc", "", "call this agent method on -exec machine via the admin socket (JSON via -params)") rpcParams := flag.String("params", "", "JSON params for -rpc") execTimeout := flag.Int("timeout", 120, "seconds before -exec gives up") addAgent := flag.String("add-agent", "", "add a machine with this name and print its install command") pw := flag.String("password", "", "password for -init (random if empty)") withTOTP := flag.Bool("totp", false, "with -init: also require a 2FA code (authenticator app) at login") flag.Parse() if *doInit { initConfig(*cfgPath, *pw, *withTOTP) return } if *newKey != "" { adminNewKeyClient(*cfgPath, *newKey) return } if *showStats { adminStatsClient(*cfgPath) return } if *execAgent != "" { if *rpcMethod != "" { adminRPCClient(*cfgPath, *execAgent, *rpcMethod, *rpcParams, *execTimeout) return } adminExecClient(*cfgPath, *execAgent, *execScript, *execTimeout) return } if *addAgent != "" { addAgentCmd(*cfgPath, *addAgent) return } var cfg Config b, err := os.ReadFile(*cfgPath) if err != nil { log.Fatal(err) } if err := json.Unmarshal(b, &cfg); err != nil { log.Fatal(err) } sess := newSessionStore(filepath.Join(filepath.Dir(*cfgPath), "sessions.json")) s := &Server{chatSess: sess.set("chat"), keys: newKeyStore(filepath.Join(filepath.Dir(*cfgPath), "keys.json")), npSess: sess.set("np"), studioSess: sess.set("studio"), comfySess: sess.set("comfy"), noteSess: sess.set("note"), feedback: newFeedbackStore(filepath.Join(filepath.Dir(*cfgPath), "feedback.json")), talk: newTalkStore(filepath.Join(filepath.Dir(*cfgPath), "talk.json")), seen: newSeenStore(filepath.Join(filepath.Dir(*cfgPath), "machines_seen.json")), cfgPath: *cfgPath, cfg: cfg, sessions: sess.set("ui"), fails: map[string][]time.Time{}, hub: &Hub{agents: map[string]*conn{}, uis: map[*conn]bool{}, sessions: map[string]*session{}, rpcs: map[string]*rpcEntry{}, stats: map[string]*statBuf{}, infers: map[string]*inferJob{}, inflight: map[string]int{}}} s.talkTasks, err = newTalkTaskStore(filepath.Join(filepath.Dir(*cfgPath), "talk_tasks.json")) if err != nil { log.Fatal(err) } modCat, err := modules.Load() if err != nil { log.Fatal(err) } s.jobs = newJobStore(filepath.Join(filepath.Dir(*cfgPath), "jobs"), modCat) go s.jobs.run(s.hub) sub := web.FS() mux := http.NewServeMux() mux.Handle("/", pageHandler(http.FileServer(http.FS(sub)))) mux.HandleFunc("/catalog.json", func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Cache-Control", "no-cache") w.Write(catalog.JSON) }) mux.HandleFunc("/modules.json", func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Cache-Control", "no-cache") w.Write(catalog.ModulesJSON) }) registerInstallers(mux, deploy.FS(), cfg.PublicURL) mux.HandleFunc("/api/login", s.login) mux.HandleFunc("/api/auth-info", s.authInfo) mux.HandleFunc("/api/password", s.changePassword) mux.HandleFunc("/api/keys", s.keysAPI) mux.HandleFunc("/api/keys/", s.keysAPI) mux.HandleFunc("/api/talk/", s.talkOwnerAPI) mux.HandleFunc("/api/me", func(w http.ResponseWriter, r *http.Request) { if !s.authed(r) { http.Error(w, "no", 401) return } s.renewUI(w, r) w.Write([]byte("ok")) }) mux.HandleFunc("/api/agent-cmd", s.agentCmd) mux.HandleFunc("/api/agent/job-file", s.agentJobFile) mux.HandleFunc("/api/agents-registered", s.agentsRegistered) mux.HandleFunc("/api/machines-seen", s.machinesSeen) mux.Handle("/dl/", http.StripPrefix("/dl/", http.FileServer(http.Dir("dl")))) mux.HandleFunc("/ws/ui", s.wsUI) mux.HandleFunc("/ws/agent", s.wsAgent) s.startAdmin() np := s.notepadMux(sub) note := s.noteMux(sub, web.PrivateFS()) gw := s.gatewayMux() chat := s.chatMux(sub) studio := s.studioMux(sub) comfy := s.comfyMux() s.cctvSess = sess.set("cctv") cctv, err := s.cctvMux(cfg.CCTVUpstream) if err != nil { log.Fatal(err) } talk := s.talkMux() go s.keys.flushLoop() root := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch appForHost(r.Host) { case "notepad": np.ServeHTTP(w, r) case "note": note.ServeHTTP(w, r) case "api": gw.ServeHTTP(w, r) case "chat": chat.ServeHTTP(w, r) case "studio": studio.ServeHTTP(w, r) case "comfy": comfy.ServeHTTP(w, r) case "talk": talk.ServeHTTP(w, r) case "cctv": cctv.ServeHTTP(w, r) default: mux.ServeHTTP(w, r) } }) log.Printf("listening on %s", cfg.Listen) log.Fatal(http.ListenAndServe(cfg.Listen, root)) } func addAgentCmd(path, name string) { for _, r := range name { if !(r >= 'a' && r <= 'z' || r >= '0' && r <= '9' || r == '-') { log.Fatal("name must be lowercase letters, digits or '-'") } } b, err := os.ReadFile(path) if err != nil { log.Fatal(err) } var cfg Config if err := json.Unmarshal(b, &cfg); err != nil { log.Fatal(err) } if _, ok := cfg.AgentTokens[name]; ok { log.Fatalf("agent %q already exists (delete it from %s to regenerate)", name, path) } tok := randHex(24) cfg.AgentTokens[name] = tok out, _ := json.MarshalIndent(cfg, "", " ") if err := os.WriteFile(path, out, 0600); err != nil { log.Fatal(err) } fmt.Printf("added %s. Run this once on that machine (it asks for the sudo password a single time):\n\ncurl -fsSL %s/get | WXCH_TOKEN=%s sh\n", name, baseOrPlaceholder(cfg.PublicURL), tok) }