DebTik/internal/cli/ssh.go

package cli

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"

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

	"debtik/internal/auth"
	"debtik/internal/logging"
	"debtik/internal/system"
)

const (
	DefaultSSHHostKeyPath = "/etc/debtik/ssh_host_ed25519_key"
)

type SSHServer struct {
	addr     string
	authMgr  *auth.Manager
	executor *Executor
	listener net.Listener
	mu       sync.Mutex
	stopChan chan struct{}
	running  bool
}

func NewSSHServer(
	addr string,
	authMgr *auth.Manager,
	executor *Executor,
) *SSHServer {
	if addr == "" {
		addr = ":2222"
	}
	return &SSHServer{
		addr:     addr,
		authMgr:  authMgr,
		executor: executor,
		stopChan: make(chan struct{}),
	}
}

func (s *SSHServer) SetAddr(addr string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.addr = addr
}

func (s *SSHServer) IsRunning() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.running
}

func (s *SSHServer) Start() error {
	s.mu.Lock()
	if s.running {
		s.mu.Unlock()
		return nil
	}
	s.stopChan = make(chan struct{})

	signer, err := getOrCreateHostKey()
	if err != nil {
		s.mu.Unlock()
		return fmt.Errorf("failed to load SSH host key: %w", err)
	}

	config := &ssh.ServerConfig{
		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			user := conn.User()
			sess, err := s.authMgr.Authenticate(user, string(password))
			if err != nil {
				logging.Log("auth,warning", "SSH login failed for user '%s' from %s: %v", user, conn.RemoteAddr(), err)
				return nil, errors.New("authentication failed")
			}
			logging.Log("auth,info", "SSH login successful for user '%s' from %s (role: %s)", user, conn.RemoteAddr(), sess.Role)
			return &ssh.Permissions{
				Extensions: map[string]string{
					"username": user,
					"role":     string(sess.Role),
				},
			}, nil
		},
	}
	config.AddHostKey(signer)

	l, err := net.Listen("tcp", s.addr)
	if err != nil {
		s.mu.Unlock()
		return fmt.Errorf("failed to bind SSH server on %s: %w", s.addr, err)
	}

	s.listener = l
	s.running = true
	stopCh := s.stopChan
	s.mu.Unlock()

	log.Printf("[CLI-SSH] DebTik SSH server listening on %s", s.addr)

	for {
		rawConn, err := l.Accept()
		if err != nil {
			select {
			case <-stopCh:
				return nil
			default:
				continue
			}
		}

		go s.handleSSHConn(rawConn, config)
	}
}

func (s *SSHServer) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.running {
		return
	}
	s.running = false
	close(s.stopChan)
	if s.listener != nil {
		_ = s.listener.Close()
	}
	log.Printf("[CLI-SSH] SSH server stopped")
}

func (s *SSHServer) handleSSHConn(rawConn net.Conn, config *ssh.ServerConfig) {
	defer rawConn.Close()

	sshConn, chans, reqs, err := ssh.NewServerConn(rawConn, config)
	if err != nil {
		return
	}
	defer sshConn.Close()

	// Discard out-of-band global requests
	go ssh.DiscardRequests(reqs)

	username := sshConn.Permissions.Extensions["username"]
	pubUser, err := s.authMgr.GetPublicUser(username)
	if err != nil || pubUser == nil {
		return
	}

	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			_ = newChan.Reject(ssh.UnknownChannelType, "only session channels supported")
			continue
		}

		channel, requests, err := newChan.Accept()
		if err != nil {
			continue
		}

		go s.handleSession(channel, requests, pubUser)
	}
}

