package pvetest

import (
	"bufio"
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"

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

// sshServer answers the commands pve-confsync runs on a node: pveversion, the
// read of priv/shadow.cfg and its locked write (adapter password).
type sshServer struct {
	ln   net.Listener
	key  ssh.PublicKey
	exec func(cmd string, in io.Reader, out, errw io.Writer) int
}

func startSSH(s *Server) *sshServer {
	return newSSH(func(user, pass string) bool {
		s.mu.Lock()
		defer s.mu.Unlock()
		return user == "root" && pass == s.Password
	}, s.exec)
}

// newSSH starts an SSH server accepting the password check ok and running
// the commands with exec.
func newSSH(ok func(user, pass string) bool, exec func(cmd string, in io.Reader, out, errw io.Writer) int) *sshServer {
	_, priv, _ := ed25519.GenerateKey(rand.Reader)
	signer, _ := ssh.NewSignerFromKey(priv)
	cfg := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			if ok(c.User(), string(pass)) {
				return nil, nil
			}
			return nil, fmt.Errorf("denied")
		},
	}
	cfg.AddHostKey(signer)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	srv := &sshServer{ln: ln, key: signer.PublicKey(), exec: exec}
	go func() {
		for {
			nc, err := ln.Accept()
			if err != nil {
				return
			}
			go srv.serve(nc, cfg)
		}
	}()
	return srv
}

func (x *sshServer) close() { x.ln.Close() }

// SSHAddr returns the host and port of the fake SSH server.
func (s *Server) SSHAddr() (string, int) {
	a := s.ssh.ln.Addr().(*net.TCPAddr)
	return a.IP.String(), a.Port
}

// SSHHostKey returns the host key of the fake node (authorized_keys format).
func (s *Server) SSHHostKey() string {
	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(s.ssh.key)))
}

func (x *sshServer) serve(nc net.Conn, cfg *ssh.ServerConfig) {
	_, chans, reqs, err := ssh.NewServerConn(nc, cfg)
	if err != nil {
		nc.Close()
		return
	}
	go ssh.DiscardRequests(reqs)
	for nch := range chans {
		if nch.ChannelType() != "session" {
			_ = nch.Reject(ssh.UnknownChannelType, "session only")
			continue
		}
		ch, reqs, err := nch.Accept()
		if err != nil {
			continue
		}
		go func() {
			defer ch.Close()
			for req := range reqs {
				if req.Type != "exec" {
					_ = req.Reply(false, nil)
					continue
				}
				var p struct{ Cmd string }
				_ = ssh.Unmarshal(req.Payload, &p)
				_ = req.Reply(true, nil)
				status := x.exec(p.Cmd, ch, ch, ch.Stderr())
				b := make([]byte, 4)
				binary.BigEndian.PutUint32(b, uint32(status))
				_, _ = ch.SendRequest("exit-status", false, b)
				return
			}
		}()
	}
}

func (s *Server) exec(cmd string, in io.Reader, out, errw io.Writer) int {
	switch {
	case strings.HasPrefix(cmd, "pveversion"):
		fmt.Fprintf(out, "pve-manager/%s/abcdef (running kernel: 6.14.8-2-pve)\n", s.Version)
		return 0
	case strings.HasPrefix(cmd, "cat /etc/pve/priv/shadow.cfg"):
		s.mu.Lock()
		defer s.mu.Unlock()
		for _, u := range sortedKeys(s.Shadow) {
			fmt.Fprintf(out, "%s:%s:\n", u, s.Shadow[u])
		}
		return 0
	case strings.Contains(cmd, "confsync-shadow-write"):
		data, _ := io.ReadAll(in)
		s.mu.Lock()
		defer s.mu.Unlock()
		s.Calls = append(s.Calls, "SSH write shadow.cfg")
		if s.FailOn == "SSH write shadow.cfg" {
			fmt.Fprintln(errw, "injected failure")
			return 1
		}
		sc := bufio.NewScanner(bytes.NewReader(data))
		for sc.Scan() {
			p := strings.SplitN(sc.Text(), "\t", 3)
			switch {
			case len(p) == 3 && p[0] == "set":
				s.Shadow[p[1]] = p[2]
			case len(p) >= 2 && p[0] == "del":
				delete(s.Shadow, p[1])
			default:
				fmt.Fprintln(errw, "invalid operation")
				return 1
			}
		}
		fmt.Fprintln(out, "ok")
		return 0
	}
	fmt.Fprintf(errw, "fake: unknown command %q\n", cmd)
	return 127
}