dzxj_dtu/tools/rkssh/main.go

package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"sync"
	"syscall"

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

type ptyRequest struct {
	Term  string
	Cols  uint32
	Rows  uint32
	Wpx   uint32
	Hpx   uint32
	Modes string
}

type winChange struct {
	Cols uint32
	Rows uint32
	Wpx  uint32
	Hpx  uint32
}

type execRequest struct{ Command string }

type envRequest struct{ Name, Value string }

type session struct {
	ch     ssh.Channel
	reqs   <-chan *ssh.Request
	shell  string
	env    []string
	hasPty bool
	cols   uint32
	rows   uint32
	mu     sync.Mutex
	master *os.File
}

func main() {
	addr := flag.String("addr", "0.0.0.0:22", "listen address")
	user := flag.String("user", "root", "allowed username")
	pass := flag.String("pass", "123456", "password")
	hostKeyPath := flag.String("hostkey", "/data/dropbear/go_hostkey", "host key file")
	shell := flag.String("shell", "/system/bin/sh", "login shell")
	flag.Parse()

	signer, err := loadOrCreateHostKey(*hostKeyPath)
	if err != nil {
		log.Fatalf("host key: %v", err)
	}

	cfg := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, p []byte) (*ssh.Permissions, error) {
			if c.User() == *user && string(p) == *pass {
				return nil, nil
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
	}
	cfg.AddHostKey(signer)

	ln, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("listen %s: %v", *addr, err)
	}
	log.Printf("rkssh listening on %s (user=%s shell=%s)", *addr, *user, *shell)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("accept: %v", err)
			continue
		}
		go handleConn(conn, cfg, *shell)
	}
}

func handleConn(nConn net.Conn, cfg *ssh.ServerConfig, shell string) {
	sshConn, chans, reqs, err := ssh.NewServerConn(nConn, cfg)
	if err != nil {
		log.Printf("handshake from %s failed: %v", nConn.RemoteAddr(), err)
		nConn.Close()
		return
	}
	log.Printf("login: %s from %s", sshConn.User(), sshConn.RemoteAddr())
	defer sshConn.Close()
	go ssh.DiscardRequests(reqs)

	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		ch, chReqs, err := newChan.Accept()
		if err != nil {
			continue
		}
		s := &session{ch: ch, reqs: chReqs, shell: shell, cols: 80, rows: 24}
		go s.run()
	}
}

