ts3-monolith-go/internal/ts3/query.go

package ts3

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"time"
)

type Query struct {
	Addr      string
	User      string
	Password  string
	VoicePort int
	Timeout   time.Duration
}

type Client struct {
	ID         int
	DatabaseID int
	Nick       string
	UID        string
	IP         string
}

type queryConn struct {
	c net.Conn
	r *bufio.Reader
	w *bufio.Writer
}

type QueryError struct {
	ID  int
	Msg string
}

func (e *QueryError) Error() string { return fmt.Sprintf("TeamSpeak query error %d: %s", e.ID, e.Msg) }

func (q *Query) connect(ctx context.Context) (*queryConn, error) {
	d := net.Dialer{Timeout: q.Timeout}
	c, err := d.DialContext(ctx, "tcp", q.Addr)
	if err != nil {
		return nil, err
	}
	qc := &queryConn{c: c, r: bufio.NewReaderSize(c, 64*1024), w: bufio.NewWriterSize(c, 16*1024)}
	if _, err = qc.run(ctx, "login client_login_name="+escape(q.User)+" client_login_password="+escape(q.Password)); err != nil {
		c.Close()
		return nil, err
	}
	if _, err = qc.run(ctx, "use port="+strconv.Itoa(q.VoicePort)); err != nil {
		c.Close()
		return nil, err
	}
	return qc, nil
}

func (q *Query) withConn(ctx context.Context, fn func(*queryConn) error) error {
	qc, err := q.connect(ctx)
	if err != nil {
		return err
	}
	defer qc.close()
	return fn(qc)
}

func (qc *queryConn) close() {
	_ = qc.c.SetWriteDeadline(time.Now().Add(time.Second))
	_, _ = io.WriteString(qc.c, "quit\n")
	_ = qc.c.Close()
}

func (qc *queryConn) run(ctx context.Context, command string) ([]string, error) {
	if deadline, ok := ctx.Deadline(); ok {
		_ = qc.c.SetDeadline(deadline)
	}
	if _, err := qc.w.WriteString(command + "\n"); err != nil {
		return nil, err
	}
	if err := qc.w.Flush(); err != nil {
		return nil, err
	}

	var lines []string
	for {
		line, err := qc.r.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "error ") {
			fields := parseFields(strings.TrimPrefix(line, "error "))
			id, _ := strconv.Atoi(fields["id"])
			if id != 0 {
				return lines, &QueryError{ID: id, Msg: fields["msg"]}
			}
			return lines, nil
		}
		// Initial TS3 greeting arrives before the first command reply. It is not data.
		if line == "TS3" || strings.HasPrefix(line, "Welcome to the TeamSpeak 3 ServerQuery interface") {
			continue
		}
		lines = append(lines, line)
	}
}

func (q *Query) ListClients(ctx context.Context) ([]Client, error) {
	var out []Client
	err := q.withConn(ctx, func(qc *queryConn) error {
		var err error
		out, err = listClients(ctx, qc)
		return err
	})
	return out, err
}

func listClients(ctx context.Context, qc *queryConn) ([]Client, error) {
	lines, err := qc.run(ctx, "clientlist -ip -uid")
	if err != nil {
		return nil, err
	}
	var out []Client
	for _, line := range lines {
		for _, row := range strings.Split(line, "|") {
			f := parseFields(row)
			if f["client_type"] != "0" {
				continue
			}
			id, err1 := strconv.Atoi(f["clid"])
			dbid, err2 := strconv.Atoi(f["client_database_id"])
			if err1 != nil || err2 != nil {
				continue
			}
			out = append(out, Client{ID: id, DatabaseID: dbid, Nick: f["client_nickname"], UID: f["client_unique_identifier"], IP: f["connection_client_ip"]})
		}
	}
	return out, nil
}

func (q *Query) ClientsByIP(ctx context.Context, ip net.IP) ([]Client, error) {
	clients, err := q.ListClients(ctx)
	if err != nil {
		return nil, err
	}
	var out []Client
	for _, c := range clients {
		if sameIP(ip, net.ParseIP(c.IP)) {
			out = append(out, c)
		}
	}
	return out, nil
}

func (q *Query) CountOnline(ctx context.Context) (int, error) {
	clients, err := q.ListClients(ctx)
	if err != nil {
		return 0, err
	}
	return len(clients), nil
}

