Golang mig.Command类代码示例

// receiveResult listens on a temporary channels for results coming from modules. It aggregated them, and
// when all are received, it build a response that is passed to the Result channel
func receiveModuleResults(ctx Context, cmd mig.Command, resultChan chan moduleResult, opsCounter int) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("receiveModuleResults() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "leaving receiveModuleResults()"}.Debug()
	}()
	ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "entering receiveModuleResults()"}.Debug()

	resultReceived := 0

	// for each result received, populate the content of cmd.Results with it
	// stop when we received all the expected results
	for result := range resultChan {
		ctx.Channels.Log <- mig.Log{OpID: result.id, CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "received results from module"}.Debug()
		cmd.Status = result.status
		cmd.Results = append(cmd.Results, result.output)
		resultReceived++
		if resultReceived >= opsCounter {
			break
		}
	}

	// forward the updated command
	ctx.Channels.Results <- cmd

	// close the channel, we're done here
	close(resultChan)
	return
}

// terminateCommand is called when a command result is dropped into ctx.Directories.Command.Returned
// it stores the result of a command and mark it as completed/failed and then
// send a message to the Action completion routine to update the action status
func terminateCommand(cmdPath string, ctx Context) (err error) {
	var cmd mig.Command
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("terminateCommand() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: cmd.Action.ID, CommandID: cmd.ID, Desc: "leaving terminateCommand()"}.Debug()
	}()

	// load and parse the command. If this fail, skip it and continue.
	cmd, err = mig.CmdFromFile(cmdPath)
	if err != nil {
		panic(err)
	}

	cmd.FinishTime = time.Now().UTC()

	// update command in database
	err = ctx.DB.Col.Cmd.Update(bson.M{"id": cmd.ID}, cmd)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: cmd.Action.ID, CommandID: cmd.ID, Desc: "command updated in database"}.Debug()

	// write to disk to Command Done directory
	data, err := json.Marshal(cmd)
	if err != nil {
		panic(err)
	}
	dest := fmt.Sprintf("%s/%d-%d.json", ctx.Directories.Command.Done, cmd.Action.ID, cmd.ID)
	err = safeWrite(ctx, dest, data)
	if err != nil {
		panic(err)
	}

	// remove command from inflight dir
	inflightPath := fmt.Sprintf("%s/%d-%d.json", ctx.Directories.Command.InFlight, cmd.Action.ID, cmd.ID)
	os.Remove(inflightPath)

	// remove command from Returned dir
	os.Remove(cmdPath)

	return
}

func createCommand(ctx Context, action mig.Action, agent mig.Agent, emptyResults []modules.Result) (err error) {
	cmdid := mig.GenID()
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("createCommand() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: action.ID, CommandID: cmdid, Desc: "leaving createCommand()"}.Debug()
	}()
	var cmd mig.Command
	cmd.Status = "sent"
	cmd.Action = action
	cmd.Agent = agent
	cmd.ID = cmdid
	cmd.StartTime = time.Now().UTC()
	cmd.Results = emptyResults
	ctx.Channels.CommandReady <- cmd
	return
}

// sendResults builds a message body and send the command results back to the scheduler
func sendResults(ctx Context, result mig.Command) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("sendResults() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: "leaving sendResults()"}.Debug()
	}()

	ctx.Channels.Log <- mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: "sending command results"}
	result.AgentQueueLoc = ctx.Agent.QueueLoc
	body, err := json.Marshal(result)
	if err != nil {
		panic(err)
	}

	routingKey := fmt.Sprintf("mig.sched.%s", ctx.Agent.QueueLoc)
	err = publish(ctx, "mig", routingKey, body)
	if err != nil {
		panic(err)
	}

	return
}

