// 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 }