// Package join lets a machine be set up without anything logging into it.
//
// The machine doing the setup — running as root, holding the authority —
// listens for a while and hands out one code per machine it expects. On the
// other machine `sudo llmgw join <code>` dials out to it, checks it is talking
// to the listener the code names, proves itself with the code's secret, and
// from then on runs what it is sent, as ssh would, until the setup closes
// the connection. To the setup it is one more onboard.Remote: preflight,
// Patroni and node onboarding run over it unchanged.
//
// For sites where nothing may log into a machine from outside: the
// connection is outbound from the machine being set up, no password is
// typed anywhere, and the code is good for one connection and a limited time.
package join

import (
	"bufio"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net"
	"strings"
	"sync"
	"time"
)

// DefaultPort is where a setup listens for machines joining it.
const DefaultPort = 7444

// codePrefix marks a join code, and its version.
const codePrefix = "llmgw-join1-"

// Code is what `llmgw join` is given: where to dial, how to know it reached
// the right place, and the secret that admits it.
//
// Short, because somebody copies it from one screen into another: the first
// version was JSON and over 250 characters, which a terminal cut off at its
// edge. Packed it is under 90.
type Code struct {
	Addr string
	// Fingerprint is the first 16 bytes of the SHA-256 of the listener's
	// certificate: 128 bits, more than enough to find nothing else by.
	Fingerprint [16]byte
	Secret      [16]byte
	Expires     int64
	// Role says what the machine is being set up as, for its operator to
	// read before it runs: "database", "node" or "control plane".
	Role string
}

var roles = []string{"", "database", "node", "control plane"}

// String is the code as typed: one word, safe in a shell.
func (c Code) String() string {
	b := make([]byte, 0, 40+len(c.Addr))
	b = append(b, c.Fingerprint[:]...)
	b = append(b, c.Secret[:]...)
	b = binary.BigEndian.AppendUint32(b, uint32(c.Expires))
	role := byte(0)
	for i, r := range roles {
		if r == c.Role {
			role = byte(i)
		}
	}
	b = append(b, role)
	b = append(b, c.Addr...)
	return codePrefix + base64.RawURLEncoding.EncodeToString(b)
}

// Parse reads a code back.
func Parse(s string) (Code, error) {
	raw, ok := strings.CutPrefix(strings.TrimSpace(s), codePrefix)
	if !ok {
		return Code{}, errors.New("not a join code: it starts " + codePrefix)
	}
	b, err := base64.RawURLEncoding.DecodeString(raw)
	if err != nil || len(b) < 16+16+4+1+3 {
		return Code{}, errors.New("the join code is damaged: copy it whole")
	}
	var c Code
	copy(c.Fingerprint[:], b[:16])
	copy(c.Secret[:], b[16:32])
	c.Expires = int64(binary.BigEndian.Uint32(b[32:36]))
	if int(b[36]) < len(roles) {
		c.Role = roles[b[36]]
	}
	c.Addr = string(b[37:])
	if _, _, err := net.SplitHostPort(c.Addr); err != nil {
		return Code{}, errors.New("the join code is damaged: copy it whole")
	}
	if time.Now().Unix() > c.Expires {
		return Code{}, fmt.Errorf("the join code expired at %s; ask for a new one",
			time.Unix(c.Expires, 0).Format(time.RFC3339))
	}
	return c, nil
}

// The wire: one JSON object a line, each way.
type hello struct {
	Secret   []byte `json:"secret"`
	Hostname string `json:"hostname"`
	// Addr is the address the machine dialled out from: how the others
	// reach it, when nothing else says.
	Addr string `json:"addr"`
}

type request struct {
	Script string `json:"script,omitempty"`
	Stdin  []byte `json:"stdin,omitempty"`
	Bye    bool   `json:"bye,omitempty"`
}

type response struct {
	Stdout []byte `json:"stdout,omitempty"`
	Stderr []byte `json:"stderr,omitempty"`
	Exit   int    `json:"exit"`
	Error  string `json:"error,omitempty"`
}

// Listener waits for machines to join.
type Listener struct {
	ln   net.Listener
	addr string
	// now is time.Now, replaced by tests.
	now         func() time.Time
	fingerprint [16]byte
	ttl         time.Duration

	mu      sync.Mutex
	pending map[[16]byte]waiting // by secret
}

// waiting is a code not yet used: where its machine is handed over, and
// until when it may come.
type waiting struct {
	ch      chan *Remote
	expires time.Time
}

// Listen starts listening on bind (":7444") and advertises addr ("10.0.0.5:7444")
// in its codes. Codes are good for ttl.
func Listen(bind, addr string, ttl time.Duration) (*Listener, error) {
	cert, fp, err := sessionCert()
	if err != nil {
		return nil, err
	}
	ln, err := tls.Listen("tcp", bind, &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS13})
	if err != nil {
		return nil, fmt.Errorf("listening for machines to join on %s: %w", bind, err)
	}
	l := &Listener{ln: ln, addr: addr, fingerprint: fp, ttl: ttl, pending: map[[16]byte]waiting{}, now: time.Now}
	go l.accept()
	return l, nil
}