func (q *Query) SendPrivateMessage(ctx context.Context, clid int, message string) error {
	return q.withConn(ctx, func(qc *queryConn) error {
		_, err := qc.run(ctx, "sendtextmessage targetmode=1 target="+strconv.Itoa(clid)+" msg="+escape(message))
		return err
	})
}

func (q *Query) CreateChannelForClient(ctx context.Context, expected Client, adminGroupID int, name, topic, password string) (int, error) {
	var cid int
	err := q.withConn(ctx, func(qc *queryConn) error {
		clients, err := listClients(ctx, qc)
		if err != nil {
			return err
		}
		found := false
		for _, c := range clients {
			sameIdentity := c.DatabaseID == expected.DatabaseID
			if expected.UID != "" {
				sameIdentity = sameIdentity && c.UID == expected.UID
			}
			if sameIdentity && sameIP(net.ParseIP(c.IP), net.ParseIP(expected.IP)) {
				expected = c // refresh clid in case the client reconnected
				found = true
				break
			}
		}
		if !found {
			return errors.New("client disconnected or changed")
		}

		cmd := "channelcreate channel_name=" + escape(name) + " channel_topic=" + escape(topic) + " channel_codec=4 channel_flag_permanent=1"
		if password != "" {
			cmd += " channel_password=" + escape(password)
		}
		lines, err := qc.run(ctx, cmd)
		if err != nil {
			return err
		}
		for _, line := range lines {
			f := parseFields(line)
			if v := f["cid"]; v != "" {
				cid, _ = strconv.Atoi(v)
			}
		}
		if cid == 0 {
			return errors.New("TeamSpeak did not return a channel id")
		}

		rollback := func() { _, _ = qc.run(ctx, "channeldelete cid="+strconv.Itoa(cid)+" force=1") }
		if _, err = qc.run(ctx, "setclientchannelgroup cgid="+strconv.Itoa(adminGroupID)+" cid="+strconv.Itoa(cid)+" cldbid="+strconv.Itoa(expected.DatabaseID)); err != nil {
			rollback()
			return err
		}
		if _, err = qc.run(ctx, "clientmove clid="+strconv.Itoa(expected.ID)+" cid="+strconv.Itoa(cid)); err != nil {
			rollback()
			return err
		}
		return nil
	})
	return cid, err
}

func parseFields(s string) map[string]string {
	out := make(map[string]string)
	for _, p := range strings.Fields(s) {
		k, v, ok := strings.Cut(p, "=")
		if !ok {
			continue
		}
		out[k] = unescape(v)
	}
	return out
}

func escape(s string) string {
	var b strings.Builder
	for _, r := range s {
		switch r {
		case '\\':
			b.WriteString(`\\`)
		case '/':
			b.WriteString(`\/`)
		case ' ':
			b.WriteString(`\s`)
		case '|':
			b.WriteString(`\p`)
		case '\a':
			b.WriteString(`\a`)
		case '\b':
			b.WriteString(`\b`)
		case '\f':
			b.WriteString(`\f`)
		case '\n':
			b.WriteString(`\n`)
		case '\r':
			b.WriteString(`\r`)
		case '\t':
			b.WriteString(`\t`)
		case '\v':
			b.WriteString(`\v`)
		default:
			b.WriteRune(r)
		}
	}
	return b.String()
}

func unescape(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		if s[i] != '\\' || i+1 >= len(s) {
			b.WriteByte(s[i])
			continue
		}
		i++
		switch s[i] {
		case '\\':
			b.WriteByte('\\')
		case '/':
			b.WriteByte('/')
		case 's':
			b.WriteByte(' ')
		case 'p':
			b.WriteByte('|')
		case 'a':
			b.WriteByte('\a')
		case 'b':
			b.WriteByte('\b')
		case 'f':
			b.WriteByte('\f')
		case 'n':
			b.WriteByte('\n')
		case 'r':
			b.WriteByte('\r')
		case 't':
			b.WriteByte('\t')
		case 'v':
			b.WriteByte('\v')
		default:
			b.WriteByte(s[i])
		}
	}
	return b.String()
}

func sameIP(a, b net.IP) bool {
	if a == nil || b == nil {
		return false
	}
	return a.Equal(b)
}