func (s *SSHServer) handleSession(channel ssh.Channel, requests <-chan *ssh.Request, user *auth.PublicUser) {
	defer channel.Close()

	term := terminal.NewTerminal(channel, "")
	term.SetSize(120, 40)

	for req := range requests {
		switch req.Type {
		case "pty-req":
			termWidth := 120
			termHeight := 40
			var pty struct {
				Term                      string
				Cols, Rows, Width, Height uint32
				Modes                     string
			}
			if ssh.Unmarshal(req.Payload, &pty) == nil {
				cols := int(pty.Cols)
				rows := int(pty.Rows)
				if cols > 20 && rows > 5 {
					termWidth = cols
					termHeight = rows
				}
			}
			term.SetSize(termWidth, termHeight)
			_ = req.Reply(true, nil)

		case "window-change":
			if len(req.Payload) >= 8 {
				cols := int(binary.BigEndian.Uint32(req.Payload[:4]))
				rows := int(binary.BigEndian.Uint32(req.Payload[4:8]))
				if cols > 20 && rows > 5 {
					term.SetSize(cols, rows)
				}
			}
			_ = req.Reply(true, nil)

		case "shell":
			_ = req.Reply(true, nil)
			s.runInteractiveShell(term, user)
			return

		case "exec":
			var command struct{ Command string }
			if ssh.Unmarshal(req.Payload, &command) != nil {
				_ = req.Reply(false, nil)
				return
			}
			_ = req.Reply(true, nil)
			output, _ := s.executor.Execute(command.Command, user)
			_, _ = channel.Write([]byte(output))
			_, _ = channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{commandExitStatus(output)}))
			return

		default:
			_ = req.Reply(false, nil)
		}
	}
}

func (s *SSHServer) runInteractiveShell(term *terminal.Terminal, user *auth.PublicUser) {
	st, _ := system.GetStatus()
	hostname := "DebTik"
	if st != nil && st.Hostname != "" {
		hostname = st.Hostname
	}

	banner := Banner("1.0", user.Username, string(user.Role), hostname)
	_, _ = term.Write([]byte(banner))

	prompt := fmt.Sprintf("[%s%s%s@%s%s%s] > ", ColorGreen+ColorBold, user.Username, ColorReset, ColorCyan+ColorBold, hostname, ColorReset)
	term.SetPrompt(prompt)

	completer := NewCompleter(term, s.executor)
	term.AutoCompleteCallback = completer.Callback

	for {
		line, err := term.ReadLine()
		if err != nil {
			if err == io.EOF {
				break
			}
			return
		}

		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}

		if trimmed == "password" {
			currPass, err := term.ReadPassword("Current password: ")
			if err != nil {
				return
			}
			newPass, err := term.ReadPassword("New password: ")
			if err != nil {
				return
			}
			if len(newPass) < 4 {
				_, _ = term.Write([]byte(ColorRed + "Password must be at least 4 characters\n" + ColorReset))
				continue
			}
			confirmPass, err := term.ReadPassword("Confirm new password: ")
			if err != nil {
				return
			}
			if newPass != confirmPass {
				_, _ = term.Write([]byte(ColorRed + "Passwords do not match\n" + ColorReset))
				continue
			}
			if err := s.authMgr.ChangePassword(user.Username, currPass, newPass, false); err != nil {
				_, _ = term.Write([]byte(ColorRed + "Error: " + err.Error() + "\n" + ColorReset))
				continue
			}
			_, _ = term.Write([]byte(ColorGreen + "Password changed successfully!\n" + ColorReset))
			continue
		}

		output, shouldExit := s.executor.Execute(trimmed, user)
		if output != "" {
			_, _ = term.Write([]byte(output))
		}
		if shouldExit {
			break
		}
	}
}

func getOrCreateHostKey() (ssh.Signer, error) {
	keyPath := DefaultSSHHostKeyPath
	if data, err := os.ReadFile(keyPath); err == nil {
		signer, err := ssh.ParsePrivateKey(data)
		if err == nil {
			return signer, nil
		}
	}

	// Generate persistent ED25519 host key
	_ = os.MkdirAll(filepath.Dir(keyPath), 0755)
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	privBytes, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		return nil, err
	}

	pemBlock := pem.EncodeToMemory(privBytes)
	_ = os.WriteFile(keyPath, pemBlock, 0600)

	return ssh.ParsePrivateKey(pemBlock)
}