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()

	addr, err := s.terminalAddr(c, siteID)
	if err != nil {
		problem(err.Error())
		return
	}
	s.mu.Lock()
	signer, err := s.terminalSignerLocked()
	siteName := ""
	if st := s.findSite(siteID); st != nil {
		siteName = st.Name
	}
	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
	}

	client, err := ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "hammerhead", Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback: s.trustHostKey(siteID, siteName), Timeout: termDialTimeout})
	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 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))
		}
	}
}