package hammerhead

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"embed"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.veillax.com/Veillax/reef/model"
	"golang.org/x/crypto/ssh"
	"golang.org/x/net/websocket"
)

// The web terminal opens a shell on a site from the dashboard. Hammerhead is the SSH client: it
// signs in to the SSH port that an edge exposes for the site (a TCP resource whose target is
// Bamboo's built-in SSH server) with a key of its own, and relays the shell to the browser over a
// WebSocket. Bamboo only lets in keys listed in its authorized_keys, so a site is reachable this
// way only after someone puts Hammerhead's public key there.

//go:embed web/term
var termFiles embed.FS

const (
	sshTargetPort    = 2222 // where Bamboo's SSH server listens by default
	termIdle         = 30 * time.Minute
	termMaxSessions  = 8
	termMaxFrame     = 64 << 10
	termDialTimeout  = 10 * time.Second
	termDefaultTerm  = "xterm-256color"
	termResizeMaxCol = 500
	termResizeMaxRow = 200
)

// TerminalSettings is what the web terminal remembers: its own SSH key, and the host key each
// site showed the first time (trusted on first use, like an ordinary ssh client).
type TerminalSettings struct {
	Key   string            `json:"key,omitempty"`   // PKCS#8 PEM private key; sealed at rest, never sent to the browser
	Hosts map[string]string `json:"hosts,omitempty"` // site ID -> host key fingerprint
}

// terminalStatic serves the embedded xterm.js files.
func terminalStatic() http.Handler {
	sub, _ := fs.Sub(termFiles, "web/term")
	fileServer := http.StripPrefix("/static/term/", http.FileServer(http.FS(sub)))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "public, max-age=86400")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		fileServer.ServeHTTP(w, r)
	})
}

// terminalSigner returns Hammerhead's SSH key, making it on first use. Caller holds s.mu.
func (s *Server) terminalSignerLocked() (ssh.Signer, error) {
	if s.st.Settings.Terminal.Key == "" {
		_, priv, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return nil, err
		}
		der, err := x509.MarshalPKCS8PrivateKey(priv)
		if err != nil {
			return nil, err
		}
		s.st.Settings.Terminal.Key = string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der}))
		s.save()
	}
	return ssh.ParsePrivateKey([]byte(s.st.Settings.Terminal.Key))
}

// terminalKey tells an administrator which public key to put in a site's authorized_keys.
func (s *Server) terminalKey(w http.ResponseWriter, r *http.Request, c *orgCtx) {
	s.mu.Lock()
	signer, err := s.terminalSignerLocked()
	s.mu.Unlock()
	if err != nil {
		fail(w, 500, err)
		return
	}
	line := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(signer.PublicKey()))) + " hammerhead"
	writeJSON(w, 200, map[string]any{"publicKey": line, "fingerprint": ssh.FingerprintSHA256(signer.PublicKey())})
}

// isSSHResource reports whether m is a TCP port that leads to the SSH server of the given site.
func isSSHResource(m *model.Mapping, siteID string) bool {
	if m.Proto != model.ProtoTCP || m.IsRedirect() || !m.Enabled {
		return false
	}
	for _, t := range m.Targets {
		if t.SiteID != siteID {
			continue
		}
		if t.Port == sshTargetPort {
			return true
		}
		for _, tag := range m.Tags {
			if tag == "ssh" {
				return true
			}
		}
	}
	return false
}

// terminalAddr finds where to connect for a site: the public address of the SSH port that
// leads to it. A port bound to every address is reached through the edge's own address.
func (s *Server) terminalAddr(c *orgCtx, siteID string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	st := s.findSite(siteID)
	if st == nil || st.OrgID != c.Org.ID || st.Local {
		return "", errors.New("no such site")
	}
	for _, m := range s.st.Mappings {
		if m.OrgID != c.Org.ID || !isSSHResource(m, siteID) {
			continue
		}
		host := m.BindIP
		if ip := net.ParseIP(host); host == "" || (ip != nil && ip.IsUnspecified()) {
			e := s.findEdge(m.EdgeID)
			if e == nil {
				continue
			}
			h, _, err := net.SplitHostPort(e.TunnelAddr)
			if err != nil || h == "" {
				continue
			}
			host = h
		}
		return net.JoinHostPort(host, strconv.Itoa(m.Port)), nil
	}
	return "", errors.New("This site has no SSH port yet. Open one from the SSH button on the Sites page.")
}

type hostKeyChanged struct{ site, was, now string }

func (e *hostKeyChanged) Error() string {
	return fmt.Sprintf("The host key of %s has changed (was %s, now %s). If you reinstalled Bamboo on it, forget the saved key and try again.", e.site, e.was, e.now)
}

