mirror of
https://github.com/YuzuZensai/TrollSSH.git
synced 2026-09-14 03:29:01 +00:00
✨ feat: rewrite in Go
This commit is contained in:
+492
@@ -0,0 +1,492 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
const clearScreen = "\x1b[2J\x1b[0f"
|
||||
|
||||
type ConnectionTracker struct {
|
||||
mu sync.Mutex
|
||||
counts map[string]int
|
||||
total int
|
||||
}
|
||||
|
||||
func newConnectionTracker() *ConnectionTracker {
|
||||
return &ConnectionTracker{counts: make(map[string]int)}
|
||||
}
|
||||
|
||||
func (t *ConnectionTracker) increment(ip string) int {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.counts[ip]++
|
||||
t.total++
|
||||
return t.counts[ip]
|
||||
}
|
||||
|
||||
func (t *ConnectionTracker) decrement(ip string) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if _, ok := t.counts[ip]; !ok {
|
||||
return
|
||||
}
|
||||
t.counts[ip]--
|
||||
if t.total > 0 {
|
||||
t.total--
|
||||
}
|
||||
if t.counts[ip] <= 0 {
|
||||
delete(t.counts, ip)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ConnectionTracker) totalCount() int {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.total
|
||||
}
|
||||
|
||||
func (t *ConnectionTracker) hasReachedLimits(ip string, maxPerIP, maxTotal int) bool {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.total >= maxTotal || t.counts[ip] >= maxPerIP
|
||||
}
|
||||
|
||||
type frameSet struct {
|
||||
data *FramesContainer
|
||||
renderer *FrameRenderer
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
config Config
|
||||
sshConfig *ssh.ServerConfig
|
||||
sets []frameSet
|
||||
tracker *ConnectionTracker
|
||||
fakeLogin *string
|
||||
goodbye *string
|
||||
listener net.Listener
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
type ServerDeps struct {
|
||||
Config Config
|
||||
HostKeys []ssh.Signer
|
||||
BannerText *string
|
||||
FakeLoginText *string
|
||||
GoodbyeText *string
|
||||
VideoSets []*FramesContainer
|
||||
}
|
||||
|
||||
func clampDimension(value, max int) int {
|
||||
if value < 1 {
|
||||
return 1
|
||||
}
|
||||
if value > max {
|
||||
return max
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func createServer(deps ServerDeps) *Server {
|
||||
config := deps.Config
|
||||
|
||||
sets := make([]frameSet, len(deps.VideoSets))
|
||||
for i, data := range deps.VideoSets {
|
||||
sets[i] = frameSet{
|
||||
data: data,
|
||||
renderer: newFrameRenderer(data.Frames, asciiOptions{
|
||||
brightnessThreshold: config.BrightnessThreshold,
|
||||
charset: config.Charset,
|
||||
invert: config.Invert,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
sshConfig := &ssh.ServerConfig{
|
||||
MaxAuthTries: config.MaxAuthAttempts,
|
||||
PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
|
||||
ip := hostOnly(conn.RemoteAddr().String())
|
||||
if config.LogCredentials {
|
||||
logInfo(fmt.Sprintf(
|
||||
`Auth attempt from %s method=password user="%s" pass="%s"`,
|
||||
ip, sanitizeN(conn.User(), 128), sanitizeN(string(password), 128),
|
||||
))
|
||||
}
|
||||
if conn.User() == "" || len(password) == 0 {
|
||||
return nil, errors.New("password rejected")
|
||||
}
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
sshConfig.KeyExchanges = []string{
|
||||
"mlkem768x25519-sha256",
|
||||
"curve25519-sha256",
|
||||
"curve25519-sha256@libssh.org",
|
||||
"ecdh-sha2-nistp256",
|
||||
"ecdh-sha2-nistp384",
|
||||
"ecdh-sha2-nistp521",
|
||||
"diffie-hellman-group14-sha256",
|
||||
}
|
||||
if deps.BannerText != nil {
|
||||
banner := *deps.BannerText
|
||||
sshConfig.BannerCallback = func(_ ssh.ConnMetadata) string { return banner }
|
||||
}
|
||||
for _, key := range deps.HostKeys {
|
||||
sshConfig.AddHostKey(key)
|
||||
}
|
||||
|
||||
return &Server{
|
||||
config: config,
|
||||
sshConfig: sshConfig,
|
||||
sets: sets,
|
||||
tracker: newConnectionTracker(),
|
||||
fakeLogin: deps.FakeLoginText,
|
||||
goodbye: deps.GoodbyeText,
|
||||
}
|
||||
}
|
||||
|
||||
func hostOnly(addr string) string {
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return addr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
func (s *Server) Listen(host string, port int) error {
|
||||
listener, err := net.Listen("tcp", net.JoinHostPort(host, fmt.Sprint(port)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.listener = listener
|
||||
logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port))
|
||||
for {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
go s.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) Close() {
|
||||
s.closeOnce.Do(func() {
|
||||
if s.listener != nil {
|
||||
s.listener.Close()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) handleConn(conn net.Conn) {
|
||||
ip := hostOnly(conn.RemoteAddr().String())
|
||||
|
||||
if s.tracker.hasReachedLimits(ip, s.config.MaxConnections, s.config.MaxTotalConnections) {
|
||||
conn.Close()
|
||||
logWarn("Connection rejected (limit reached) from", ip)
|
||||
return
|
||||
}
|
||||
|
||||
activeForIP := s.tracker.increment(ip)
|
||||
defer s.tracker.decrement(ip)
|
||||
|
||||
if s.config.HandshakeTimeout > 0 {
|
||||
conn.SetDeadline(time.Now().Add(s.config.HandshakeTimeout))
|
||||
}
|
||||
|
||||
sshConn, chans, reqs, err := ssh.NewServerConn(conn, s.sshConfig)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "i/o timeout") {
|
||||
logWarn("Handshake timeout for", ip)
|
||||
} else {
|
||||
logWarn(fmt.Sprintf("Client error from %s:", ip), sanitize(err.Error()))
|
||||
}
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
conn.SetDeadline(time.Time{})
|
||||
logDebug("Handshake from", ip)
|
||||
defer sshConn.Close()
|
||||
|
||||
setIndex := rand.Intn(len(s.sets))
|
||||
logInfo(fmt.Sprintf(
|
||||
"New connection from %s (ip=%d, total=%d) -> playing %q",
|
||||
ip, activeForIP, s.tracker.totalCount(), s.sets[setIndex].data.Name,
|
||||
))
|
||||
|
||||
go ssh.DiscardRequests(reqs)
|
||||
|
||||
for newChannel := range chans {
|
||||
if newChannel.ChannelType() != "session" {
|
||||
newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
|
||||
continue
|
||||
}
|
||||
channel, requests, err := newChannel.Accept()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
go s.handleSession(sshConn, channel, requests, ip, setIndex)
|
||||
}
|
||||
logInfo("Client closed connection from", ip)
|
||||
}
|
||||
|
||||
type termSize struct {
|
||||
mu sync.Mutex
|
||||
width int
|
||||
height int
|
||||
}
|
||||
|
||||
func (t *termSize) set(w, h, maxDim int) {
|
||||
t.mu.Lock()
|
||||
t.width = clampDimension(w, maxDim)
|
||||
t.height = clampDimension(h, maxDim)
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
func (t *termSize) get() (int, int) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.width, t.height
|
||||
}
|
||||
|
||||
func parseDims(payload []byte) (cols, rows int, ok bool) {
|
||||
if len(payload) < 8 {
|
||||
return 0, 0, false
|
||||
}
|
||||
// pty-req prefixes cols/rows with a TERM string; window-change does not.
|
||||
offset := 0
|
||||
strLen := binary.BigEndian.Uint32(payload)
|
||||
if int(strLen)+12 <= len(payload) {
|
||||
offset = 4 + int(strLen)
|
||||
}
|
||||
if len(payload) < offset+8 {
|
||||
return 0, 0, false
|
||||
}
|
||||
cols = int(binary.BigEndian.Uint32(payload[offset:]))
|
||||
rows = int(binary.BigEndian.Uint32(payload[offset+4:]))
|
||||
return cols, rows, true
|
||||
}
|
||||
|
||||
func (s *Server) handleSession(
|
||||
sshConn *ssh.ServerConn,
|
||||
channel ssh.Channel,
|
||||
requests <-chan *ssh.Request,
|
||||
ip string,
|
||||
initialSetIndex int,
|
||||
) {
|
||||
size := &termSize{}
|
||||
size.set(80, 24, s.config.MaxDimension)
|
||||
|
||||
started := false
|
||||
for req := range requests {
|
||||
switch req.Type {
|
||||
case "pty-req":
|
||||
logDebug("Opening pty for session", ip)
|
||||
if cols, rows, ok := parseDims(req.Payload); ok {
|
||||
size.set(cols, rows, s.config.MaxDimension)
|
||||
}
|
||||
req.Reply(true, nil)
|
||||
case "window-change":
|
||||
if len(req.Payload) >= 8 {
|
||||
cols := int(binary.BigEndian.Uint32(req.Payload))
|
||||
rows := int(binary.BigEndian.Uint32(req.Payload[4:]))
|
||||
size.set(cols, rows, s.config.MaxDimension)
|
||||
}
|
||||
if req.WantReply {
|
||||
req.Reply(true, nil)
|
||||
}
|
||||
case "exec":
|
||||
command := ""
|
||||
if len(req.Payload) >= 4 {
|
||||
n := binary.BigEndian.Uint32(req.Payload)
|
||||
if int(n)+4 <= len(req.Payload) {
|
||||
command = string(req.Payload[4 : 4+n])
|
||||
}
|
||||
}
|
||||
logInfo(fmt.Sprintf("Client %s attempted exec: %q", ip, sanitizeN(command, 512)))
|
||||
req.Reply(true, nil)
|
||||
if !started {
|
||||
started = true
|
||||
go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false)
|
||||
}
|
||||
case "shell":
|
||||
logDebug("Opening shell for session", ip)
|
||||
req.Reply(true, nil)
|
||||
if !started {
|
||||
started = true
|
||||
go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false)
|
||||
}
|
||||
default:
|
||||
if req.WantReply {
|
||||
req.Reply(false, nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) pickNextSetIndex(exclude int) int {
|
||||
if len(s.sets) <= 1 {
|
||||
return exclude
|
||||
}
|
||||
next := exclude
|
||||
for next == exclude {
|
||||
next = rand.Intn(len(s.sets))
|
||||
}
|
||||
return next
|
||||
}
|
||||
|
||||
func (s *Server) playVideo(
|
||||
sshConn *ssh.ServerConn,
|
||||
channel ssh.Channel,
|
||||
size *termSize,
|
||||
ip string,
|
||||
setIndex int,
|
||||
keepAspectRatio bool,
|
||||
) {
|
||||
config := s.config
|
||||
current := s.sets[setIndex]
|
||||
|
||||
w, h := size.get()
|
||||
logDebug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip))
|
||||
|
||||
if s.fakeLogin != nil {
|
||||
channel.Write([]byte(clearScreen))
|
||||
channel.Write([]byte(*s.fakeLogin))
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
var doneOnce sync.Once
|
||||
closeSession := func() {
|
||||
doneOnce.Do(func() { close(done) })
|
||||
}
|
||||
|
||||
switchCh := make(chan int, 8)
|
||||
go func() {
|
||||
buf := make([]byte, 256)
|
||||
var lastSwitch time.Time
|
||||
for {
|
||||
n, err := channel.Read(buf)
|
||||
if err != nil {
|
||||
closeSession()
|
||||
return
|
||||
}
|
||||
if !config.AllowUserControl {
|
||||
continue
|
||||
}
|
||||
str := string(buf[:n])
|
||||
delta := 0
|
||||
if strings.Contains(str, "\x1b[C") || strings.Contains(str, "\x1b[A") {
|
||||
delta = 1
|
||||
} else if strings.Contains(str, "\x1b[D") || strings.Contains(str, "\x1b[B") {
|
||||
delta = -1
|
||||
}
|
||||
if delta == 0 {
|
||||
continue
|
||||
}
|
||||
now := time.Now()
|
||||
if now.Sub(lastSwitch) < config.SwitchDebounce {
|
||||
continue
|
||||
}
|
||||
lastSwitch = now
|
||||
select {
|
||||
case switchCh <- delta:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-time.After(config.LoginDelay):
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
|
||||
frameInterval := func() time.Duration {
|
||||
return time.Duration(float64(time.Second) / current.data.FPS)
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(frameInterval())
|
||||
defer ticker.Stop()
|
||||
|
||||
currentFrame := 0
|
||||
loopCount := 0
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
|
||||
case delta := <-switchCh:
|
||||
if len(s.sets) <= 1 {
|
||||
continue
|
||||
}
|
||||
setIndex = (setIndex + delta + len(s.sets)) % len(s.sets)
|
||||
current = s.sets[setIndex]
|
||||
currentFrame = 0
|
||||
logDebug(fmt.Sprintf("%s switched to %q", ip, current.data.Name))
|
||||
ticker.Reset(frameInterval())
|
||||
|
||||
case <-ticker.C:
|
||||
w, h := size.get()
|
||||
ascii, err := current.renderer.render(currentFrame, w, h, keepAspectRatio)
|
||||
if err != nil {
|
||||
logError("Render error for", ip, sanitize(err.Error()))
|
||||
sshConn.Close()
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := channel.Write([]byte(clearScreen + ascii)); err != nil {
|
||||
closeSession()
|
||||
return
|
||||
}
|
||||
|
||||
currentFrame++
|
||||
if currentFrame < len(current.data.Frames) {
|
||||
continue
|
||||
}
|
||||
|
||||
currentFrame = 0
|
||||
loopCount++
|
||||
if config.MaxLoop > 0 && loopCount >= config.MaxLoop {
|
||||
channel.Write([]byte(clearScreen))
|
||||
if s.goodbye != nil {
|
||||
channel.Write([]byte(*s.goodbye))
|
||||
}
|
||||
time.Sleep(1 * time.Second)
|
||||
logInfo("Playback finished, closing session", ip)
|
||||
channel.Close()
|
||||
sshConn.Close()
|
||||
return
|
||||
}
|
||||
|
||||
if config.PlaybackMode == PlaybackRandom {
|
||||
setIndex = s.pickNextSetIndex(setIndex)
|
||||
current = s.sets[setIndex]
|
||||
logInfo(fmt.Sprintf(
|
||||
"Playthrough done for %s, switching to %q", ip, current.data.Name,
|
||||
))
|
||||
ticker.Reset(frameInterval())
|
||||
} else if config.MaxLoop > 0 {
|
||||
logInfo(fmt.Sprintf(
|
||||
"Playthrough done for %s, looping %q (%d/%d)",
|
||||
ip, current.data.Name, loopCount, config.MaxLoop,
|
||||
))
|
||||
} else {
|
||||
logInfo(fmt.Sprintf(
|
||||
"Playthrough done for %s, looping %q (%d)",
|
||||
ip, current.data.Name, loopCount,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user