// sessionCert is a certificate for this listener alone, trusted by the
// fingerprint in its codes rather than by any authority: it lives as long as
// the setup does.
func sessionCert() (tls.Certificate, [16]byte, error) {
	var fp [16]byte
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return tls.Certificate{}, fp, err
	}
	tmpl := &x509.Certificate{
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		Subject:      pkix.Name{CommonName: "llmgw setup"},
		NotBefore:    time.Now().Add(-time.Minute),
		NotAfter:     time.Now().Add(48 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return tls.Certificate{}, fp, err
	}
	sum := sha256.Sum256(der)
	copy(fp[:], sum[:16])
	return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, fp, nil
}

// Code makes a code for one machine, and the channel its connection arrives
// on.
func (l *Listener) Code(role string) (Code, <-chan *Remote) {
	var secret [16]byte
	_, _ = rand.Read(secret[:])
	ch := make(chan *Remote, 1)
	expires := time.Now().Add(l.ttl)
	l.mu.Lock()
	l.pending[secret] = waiting{ch: ch, expires: expires}
	l.mu.Unlock()
	return Code{Addr: l.addr, Fingerprint: l.fingerprint, Secret: secret,
		Expires: expires.Unix(), Role: role}, ch
}

// Close stops listening; machines already joined stay connected until their
// Remote is closed.
func (l *Listener) Close() error { return l.ln.Close() }

func (l *Listener) accept() {
	for {
		c, err := l.ln.Accept()
		if err != nil {
			return
		}
		go l.admit(c)
	}
}

func (l *Listener) admit(c net.Conn) {
	_ = c.SetDeadline(time.Now().Add(30 * time.Second))
	r := bufio.NewReader(c)
	var h hello
	line, err := r.ReadBytes('\n')
	if err != nil || json.Unmarshal(line, &h) != nil {
		_ = c.Close()
		return
	}
	l.mu.Lock()
	var ch chan *Remote
	expired := false
	now := l.now()
	for secret, p := range l.pending {
		if subtle.ConstantTimeCompare(secret[:], h.Secret) == 1 {
			// One machine a code: whoever comes second is turned away. And
			// not after it expires, which the joining side checking was no
			// guarantee of: a code is a credential for as long as this
			// listener runs, unless this side looks too.
			delete(l.pending, secret)
			if now.After(p.expires) {
				expired = true
			} else {
				ch = p.ch
			}
		}
	}
	l.mu.Unlock()
	if ch == nil {
		msg := "this code is not one the setup is waiting for, or it was used already"
		if expired {
			msg = "this code expired; ask for a new one"
		}
		_, _ = c.Write([]byte(`{"error":"` + msg + `"}` + "\n"))
		_ = c.Close()
		return
	}
	_ = c.SetDeadline(time.Time{})
	if _, err := c.Write([]byte("{}\n")); err != nil {
		_ = c.Close()
		return
	}
	// The address it connected from, not the one it says it has: the setup
	// lets that address into the store, and a claim is anyone's to make.
	addr, _, _ := net.SplitHostPort(c.RemoteAddr().String())
	ch <- &Remote{conn: c, r: r, Hostname: h.Hostname, Addr: addr, Claimed: h.Addr}
}

// Remote is a joined machine, as the setup runs things on it.
type Remote struct {
	Hostname string
	// Addr is where it connected from; Claimed what it said its address is,
	// for showing beside it when the two differ.
	Addr    string
	Claimed string

	mu   sync.Mutex
	conn net.Conn
	r    *bufio.Reader
}

// Run runs script with sh on the machine, as root.
func (m *Remote) Run(ctx context.Context, script string, stdin []byte) ([]byte, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.conn == nil {
		return nil, fmt.Errorf("%s is no longer connected", m.Hostname)
	}
	if dl, ok := ctx.Deadline(); ok {
		_ = m.conn.SetDeadline(dl)
		defer m.conn.SetDeadline(time.Time{})
	}
	b, _ := json.Marshal(request{Script: script, Stdin: stdin})
	if _, err := m.conn.Write(append(b, '\n')); err != nil {
		return nil, fmt.Errorf("%s: %w", m.Hostname, err)
	}
	line, err := m.r.ReadBytes('\n')
	if err != nil {
		return nil, fmt.Errorf("%s went away: %w", m.Hostname, err)
	}
	var resp response
	if err := json.Unmarshal(line, &resp); err != nil {
		return nil, fmt.Errorf("%s answered something unreadable", m.Hostname)
	}
	if resp.Error != "" {
		return append(resp.Stdout, resp.Stderr...), errors.New(resp.Error)
	}
	if resp.Exit != 0 {
		return append(resp.Stdout, resp.Stderr...), fmt.Errorf("exit status %d", resp.Exit)
	}
	return resp.Stdout, nil
}

// Close tells the machine the setup is done, which ends its `llmgw join`.
func (m *Remote) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.conn == nil {
		return nil
	}
	b, _ := json.Marshal(request{Bye: true})
	_, _ = m.conn.Write(append(b, '\n'))
	err := m.conn.Close()
	m.conn = nil
	return err
}