// trustHostKey remembers a site's host key the first time and refuses a different one afterwards.
func (s *Server) trustHostKey(siteID, siteName string) ssh.HostKeyCallback {
	return func(_ string, _ net.Addr, key ssh.PublicKey) error {
		fp := ssh.FingerprintSHA256(key)
		s.mu.Lock()
		defer s.mu.Unlock()
		t := &s.st.Settings.Terminal
		switch known := t.Hosts[siteID]; {
		case known == "":
			if t.Hosts == nil {
				t.Hosts = map[string]string{}
			}
			t.Hosts[siteID] = fp
			s.save()
			return nil
		case known != fp:
			return &hostKeyChanged{siteName, known, fp}
		}
		return nil
	}
}

func (s *Server) forgetHostKey(w http.ResponseWriter, r *http.Request, c *orgCtx) {
	id := r.PathValue("id")
	s.mu.Lock()
	defer s.mu.Unlock()
	if st := s.findSite(id); st == nil || st.OrgID != c.Org.ID {
		fail(w, 404, errors.New("no such site"))
		return
	}
	delete(s.st.Settings.Terminal.Hosts, id)
	s.save()
	s.audit(c.Org.ID, c.Actor, "terminal.forget-host-key", id)
	ok(w)
}

// sameOrigin lets a browser open the terminal only from this dashboard: a page on another site
// cannot use a signed-in administrator's cookie to start a shell. Requests without an Origin
// header do not come from a browser page.
func sameOrigin(r *http.Request) bool {
	o := r.Header.Get("Origin")
	if o == "" {
		return true
	}
	u, err := url.Parse(o)
	return err == nil && strings.EqualFold(u.Host, r.Host)
}

type termFrame struct {
	text bool
	data []byte
}

// frames reads and writes WebSocket frames keeping their type: binary frames carry the terminal's
// bytes, text frames carry control messages (resize, exit).
var frames = websocket.Codec{
	Marshal: func(v any) ([]byte, byte, error) {
		f := v.(termFrame)
		if f.text {
			return f.data, websocket.TextFrame, nil
		}
		return f.data, websocket.BinaryFrame, nil
	},
	Unmarshal: func(data []byte, pt byte, v any) error {
		*(v.(*termFrame)) = termFrame{text: pt == websocket.TextFrame, data: data}
		return nil
	},
}

type termControl struct {
	Type string `json:"type"`
	Cols int    `json:"cols,omitempty"`
	Rows int    `json:"rows,omitempty"`
	Code int    `json:"code,omitempty"`

	Message        string `json:"message,omitempty"`        // type "error": what went wrong, in words for the person at the keyboard
	HostKeyChanged bool   `json:"hostKeyChanged,omitempty"` // type "error": the site's host key is not the one saved, so offer to forget it
}

func clamp(v, lo, hi int) int {
	if v < lo {
		return lo
	}
	if v > hi {
		return hi
	}
	return v
}

var termSlots = make(chan struct{}, termMaxSessions)

// siteTerminal is the WebSocket that carries one shell session. The connection is accepted first
// and any problem (no SSH port, a refused key, a changed host key) is sent as an "error" message,
// because a browser cannot read the body of a failed WebSocket handshake.
func (s *Server) siteTerminal(w http.ResponseWriter, r *http.Request, c *orgCtx) {
	if !sameOrigin(r) {
		fail(w, 403, errors.New("The terminal can only be opened from the dashboard itself."))
		return
	}
	siteID := r.PathValue("id")
	cols, _ := strconv.Atoi(r.URL.Query().Get("cols"))
	rows, _ := strconv.Atoi(r.URL.Query().Get("rows"))
	cols, rows = clamp(cols, 20, termResizeMaxCol), clamp(rows, 5, termResizeMaxRow)

	srv := websocket.Server{
		Handshake: func(*websocket.Config, *http.Request) error { return nil }, // the origin was checked above
		Handler: func(ws *websocket.Conn) {
			ws.MaxPayloadBytes = termMaxFrame
			s.runTerminal(ws, c, siteID, cols, rows)
		},
	}
	srv.ServeHTTP(w, r)
}