func (s *session) run() {
	defer s.ch.Close()
	started := false

	for req := range s.reqs {
		switch req.Type {
		case "pty-req":
			var p ptyRequest
			if err := ssh.Unmarshal(req.Payload, &p); err != nil {
				req.Reply(false, nil)
				continue
			}
			s.mu.Lock()
			s.hasPty = true
			if p.Cols > 0 {
				s.cols = p.Cols
			}
			if p.Rows > 0 {
				s.rows = p.Rows
			}
			s.mu.Unlock()
			if p.Term != "" {
				s.env = append(s.env, "TERM="+p.Term)
			}
			req.Reply(true, nil)

		case "window-change":
			var w winChange
			if err := ssh.Unmarshal(req.Payload, &w); err == nil {
				s.resize(w.Cols, w.Rows)
			}
			req.Reply(false, nil)

		case "env":
			var e envRequest
			if err := ssh.Unmarshal(req.Payload, &e); err == nil {
				s.env = append(s.env, e.Name+"="+e.Value)
				req.Reply(true, nil)
			} else {
				req.Reply(false, nil)
			}

		case "shell":
			if started {
				req.Reply(false, nil)
				continue
			}
			started = true
			req.Reply(true, nil)
			go s.exec()

		case "exec":
			if started {
				req.Reply(false, nil)
				continue
			}
			var e execRequest
			if err := ssh.Unmarshal(req.Payload, &e); err != nil {
				req.Reply(false, nil)
				continue
			}
			started = true
			req.Reply(true, nil)
			go s.exec("-c", e.Command)

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

func (s *session) resize(cols, rows uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if cols > 0 {
		s.cols = cols
	}
	if rows > 0 {
		s.rows = rows
	}
	if s.master != nil {
		_ = unix.IoctlSetWinsize(int(s.master.Fd()), unix.TIOCSWINSZ,
			&unix.Winsize{Row: uint16(s.rows), Col: uint16(s.cols)})
	}
}

func (s *session) exec(shellArgs ...string) {
	s.mu.Lock()
	hasPty, cols, rows := s.hasPty, s.cols, s.rows
	s.mu.Unlock()

	cmd := exec.Command(s.shell, shellArgs...)
	cmd.Env = append(defaultEnv(), s.env...)
	cmd.Dir = "/data"
	status := uint32(0)

	if hasPty {
		master, slave, err := openPty(cols, rows)
		if err != nil {
			fmt.Fprintf(s.ch, "rkssh: pty error: %v\r\n", err)
			s.exit(1)
			return
		}
		s.mu.Lock()
		s.master = master
		s.mu.Unlock()
		cmd.Stdin, cmd.Stdout, cmd.Stderr = slave, slave, slave
		cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true, Setctty: true, Ctty: 0}
		if err := cmd.Start(); err != nil {
			fmt.Fprintf(s.ch, "rkssh: start error: %v\r\n", err)
			master.Close()
			slave.Close()
			s.exit(1)
			return
		}
		slave.Close()
		go func() { io.Copy(master, s.ch); master.Close() }()
		io.Copy(s.ch, master)
		if err := cmd.Wait(); err != nil {
			status = exitCode(err)
		}
		s.exit(status)
		return
	}

	stdin, _ := cmd.StdinPipe()
	stdout, _ := cmd.StdoutPipe()
	cmd.Stderr = cmd.Stdout
	if err := cmd.Start(); err != nil {
		fmt.Fprintf(s.ch, "rkssh: start error: %v\n", err)
		s.exit(1)
		return
	}
	go func() { io.Copy(stdin, s.ch); stdin.Close() }()
	io.Copy(s.ch, stdout)
	if err := cmd.Wait(); err != nil {
		status = exitCode(err)
	}
	s.exit(status)
}

func (s *session) exit(code uint32) {
	s.ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{code}))
	s.ch.Close()
}

func exitCode(err error) uint32 {
	if ee, ok := err.(*exec.ExitError); ok {
		if ws, ok := ee.Sys().(syscall.WaitStatus); ok {
			return uint32(ws.ExitStatus())
		}
	}
	return 1
}

func defaultEnv() []string {
	return []string{
		"PATH=/system/bin:/system/xbin:/vendor/bin:/data/dropbear:/data/local/tmp",
		"HOME=/data",
		"TERM=xterm-256color",
		"LANG=C.UTF-8",
	}
}

func openPty(cols, rows uint32) (*os.File, *os.File, error) {
	if cols == 0 {
		cols = 80
	}
	if rows == 0 {
		rows = 24
	}
	master, err := os.OpenFile("/dev/ptmx", os.O_RDWR|syscall.O_NOCTTY, 0)
	if err != nil {
		return nil, nil, err
	}
	if err := unix.IoctlSetPointerInt(int(master.Fd()), unix.TIOCSPTLCK, 0); err != nil {
		master.Close()
		return nil, nil, err
	}
	n, err := unix.IoctlGetInt(int(master.Fd()), unix.TIOCGPTN)
	if err != nil {
		master.Close()
		return nil, nil, err
	}
	_ = unix.IoctlSetWinsize(int(master.Fd()), unix.TIOCSWINSZ,
		&unix.Winsize{Row: uint16(rows), Col: uint16(cols)})
	slave, err := os.OpenFile(fmt.Sprintf("/dev/pts/%d", n), os.O_RDWR|syscall.O_NOCTTY, 0)
	if err != nil {
		master.Close()
		return nil, nil, err
	}
	return master, slave, nil
}

func loadOrCreateHostKey(path string) (ssh.Signer, error) {
	if data, err := os.ReadFile(path); err == nil {
		if signer, err := ssh.ParsePrivateKey(data); err == nil {
			return signer, nil
		}
	}
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}
	block, err := ssh.MarshalPrivateKey(priv, "rkssh")
	if err != nil {
		return nil, err
	}
	if err := os.WriteFile(path, pem.EncodeToMemory(block), 0o600); err != nil {
		log.Printf("warn: cannot persist host key %s: %v", path, err)
	}
	return ssh.NewSignerFromKey(priv)
}