// Package sshd is a small SSH server for Bamboo, so a machine behind NAT can be reached with an
// ordinary ssh client through a tunnelled port, without installing or configuring sshd. It is
// opt-in and deliberately narrow: public keys only (listed in an authorized_keys file), session
// channels only (a shell or one command, no port forwarding, no agent forwarding, no SFTP), and
// it listens on the loopback address unless told otherwise. Sessions run as the user Bamboo runs as.
package sshd

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

const (
	// DefaultListen is loopback only: a Hammerhead port that targets this site reaches it
	// through Bamboo, and nothing else on the network can.
	DefaultListen = "127.0.0.1:2222"

	maxConns         = 16
	handshakeTimeout = 20 * time.Second
)

// Supported says whether this platform can run shells (it cannot on Windows).
const Supported = runSupported

// Config says where the server listens and which keys may sign in.
type Config struct {
	Listen         string // host:port, default DefaultListen
	AuthorizedKeys string // path of an authorized_keys file; read on every sign-in, so edits apply at once
	HostKey        string // path of the server's ed25519 key; created on first start
	Logf           func(format string, args ...any)
}

// Server accepts SSH connections.
type Server struct {
	cfg  Config
	conf *ssh.ServerConfig
	host ssh.Signer
	sem  chan struct{}
}

// New loads (or creates) the host key and checks that at least one authorized key exists, so
// the server never starts in a state where nobody, or everybody, could sign in.
func New(cfg Config) (*Server, error) {
	if cfg.Listen == "" {
		cfg.Listen = DefaultListen
	}
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	if cfg.AuthorizedKeys == "" {
		return nil, errors.New("no authorized_keys file configured")
	}
	keys, err := loadAuthorizedKeys(cfg.AuthorizedKeys)
	if err != nil {
		return nil, err
	}
	if len(keys) == 0 {
		return nil, fmt.Errorf("%s has no public keys: add the keys that may sign in, one per line", cfg.AuthorizedKeys)
	}
	signer, err := loadOrCreateHostKey(cfg.HostKey)
	if err != nil {
		return nil, err
	}
	s := &Server{cfg: cfg, host: signer, sem: make(chan struct{}, maxConns)}
	s.conf = &ssh.ServerConfig{
		MaxAuthTries: 3,
		// Only public keys. There is no password or keyboard-interactive callback at all.
		PublicKeyCallback: func(c ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			keys, err := loadAuthorizedKeys(s.cfg.AuthorizedKeys)
			if err != nil {
				s.cfg.Logf("ssh: %v", err)
				return nil, errors.New("no access")
			}
			for _, k := range keys {
				if bytes.Equal(k.Marshal(), key.Marshal()) {
					return &ssh.Permissions{Extensions: map[string]string{"fp": ssh.FingerprintSHA256(key)}}, nil
				}
			}
			return nil, errors.New("key not authorized")
		},
	}
	s.conf.AddHostKey(signer)
	return s, nil
}

// Fingerprint is the server's host key fingerprint, for the log so it can be checked on first use.
func (s *Server) Fingerprint() string { return ssh.FingerprintSHA256(s.host.PublicKey()) }

// Run listens until ctx is cancelled.
func (s *Server) Run(ctx context.Context) error {
	ln, err := net.Listen("tcp", s.cfg.Listen)
	if err != nil {
		return err
	}
	return s.Serve(ctx, ln)
}

// Serve accepts connections on ln until ctx is cancelled.
func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
	go func() { <-ctx.Done(); ln.Close() }()
	var wg sync.WaitGroup
	defer wg.Wait()
	for {
		c, err := ln.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return nil
			}
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				continue
			}
			return err
		}
		select {
		case s.sem <- struct{}{}:
		default:
			c.Close() // too many sessions already
			continue
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			defer func() { <-s.sem }()
			s.handle(ctx, c)
		}()
	}
}

