<commit_before>package supervisor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/starkandwayne/goutils/log"
	"github.com/starkandwayne/shield/agent"

	"github.com/pborman/uuid"
)

type UpdateOp int

const (
	STOPPED UpdateOp = iota
	FAILED
	OUTPUT
	RESTORE_KEY
)

type WorkerUpdate struct {
	Task      uuid.UUID
	Op        UpdateOp
	StoppedAt time.Time
	Output    string
}

type WorkerRequest struct {
	Operation      string `json:"operation"`
	TargetPlugin   string `json:"target_plugin"`
	TargetEndpoint string `json:"target_endpoint"`
	StorePlugin    string `json:"store_plugin"`
	StoreEndpoint  string `json:"store_endpoint"`
}

func worker(id uint, privateKeyFile string, work chan Task, updates chan WorkerUpdate) {
	config, err := agent.ConfigureSSHClient(privateKeyFile)
	if err != nil {
		log.Errorf("worker %d unable to read user key %s: %s; bailing out.\n",
			id, privateKeyFile, err)
		return
	}

	for t := range work {
		client := agent.NewClient(config)

		remote := t.Agent
		if remote == "" {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!!  no remote agent specified for task %s\n", t.UUID)}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}

		err = client.Dial(remote)
		if err != nil {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!!  shield worker %d unable to connect to %s (%s)\n", id, remote, err)}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}
		defer client.Close()

		// start a command and stream output
		final := make(chan string)
		partial := make(chan string)

		go func(out chan string, up chan WorkerUpdate, t Task, in chan string) {
			var buffer []string
			for {
				s, ok := <-in
				if !ok {
					break
				}

				switch s[0:2] {
				case "O:":
					buffer = append(buffer, s[2:])
				case "E:":
					up <- WorkerUpdate{
						Task:   t.UUID,
						Op:     OUTPUT,
						Output: s[2:],
					}
				}
			}
			out <- strings.Join(buffer, "")
			close(out)
		}(final, updates, t, partial)

		command, err := json.Marshal(WorkerRequest{
			Operation:      t.Op.String(),
			TargetPlugin:   t.TargetPlugin,
			TargetEndpoint: t.TargetEndpoint,
			StorePlugin:    t.StorePlugin,
			StoreEndpoint:  t.StoreEndpoint,
		})
		if err != nil {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!! shield worker %d was unable to json encode the request bound for remote agent %s (%s)", id, remote, err),
			}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}
		// exec the command<commit_msg>More restore-key hookup<commit_after>package supervisor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/starkandwayne/goutils/log"
	"github.com/starkandwayne/shield/agent"

	"github.com/pborman/uuid"
)

type UpdateOp int

const (
	STOPPED UpdateOp = iota
	FAILED
	OUTPUT
	RESTORE_KEY
)

type WorkerUpdate struct {
	Task      uuid.UUID
	Op        UpdateOp
	StoppedAt time.Time
	Output    string
}

type WorkerRequest struct {
	Operation      string `json:"operation"`
	TargetPlugin   string `json:"target_plugin"`
	TargetEndpoint string `json:"target_endpoint"`
	StorePlugin    string `json:"store_plugin"`
	StoreEndpoint  string `json:"store_endpoint"`
	RestoreKey     string `json:"restore_key"`
}

func worker(id uint, privateKeyFile string, work chan Task, updates chan WorkerUpdate) {
	config, err := agent.ConfigureSSHClient(privateKeyFile)
	if err != nil {
		log.Errorf("worker %d unable to read user key %s: %s; bailing out.\n",
			id, privateKeyFile, err)
		return
	}

	for t := range work {
		client := agent.NewClient(config)

		remote := t.Agent
		if remote == "" {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!!  no remote agent specified for task %s\n", t.UUID)}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}

		err = client.Dial(remote)
		if err != nil {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!!  shield worker %d unable to connect to %s (%s)\n", id, remote, err)}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}
		defer client.Close()

		// start a command and stream output
		final := make(chan string)
		partial := make(chan string)

		go func(out chan string, up chan WorkerUpdate, t Task, in chan string) {
			var buffer []string
			for {
				s, ok := <-in
				if !ok {
					break
				}

				switch s[0:2] {
				case "O:":
					buffer = append(buffer, s[2:])
				case "E:":
					up <- WorkerUpdate{
						Task:   t.UUID,
						Op:     OUTPUT,
						Output: s[2:],
					}
				}
			}
			out <- strings.Join(buffer, "")
			close(out)
		}(final, updates, t, partial)

		command, err := json.Marshal(WorkerRequest{
			Operation:      t.Op.String(),
			TargetPlugin:   t.TargetPlugin,
			TargetEndpoint: t.TargetEndpoint,
			StorePlugin:    t.StorePlugin,
			StoreEndpoint:  t.StoreEndpoint,
			RestoreKey:     t.RestoreKey,
		})
		if err != nil {
			updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT,
				Output: fmt.Sprintf("TASK FAILED!! shield worker %d was unable to json encode the request bound for remote agent %s (%s)", id, remote, err),
			}
			updates <- WorkerUpdate{Task: t.UUID, Op: FAILED}
			continue
		}
		// exec the command
{'fraction_non_alphanumeric': 0.1260033600896024, 'fraction_numerical': 0.0014933731566175099, 'mean_word_length': 3.159937888198758, 'pattern_counts': {'":': 4, '<': 21, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3702047', 'n_tokens_mistral': 2127, 'n_tokens_neox': 1840, 'n_words': 520}