Skip to content

Commit

Permalink
feat(agent): add script sensors
Browse files Browse the repository at this point in the history
- This commit adds the ability to utilise scripts to generate additional
  sensors.
- Each script can produce as many sensors as needed.
- Each script runs on its own interval specified as a cron syntax.
- Scripts need to be placed in a certain location and output a certain
  format.
  • Loading branch information
joshuar committed Nov 26, 2023
1 parent 7feae07 commit ece4ddd
Show file tree
Hide file tree
Showing 4 changed files with 284 additions and 3 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ require (
github.com/lxzan/gws v1.6.14
github.com/nutsdb/nutsdb v0.13.0
github.com/perimeterx/marshmallow v1.1.5
github.com/robfig/cron/v3 v3.0.1
github.com/rs/zerolog v1.31.0
github.com/spf13/pflag v1.0.5 // indirect
github.com/spf13/viper v1.17.0
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,8 @@ github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:Om
github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b h1:0LFwY6Q3gMACTjAbMZBjXAqTOzOwFaj2Ld6cjeQ7Rig=
github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
Expand Down
84 changes: 81 additions & 3 deletions internal/agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,14 @@ import (
"sync"
"syscall"

"github.com/joshuar/go-hass-agent/internal/device"
"github.com/joshuar/go-hass-agent/internal/hass/api"

"github.com/joshuar/go-hass-agent/internal/agent/config"
"github.com/joshuar/go-hass-agent/internal/agent/ui"
fyneui "github.com/joshuar/go-hass-agent/internal/agent/ui/fyneUI"
"github.com/joshuar/go-hass-agent/internal/device"
"github.com/joshuar/go-hass-agent/internal/hass/api"
"github.com/joshuar/go-hass-agent/internal/scripts"
"github.com/joshuar/go-hass-agent/internal/tracker"
"github.com/robfig/cron/v3"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
)
Expand Down Expand Up @@ -117,6 +118,11 @@ func Run(options AgentOptions) {
agent.startWorkers(ctx)
}()
wg.Add(1)
go func() {
defer wg.Done()
agent.runScripts(ctx)
}()
wg.Add(1)
go func() {
defer wg.Done()
agent.runNotificationsWorker(ctx, options)
Expand Down Expand Up @@ -294,3 +300,75 @@ func (agent *Agent) startWorkers(ctx context.Context) {
close(workerCh)
wg.Wait()
}

func (agent *Agent) runScripts(ctx context.Context) {
scriptPath, err := agent.config.StoragePath("scripts")
if err != nil {
log.Error().Err(err).Msg("Could not retrieve script path from config.")
return
}
allScripts, err := scripts.FindScripts(scriptPath)
if err != nil || len(allScripts) == 0 {
log.Error().Err(err).Msg("Could not find any script files.")
return
}
c := cron.New()
var outCh []<-chan tracker.Sensor
for _, s := range allScripts {
schedule := s.Schedule()
if schedule != "" {
_, err := c.AddJob(schedule, s)
if err != nil {
log.Warn().Err(err).Str("script", s.Path()).
Msg("Unable to schedule script.")
break
}
outCh = append(outCh, s.Output)
log.Debug().Str("schedule", schedule).Str("script", s.Path()).
Msg("Added script sensor.")
}
}
log.Debug().Msg("Starting cron scheduler for script sensors.")
c.Start()
go func() {
for s := range mergeSensorCh(ctx, outCh...) {
if err := agent.sensors.UpdateSensors(ctx, s); err != nil {
log.Error().Err(err).Msg("Could not update script sensor.")
}
}
}()
<-ctx.Done()
log.Debug().Msg("Stopping cron scheduler for script sensors.")
cronCtx := c.Stop()
<-cronCtx.Done()
}

func mergeSensorCh(ctx context.Context, sensorCh ...<-chan tracker.Sensor) <-chan tracker.Sensor {
var wg sync.WaitGroup
out := make(chan tracker.Sensor)

// Start an output goroutine for each input channel in sensorCh. output
// copies values from c to out until c is closed, then calls wg.Done.
output := func(c <-chan tracker.Sensor) {
defer wg.Done()
for n := range c {
select {
case out <- n:
case <-ctx.Done():
return
}
}
}
wg.Add(len(sensorCh))
for _, c := range sensorCh {
go output(c)
}

// Start a goroutine to close out once all the output goroutines are
// done. This must start after the wg.Add call.
go func() {
wg.Wait()
close(out)
}()
return out
}
200 changes: 200 additions & 0 deletions internal/scripts/scripts.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
// Copyright (c) 2023 Joshua Rich <[email protected]>
//
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

package scripts

import (
"encoding/json"
"errors"
"os"
"os/exec"
"path/filepath"

"github.com/iancoleman/strcase"
"github.com/joshuar/go-hass-agent/internal/hass/sensor"
"github.com/joshuar/go-hass-agent/internal/tracker"
"github.com/rs/zerolog/log"
"gopkg.in/yaml.v3"
)

type script struct {
path string
schedule string
Output chan tracker.Sensor
}

func (s *script) execute() (*scriptOutput, error) {
cmd := exec.Command(s.path)
o, err := cmd.Output()
if err != nil {
return nil, err
}
output := &scriptOutput{}
err = output.Unmarshal(o)
if err != nil {
return nil, err
}
return output, nil
}

// Run is the function that is called when a script is run by the scheduler on
// its specified schedule. It is implemented to satisfy the cron package
// interface, so the script can be treated as a cron job. Run will execute the
// script, collect the output and send it through a channel as a sensor object.
func (s *script) Run() {
output, err := s.execute()
if err != nil {
log.Warn().Err(err).Str("script", s.path).
Msg("Could not run script.")
return
}

for _, sensor := range output.Sensors {
s.Output <- sensor
}
}

// Schedule retrieves the cron schedule that the script should be run on.73
func (s *script) Schedule() string {
return s.schedule
}

// Path returns the path to the script on disk
func (s *script) Path() string {
return s.path
}

// NewScript returns a new script object that can scheduled with the joib
// scheduler by the agent.
func NewScript(p string) *script {
s := &script{
path: p,
Output: make(chan tracker.Sensor),
}
o, err := s.execute()
if err != nil {
log.Warn().Err(err).Str("script", p).
Msg("Cannot run script")
return nil
}
s.schedule = o.Schedule
return s
}

// scriptOutput represents the output from a script. The output must be
// formatted as either valid JSON or YAML. This output is used to define a
// sensor in Home Assistant.
type scriptOutput struct {
Schedule string `json:"schedule" yaml:"schedule"`
Sensors []*scriptSensor `json:"sensors" yaml:"sensors"`
}

// Unmarshal will attempt to take the raw output from a script execution and
// format it as either JSON or YAML. If successful, this format can then be used
// as a sensor.
func (o *scriptOutput) Unmarshal(b []byte) error {
var err error
err = json.Unmarshal(b, &o)
if err == nil {
return nil
}
err = yaml.Unmarshal(b, &o)
if err == nil {
return nil
}
return errors.New("could not unmarshal script output")
}

type scriptSensor struct {
SensorName string `json:"sensor_name" yaml:"sensor_name"`
SensorIcon string `json:"sensor_icon" yaml:"sensor_icon"`
SensorDeviceClass string `json:"sensor_device_class,omitempty" yaml:"sensor_device_class,omitempty"`
SensorStateClass string `json:"sensor_state_class,omitempty" yaml:"sensor_state_class,omitempty"`
SensorStateType string `json:"sensor_type,omitempty" yaml:"sensor_type,omitempty"`
SensorState interface{} `json:"sensor_state" yaml:"sensor_state"`
SensorUnits string `json:"sensor_units,omitempty" yaml:"sensor_units,omitempty"`
SensorAttributes interface{} `json:"sensor_attributes,omitempty" yaml:"sensor_attributes,omitempty"`
}

func (s *scriptSensor) Name() string {
return s.SensorName
}

func (s *scriptSensor) ID() string {
return strcase.ToSnake(s.SensorName)
}

func (s *scriptSensor) Icon() string {
return s.SensorIcon
}

func (s *scriptSensor) SensorType() sensor.SensorType {
switch s.SensorStateType {
case "binary":
return sensor.TypeBinary
default:
return sensor.TypeSensor
}
}

func (s *scriptSensor) DeviceClass() sensor.SensorDeviceClass {
for d := sensor.Apparent_power; d <= sensor.Wind_speed; d++ {
if s.SensorDeviceClass == d.String() {
return d
}
}
return 0
}

func (s *scriptSensor) StateClass() sensor.SensorStateClass {
switch s.SensorStateClass {
case "measurement":
return sensor.StateMeasurement
case "total":
return sensor.StateTotal
case "total_increasing":
return sensor.StateTotalIncreasing
default:
return 0
}
}

func (s *scriptSensor) State() interface{} {
return s.SensorState
}

func (s *scriptSensor) Units() string {
return s.SensorUnits
}

func (s *scriptSensor) Category() string {
return ""
}

func (s *scriptSensor) Attributes() interface{} {
return s.SensorAttributes
}

func FindScripts(path string) ([]*script, error) {
var scripts []*script
files, err := filepath.Glob(path + "/*")
if err != nil {
return nil, err
}
for _, s := range files {
if isExecutable(s) {
scripts = append(scripts, NewScript(s))
}
}
return scripts, nil
}

func isExecutable(filename string) bool {
fi, err := os.Stat(filename)
if err != nil {
return false
}
return fi.Mode().Perm()&0111 != 0
}

0 comments on commit ece4ddd

Please sign in to comment.