ghostream/stream/telnet/telnet.go

179 lines
4.0 KiB
Go

// Package telnet provides some fancy tools, like an ASCII-art stream.
package telnet
import (
"fmt"
"io"
"log"
"net"
"strings"
"time"
)
var (
// Cfg contains the different options of the telnet package, see below
// TODO Config should not be exported
Cfg *Options
currentMessage map[string]*string
clientCount map[string]int
)
// Options holds telnet package configuration
type Options struct {
Enabled bool
ListenAddress string
Width int
Height int
Delay int
}
// Serve starts the telnet server and listen to clients
func Serve(config *Options) {
Cfg = config
if !config.Enabled {
return
}
currentMessage = make(map[string]*string)
clientCount = make(map[string]int)
listener, err := net.Listen("tcp", config.ListenAddress)
if err != nil {
log.Printf("Error while listening to the address %s: %s", config.ListenAddress, err)
return
}
go func() {
for {
s, err := listener.Accept()
if err != nil {
log.Printf("Error while accepting TCP socket: %s", s)
continue
}
go func(s net.Conn) {
streamID := ""
// Request for stream ID
for {
_, err = s.Write([]byte("[GHOSTREAM]\nEnter stream ID: "))
if err != nil {
log.Println("Error while requesting stream ID to telnet client")
_ = s.Close()
return
}
buff := make([]byte, 255)
n, err := s.Read(buff)
if err != nil {
log.Println("Error while requesting stream ID to telnet client")
_ = s.Close()
return
}
// Avoid bruteforce
time.Sleep(3 * time.Second)
streamID = string(buff[:n])
streamID = strings.Replace(streamID, "\r", "", -1)
streamID = strings.Replace(streamID, "\n", "", -1)
if len(streamID) > 0 {
if strings.ToLower(streamID) == "exit" {
_, _ = s.Write([]byte("Goodbye!\n"))
_ = s.Close()
return
}
if _, ok := currentMessage[streamID]; !ok {
_, err = s.Write([]byte("Unknown stream ID.\n"))
if err != nil {
log.Println("Error while requesting stream ID to telnet client")
_ = s.Close()
return
}
continue
}
break
}
}
clientCount[streamID]++
// Hide terminal cursor
_, _ = s.Write([]byte("\033[?25l"))
for {
n, err := s.Write([]byte(*currentMessage[streamID]))
if err != nil {
log.Printf("Error while sending TCP data: %s", err)
_ = s.Close()
clientCount[streamID]--
break
}
if n == 0 {
_ = s.Close()
clientCount[streamID]--
break
}
time.Sleep(time.Duration(config.Delay) * time.Millisecond)
}
}(s)
}
}()
log.Println("Telnet server initialized")
}
// GetNumberConnectedSessions returns the numbers of clients that are viewing the stream through a telnet shell
func GetNumberConnectedSessions(streamID string) int {
if Cfg == nil || !Cfg.Enabled {
return 0
}
return clientCount[streamID]
}
// StartASCIIArtStream send all packets received by ffmpeg as ASCII Art to telnet clients
func StartASCIIArtStream(streamID string, reader io.ReadCloser) {
if !Cfg.Enabled {
_ = reader.Close()
return
}
currentMessage[streamID] = new(string)
pixelBuff := make([]byte, Cfg.Width*Cfg.Height)
textBuff := strings.Builder{}
for {
n, err := reader.Read(pixelBuff)
if err != nil {
log.Printf("An error occurred while reading input: %s", err)
break
}
if n == 0 {
// Stream is finished
break
}
// Header
textBuff.Reset()
textBuff.Grow((40*Cfg.Width+6)*Cfg.Height + 47)
for i := 0; i < 42; i++ {
textBuff.WriteByte('\n')
}
// Convert image to ASCII
for i, pixel := range pixelBuff {
if i%Cfg.Width == 0 {
// New line
textBuff.WriteString("\033[49m\n")
}
// Print two times the character to make a square
text := fmt.Sprintf("\033[48;2;%d;%d;%dm ", pixel, pixel, pixel)
textBuff.WriteString(text)
textBuff.WriteString(text)
}
textBuff.WriteString("\033[49m")
*(currentMessage[streamID]) = textBuff.String()
}
}