func (s *Server) runTerminal(ws *websocket.Conn, c *orgCtx, siteID string, cols, rows int) {
	var wmu sync.Mutex
	send := func(f termFrame) error {
		wmu.Lock()
		defer wmu.Unlock()
		return frames.Send(ws, f)
	}
	control := func(v termControl) {
		msg, _ := json.Marshal(v)
		send(termFrame{text: true, data: msg})
	}
	problem := func(msg string) {
		control(termControl{Type: "error", Message: msg})
		time.Sleep(200 * time.Millisecond) // let it reach the browser before the connection closes
	}
	defer ws.Close()

	s.mu.Lock()
	signer, err := s.terminalSignerLocked()
	siteName, viaAgent, agentCanRelay := "", false, false
	if st := s.findSite(siteID); st != nil {
		siteName, viaAgent, agentCanRelay = st.Name, st.SSH != "" && st.SSHRelay, st.SSHRelay
	}
	s.mu.Unlock()
	if err != nil {
		problem(err.Error())
		return
	}
	select {
	case termSlots <- struct{}{}:
		defer func() { <-termSlots }()
	default:
		problem("Too many terminals are open. Close one and try again.")
		return
	}

	cfg := &ssh.ClientConfig{User: "hammerhead", Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback: s.trustHostKey(siteID, siteName), Timeout: termDialTimeout}
	var client *ssh.Client
	if viaAgent {
		// Bamboo reports its SSH server, so the agent calls back and no port has to be open anywhere.
		var release func()
		client, release, err = s.dialViaAgent(siteID, cfg)
		defer func() {
			if release != nil {
				release()
			}
		}()
	} else {
		addr, aerr := s.terminalAddr(c, siteID)
		if aerr != nil {
			switch {
			case agentCanRelay: // a current Bamboo whose SSH server is not running
				problem("That machine's SSH server isn't running. Start Bamboo there with -ssh, and add Hammerhead's public key to authorized_keys.")
			default:
				problem(aerr.Error() + " Updating Bamboo on that machine lets the terminal work without any port.")
			}
			return
		}
		client, err = ssh.Dial("tcp", addr, cfg)
	}
	if err != nil {
		var hk *hostKeyChanged
		switch {
		case errors.As(err, &hk):
			control(termControl{Type: "error", Message: hk.Error(), HostKeyChanged: true})
			time.Sleep(200 * time.Millisecond)
		case errors.Is(err, errAgentSilent):
			problem("The site didn't answer. Is Bamboo running there and connected to Hammerhead?")
		case strings.Contains(err.Error(), "unable to authenticate"):
			problem("The site refused Hammerhead's key. Add Hammerhead's public key to authorized_keys on that machine.")
		default:
			problem(fmt.Sprintf("Couldn't reach the site's SSH server: is Bamboo running with -ssh? (%v)", err))
		}
		return
	}
	defer client.Close()
	sess, err := client.NewSession()
	if err != nil {
		problem(err.Error())
		return
	}
	defer sess.Close()
	if err := sess.RequestPty(termDefaultTerm, rows, cols, ssh.TerminalModes{ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400}); err != nil {
		problem(fmt.Sprintf("The site wouldn't give a terminal: %v", err))
		return
	}
	stdin, err := sess.StdinPipe()
	if err != nil {
		problem(err.Error())
		return
	}
	stdout, err := sess.StdoutPipe()
	if err != nil {
		problem(err.Error())
		return
	}
	if err := sess.Shell(); err != nil {
		problem(fmt.Sprintf("Couldn't start a shell: %v", err))
		return
	}
	s.audit(c.Org.ID, c.Actor, "terminal.open", siteName)
	started := time.Now()
	defer func() {
		s.audit(c.Org.ID, c.Actor, "terminal.close", fmt.Sprintf("%s after %s", siteName, time.Since(started).Round(time.Second)))
	}()
	control(termControl{Type: "ready"})

	done := make(chan struct{})
	var once sync.Once
	stop := func() { once.Do(func() { close(done); ws.Close() }) }
	defer stop()

	waited := make(chan struct{}) // closed once the browser has been told how the shell ended
	go func() {                   // shell output to the browser
		defer func() {
			select { // the output ends just before the shell is reaped: let the exit message go out first
			case <-waited:
			case <-time.After(2 * time.Second):
			}
			stop()
		}()
		buf := make([]byte, 16<<10)
		for {
			n, err := stdout.Read(buf)
			if n > 0 {
				if send(termFrame{data: append([]byte(nil), buf[:n]...)}) != nil {
					return
				}
			}
			if err != nil {
				return
			}
		}
	}()
	go func() { // the shell ended: say how, then close
		code := 0
		if err := sess.Wait(); err != nil {
			var ee *ssh.ExitError
			if errors.As(err, &ee) {
				code = ee.ExitStatus()
			} else if !errors.Is(err, io.EOF) {
				code = -1
			}
		}
		control(termControl{Type: "exit", Code: code})
		close(waited)
		time.Sleep(200 * time.Millisecond) // let the exit message reach the browser first
		stop()
	}()
	idle := time.AfterFunc(termIdle, func() {
		control(termControl{Type: "idle"})
		stop()
	})
	defer idle.Stop()
	for { // browser to shell
		var f termFrame
		if err := frames.Receive(ws, &f); err != nil {
			return
		}
		idle.Reset(termIdle)
		if !f.text {
			if _, err := stdin.Write(f.data); err != nil {
				return
			}
			continue
		}
		var ctl termControl
		if json.Unmarshal(f.data, &ctl) == nil && ctl.Type == "resize" {
			sess.WindowChange(clamp(ctl.Rows, 5, termResizeMaxRow), clamp(ctl.Cols, 20, termResizeMaxCol))
		}
	}
}

