ha-n95-local-control/internal/ha/send_command.go

package ha

import (
	"bytes"
	"context"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"strings"

	"git.i3omb.com/gronod/ha-n95-local-control/internal/config"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/robot"
)

type sendCommandJSON struct {
	Command   string `json:"command"`
	CleanType string `json:"clean_type"`
	Action    string `json:"action"`
	Name      string `json:"name"`
	On        string `json:"on"`
	Time      string `json:"time"`
	Repeat    string `json:"repeat"`
	XML       string `json:"xml"`
}

func SendCommandHook(cfg config.Config) func(ctx context.Context, serial string, payload []byte) (robot.Command, error) {
	return func(ctx context.Context, serial string, payload []byte) (robot.Command, error) {
		if len(payload) > 4096 {
			return robot.Command{}, errors.New("rejected:size")
		}

		var req sendCommandJSON
		if err := json.Unmarshal(payload, &req); err != nil {
			return robot.Command{}, errors.New("rejected:payload")
		}

		if req.Command == "" {
			return robot.Command{}, errors.New("rejected:payload")
		}

		switch req.Command {
		case "clean":
			if req.CleanType != "auto" && req.CleanType != "border" && req.CleanType != "spot" && req.CleanType != "singleRoom" {
				return robot.Command{}, fmt.Errorf("rejected:clean-type")
			}
			return robot.Command{Name: "clean", Args: map[string]string{"type": req.CleanType}}, nil

		case "move":
			if req.Action == "backward" {
				return robot.Command{}, errors.New("rejected:backward")
			}
			if req.Action != "forward" && req.Action != "SpinLeft" && req.Action != "SpinRight" && req.Action != "TurnAround" && req.Action != "stop" {
				return robot.Command{}, errors.New("rejected:move-action")
			}
			return robot.Command{Name: "move", Args: map[string]string{"action": req.Action}}, nil

		case "cancel_return":
			return robot.Command{Name: "cancel_return"}, nil

		case "set_time":
			return robot.Command{Name: "set_time"}, nil

		case "get_status":
			return robot.Command{Name: "get_status"}, nil

		case "get_lifespan":
			return robot.Command{Name: "get_lifespan"}, nil

		case "add_sched", "mod_sched":
			if req.Name == "" || len(req.Name) > 64 {
				return robot.Command{}, errors.New("rejected:name")
			}
			if req.On != "1" && req.On != "0" && req.On != "true" && req.On != "false" {
				return robot.Command{}, errors.New("rejected:on")
			}
			if req.On == "true" {
				req.On = "1"
			} else if req.On == "false" {
				req.On = "0"
			}
			if req.CleanType != "auto" {
				return robot.Command{}, errors.New("rejected:schedule-type")
			}
			if len(req.Repeat) != 7 {
				return robot.Command{}, errors.New("rejected:repeat")
			}
			for _, c := range req.Repeat {
				if c != '0' && c != '1' {
					return robot.Command{}, errors.New("rejected:repeat")
				}
			}
			if len(req.Time) != 5 || req.Time[2] != ':' {
				return robot.Command{}, errors.New("rejected:time")
			}
			return robot.Command{
				Name: req.Command,
				Args: map[string]string{
					"name":   req.Name,
					"on":     req.On,
					"time":   req.Time,
					"repeat": req.Repeat,
				},
			}, nil

		case "del_sched":
			return robot.Command{Name: "del_sched", Args: map[string]string{"name": req.Name}}, nil

		case "get_sched":
			return robot.Command{Name: "get_sched"}, nil

		case "resume":
			return robot.Command{}, errors.New("rejected:resume")

		case "raw":
			if !cfg.RawCommands {
				return robot.Command{}, errors.New("raw-disabled")
			}
			if len(req.XML) > 4096 {
				return robot.Command{}, errors.New("rejected:size")
			}
			dec := xml.NewDecoder(strings.NewReader(req.XML))
			tok, err := dec.Token()
			if err != nil {
				return robot.Command{}, errors.New("rejected:raw-xml")
			}
			var start xml.StartElement
			var ok bool
			if start, ok = tok.(xml.StartElement); !ok {
				return robot.Command{}, errors.New("rejected:raw-xml")
			}

			// Reject forbidden tags
			switch start.Name.Local {
			case "stream", "auth", "iq", "presence", "message", "starttls", "bind", "session":
				return robot.Command{}, errors.New("rejected:raw-xml-root")
			}

			if start.Name.Local != "ctl" {
				return robot.Command{}, errors.New("rejected:raw-xml-root")
			}

			// Capture inner XML
			var innerBuf bytes.Buffer
			enc := xml.NewEncoder(&innerBuf)
			
			var idAttr string
			var tdAttr string
			
			for _, attr := range start.Attr {
				if attr.Name.Local == "id" {
					idAttr = attr.Value
				} else if attr.Name.Local == "td" {
					tdAttr = attr.Value
				}
			}
			
			if idAttr == "" && tdAttr != "Move" {
				return robot.Command{}, errors.New("rejected:raw-id")
			}

			for {
				t, err := dec.Token()
				if err != nil {
					break
				}
				if end, ok := t.(xml.EndElement); ok && end.Name.Local == "ctl" {
					break
				}
				if err := enc.EncodeToken(t); err != nil {
					return robot.Command{}, errors.New("rejected:raw-xml-inner")
				}
			}
			enc.Flush()

			return robot.Command{
				Name: "raw",
				Args: map[string]string{
					"td": tdAttr,
					"id": idAttr,
				},
				Payload: innerBuf.Bytes(),
			}, nil

		default:
			return robot.Command{}, errors.New("rejected:" + req.Command)
		}
	}
}