func (s *Server) handle(ctx context.Context, nc net.Conn) {
	defer nc.Close()
	nc.SetDeadline(time.Now().Add(handshakeTimeout))
	conn, chans, reqs, err := ssh.NewServerConn(nc, s.conf)
	if err != nil {
		s.cfg.Logf("ssh: sign-in from %s failed: %v", nc.RemoteAddr(), err)
		return
	}
	nc.SetDeadline(time.Time{})
	defer conn.Close()
	s.cfg.Logf("ssh: %s signed in as %s (%s)", conn.User(), nc.RemoteAddr(), conn.Permissions.Extensions["fp"])
	go func() { <-ctx.Done(); conn.Close() }()
	go ssh.DiscardRequests(reqs)
	var wg sync.WaitGroup
	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "only sessions are allowed")
			continue
		}
		c, creqs, err := ch.Accept()
		if err != nil {
			continue
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			s.session(c, creqs)
		}()
	}
	wg.Wait()
}

type ptyReq struct {
	term       string
	cols, rows uint32
}

// session answers a session channel's requests: pty-req, window-change, env (ignored), shell and exec.
func (s *Server) session(ch ssh.Channel, reqs <-chan *ssh.Request) {
	defer ch.Close()
	var pty *ptyReq
	var win chan [2]uint32
	for req := range reqs {
		switch req.Type {
		case "pty-req":
			var p struct {
				Term                 string
				Cols, Rows, Wpx, Hpx uint32
				Modes                string
			}
			if ssh.Unmarshal(req.Payload, &p) != nil || !ptySupported {
				req.Reply(false, nil)
				continue
			}
			pty = &ptyReq{term: p.Term, cols: p.Cols, rows: p.Rows}
			win = make(chan [2]uint32, 4)
			req.Reply(true, nil)
		case "window-change":
			var w struct{ Cols, Rows, Wpx, Hpx uint32 }
			if ssh.Unmarshal(req.Payload, &w) == nil && win != nil {
				select {
				case win <- [2]uint32{w.Cols, w.Rows}:
				default:
				}
			}
			req.Reply(true, nil)
		case "env":
			req.Reply(true, nil) // accepted and ignored: the client cannot change the environment
		case "shell", "exec":
			var cmd string
			if req.Type == "exec" {
				var e struct{ Command string }
				if ssh.Unmarshal(req.Payload, &e) != nil {
					req.Reply(false, nil)
					return
				}
				cmd = e.Command
			}
			if !runSupported {
				req.Reply(false, nil)
				return
			}
			req.Reply(true, nil)
			code := run(ch, cmd, pty, win)
			ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Code uint32 }{uint32(code)}))
			return
		default:
			req.Reply(false, nil) // subsystem (sftp), agent and X11 forwarding, signals: not offered
		}
	}
}

// loadAuthorizedKeys reads the public keys from an authorized_keys file. Options before a key
// (command=, from=, ...) are not supported, so a line that has them is skipped rather than being
// honoured too loosely.
func loadAuthorizedKeys(path string) ([]ssh.PublicKey, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("authorized keys: %w", err)
	}
	var keys []ssh.PublicKey
	for len(b) > 0 {
		key, _, options, rest, err := ssh.ParseAuthorizedKey(b)
		if err != nil {
			break // comments and blank lines are skipped by the parser; anything else ends the list
		}
		b = rest
		if len(options) > 0 {
			continue
		}
		keys = append(keys, key)
	}
	return keys, nil
}

// loadOrCreateHostKey keeps the server's identity across restarts, so clients are not warned
// that it changed.
func loadOrCreateHostKey(path string) (ssh.Signer, error) {
	if path == "" {
		return nil, errors.New("no host key path configured")
	}
	if b, err := os.ReadFile(path); err == nil {
		return ssh.ParsePrivateKey(b)
	} else if !os.IsNotExist(err) {
		return nil, err
	}
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}
	der, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return nil, err
	}
	pemBytes := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return nil, err
	}
	if err := os.WriteFile(path, pemBytes, 0o600); err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(pemBytes)
}