// parseCommands transforms a message into a MIG Command struct, performs validation
// and run the command
func parseCommands(ctx Context, msg []byte) (err error) {
	var cmd mig.Command
	cmd.ID = 0 // safety net
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("parseCommands() -> %v", e)

			// if we have a command to return, update status and send back
			if cmd.ID > 0 {
				errLog := mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
				cmd.Results = append(cmd.Results, errLog)
				cmd.Status = "failed"
				ctx.Channels.Results <- cmd
			}
		}
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "leaving parseCommands()"}.Debug()
	}()

	// unmarshal the received command into a command struct
	// if this fails, inform the scheduler and skip this message
	err = json.Unmarshal(msg, &cmd)
	if err != nil {
		panic(err)
	}

	// get an io.Reader from the public pgp key
	keyring, keycount, err := pgp.ArmoredPubKeysToKeyring(PUBLICPGPKEYS[0:])
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("loaded %d keys", keycount)}.Debug()

	// Check the action syntax and signature
	err = cmd.Action.Validate()
	if err != nil {
		desc := fmt.Sprintf("action validation failed: %v", err)
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: desc}.Err()
		panic(desc)
	}
	err = cmd.Action.VerifySignature(keyring)
	if err != nil {
		desc := fmt.Sprintf("action signature verification failed: %v", err)
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: desc}.Err()
		panic(desc)
	}

	// Expiration is verified by the Validate() call above, but we need
	// to verify the ScheduledDate ourselves
	if time.Now().Before(cmd.Action.ValidFrom) {
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "action is scheduled for later"}.Err()
		panic("ScheduledDateInFuture")
	}

	// Each operation is ran separately by a module, a channel is created to receive the results from each module
	// a goroutine is created to read from the result channel, and when all modules are done, build the response
	resultChan := make(chan moduleResult)
	opsCounter := 0
	for counter, operation := range cmd.Action.Operations {
		// create an module operation object
		currentOp := moduleOp{id: mig.GenID(),
			mode:       operation.Module,
			params:     operation.Parameters,
			resultChan: resultChan}

		desc := fmt.Sprintf("sending operation %d to module %s", counter, operation.Module)
		ctx.Channels.Log <- mig.Log{OpID: currentOp.id, ActionID: cmd.Action.ID, CommandID: cmd.ID, Desc: desc}

		// pass the module operation object to the proper channel
		switch operation.Module {
		case "filechecker":
			// send the operation to the module
			ctx.Channels.RunAgentCommand <- currentOp
			opsCounter++
		case "shell":
			// send to the external execution path
			ctx.Channels.RunExternalCommand <- currentOp
			opsCounter++
		case "terminate":
			ctx.Channels.Terminate <- fmt.Errorf("Terminate order received from scheduler")
			opsCounter++
		default:
			ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("module '%s' is invalid", operation.Module)}
		}
	}

	// start the goroutine that will receive the results
	go receiveModuleResults(ctx, cmd, resultChan, opsCounter)

	return
}

// sendCommand is called when a command file is created in ctx.Directories.Command.Ready
// it read the command, sends it to the agent via AMQP, and update the DB
func sendCommand(cmdPath string, ctx Context) (err error) {
	var cmd mig.Command
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("sendCommand() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: cmd.Action.ID, Desc: "leaving sendCommand()"}.Debug()
	}()

	// load and parse the command. If this fail, skip it and continue.
	cmd, err = mig.CmdFromFile(cmdPath)
	if err != nil {
		panic(err)
	}

	cmd.Status = "sent"

	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: cmd.Action.ID, CommandID: cmd.ID, Desc: "command received for sending"}.Debug()

	data, err := json.Marshal(cmd)
	if err != nil {
		panic(err)
	}

	// build amqp message for sending
	msg := amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		ContentType:  "text/plain",
		Body:         []byte(data),
	}

	agtQueue := fmt.Sprintf("mig.agt.%s", cmd.AgentQueueLoc)

	// send
	err = ctx.MQ.Chan.Publish("mig", agtQueue, true, false, msg)
	if err != nil {
		panic(err)
	}

	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: cmd.Action.ID, CommandID: cmd.ID, Desc: "command sent to agent queue"}.Debug()

	// write command to InFlight directory
	dest := fmt.Sprintf("%s/%d-%d.json", ctx.Directories.Command.InFlight, cmd.Action.ID, cmd.ID)
	err = safeWrite(ctx, dest, data)
	if err != nil {
		panic(err)
	}

	// remove original command file
	err = os.Remove(cmdPath)
	if err != nil {
		panic(err)
	}

	// store command in database
	err = ctx.DB.Col.Cmd.Insert(cmd)
	if err != nil {
		panic(err)
	}

	return
}