// ---------- the call-back: a terminal without any open port ----------

// termCallBack is how long Hammerhead waits for an agent to answer a terminal request. The agent
// notices the request within a moment, because it is long-polling for its configuration.
const termCallBack = 15 * time.Second

var errAgentSilent = errors.New("the agent did not call back")

// pendingTerm is a terminal session waiting for the site's agent to call back.
type pendingTerm struct {
	site    string
	created time.Time
	claimed bool
	offered bool          // an agent has been sent this session in its configuration
	conn    chan net.Conn // the agent's call-back connection, once it arrives
	done    chan struct{} // closed when Hammerhead has finished with that connection
}

// wakeAgentsLocked ends every long poll now, so an agent hears about a change at once.
// The configuration version does not move, so nothing is saved or re-applied.
func (s *Server) wakeAgentsLocked() {
	close(s.changed)
	s.changed = make(chan struct{})
}

// waitingTerminalsLocked lists the sessions an agent should answer, and forgets stale ones.
func (s *Server) waitingTerminalsLocked(siteID string) []string {
	var ids []string
	for id, p := range s.terms {
		if time.Since(p.created) > 2*termCallBack {
			delete(s.terms, id)
			continue
		}
		if p.site == siteID && !p.claimed {
			ids = append(ids, id)
			p.offered = true
		}
	}
	return ids
}

// newTerminalsLocked reports whether the site has a session its agent has not been told about yet,
// which makes that agent's long poll return straight away instead of waiting for a change.
func (s *Server) newTerminalsLocked(siteID string) bool {
	for _, p := range s.terms {
		if p.site == siteID && !p.claimed && !p.offered {
			return true
		}
	}
	return false
}

// dialViaAgent asks the site's agent to call back and signs in to its SSH server over that
// connection. release must be called when the session is over.
func (s *Server) dialViaAgent(siteID string, cfg *ssh.ClientConfig) (*ssh.Client, func(), error) {
	id := randHex(12)
	p := &pendingTerm{site: siteID, created: time.Now(), conn: make(chan net.Conn, 1), done: make(chan struct{})}
	s.mu.Lock()
	s.terms[id] = p
	s.wakeAgentsLocked()
	s.mu.Unlock()
	var once sync.Once
	release := func() {
		once.Do(func() {
			s.mu.Lock()
			delete(s.terms, id)
			s.mu.Unlock()
			close(p.done)
		})
	}
	var conn net.Conn
	select {
	case conn = <-p.conn:
	case <-time.After(termCallBack):
		release()
		return nil, nil, errAgentSilent
	}
	conn.SetDeadline(time.Now().Add(termDialTimeout))
	cc, chans, reqs, err := ssh.NewClientConn(conn, "site", cfg)
	if err != nil {
		conn.Close()
		release()
		return nil, nil, err
	}
	conn.SetDeadline(time.Time{})
	return ssh.NewClient(cc, chans, reqs), release, nil
}

// terminalCallBack is where an agent connects to answer a terminal request. The connection is
// handed to the session that asked for it and carries that session's SSH traffic.
func (s *Server) terminalCallBack(w http.ResponseWriter, r *http.Request) {
	s.mu.Lock()
	st := s.siteByToken(r)
	var p *pendingTerm
	if st != nil {
		if q := s.terms[r.PathValue("id")]; q != nil && q.site == st.ID && !q.claimed {
			q.claimed, p = true, q
		}
	}
	s.mu.Unlock()
	switch {
	case st == nil:
		fail(w, 401, errors.New("bad site token"))
		return
	case p == nil:
		fail(w, 404, errors.New("no such terminal"))
		return
	}
	websocket.Server{
		Handshake: func(*websocket.Config, *http.Request) error { return nil }, // authenticated by the site token
		Handler: func(ws *websocket.Conn) {
			ws.PayloadType = websocket.BinaryFrame
			ws.MaxPayloadBytes = termMaxFrame
			select {
			case p.conn <- ws:
				<-p.done // the websocket closes when this handler returns
			case <-time.After(termCallBack):
			}
		},
	}.ServeHTTP(w, r)
}

// clipText cuts s to at most n bytes.
func clipText(s string, n int) string {
	if len(s) > n {
		return s[:n]
	}
	return s
}