Golang mig.GenID函数代码示例

// getHeartbeat returns a 200
func getHeartbeat(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	loc := fmt.Sprintf("%s%s", ctx.Server.Host, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getHeartbeat()"}.Debug()
	}()
	err := resource.AddItem(cljs.Item{
		Href: request.URL.String(),
		Data: []cljs.Data{
			{
				Name:  "heartbeat",
				Value: "gatorz say hi",
			},
		}})
	if err != nil {
		panic(err)
	}
	respond(200, resource, respWriter, request)
}

func createCommand(ctx Context, action mig.Action, target mig.KeepAlive) (cmdid uint64, err error) {
	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()
	}()

	cmdid = mig.GenID()

	cmd := mig.Command{
		AgentName:     target.Name,
		AgentQueueLoc: target.QueueLoc,
		Status:        "prepared",
		Action:        action,
		ID:            cmdid,
		StartTime:     time.Now().UTC(),
	}

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

	dest := fmt.Sprintf("%s/%d-%d.json", ctx.Directories.Command.Ready, action.ID, cmdid)
	err = safeWrite(ctx, dest, data)
	if err != nil {
		panic(err)
	}

	desc := fmt.Sprintf("created command for action '%s' on target '%s'", action.Name, target.Name)
	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: action.ID, CommandID: cmdid, Desc: desc}

	return
}

// getCommand takes an actionid and a commandid and returns a command
func getCommand(respWriter http.ResponseWriter, request *http.Request) {
	var err error
	opid := mig.GenID()
	resource := cljs.New(request.URL.String())
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%d", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getCommand()"}.Debug()
	}()
	var actionID, commandID int
	aid := request.URL.Query()["actionid"][0]
	if aid != "" {
		actionID, err = strconv.Atoi(aid)
		if err != nil {
			panic(err)
		}
	}
	cmdid := request.URL.Query()["commandid"][0]
	if cmdid != "" {
		commandID, err = strconv.Atoi(cmdid)
		if err != nil {
			panic(err)
		}
	}

	// retrieve the action
	cmds := []mig.Command{}
	var iter *mgo.Iter
	if commandID > 0 {
		if actionID > 0 {
			iter = ctx.DB.Col.Cmd.Find(bson.M{"id": commandID, "action.id": actionID}).Iter()
		} else {
			iter = ctx.DB.Col.Cmd.Find(bson.M{"id": commandID}).Iter()
		}
	} else {
		// nothing to search for, return empty resource
		respond(200, resource, respWriter, request, opid)
	}
	err = iter.All(&cmds)
	if err != nil {
		panic(err)
	}
	if len(cmds) == 0 {
		resource.SetError(cljs.Error{Code: fmt.Sprintf("%d", opid), Message: "No command found"})
		respond(404, resource, respWriter, request, opid)
	}
	// store the results in the resource
	for _, cmd := range cmds {
		commandItem, err := commandToItem(cmd)
		if err != nil {
			panic(err)
		}
		resource.AddItem(commandItem)
	}
	respond(200, resource, respWriter, request, opid)
}

// issueKillAction issues an `agentdestroy` action targetted to a specific agent
// and updates the status of the agent in the database
func issueKillAction(agent mig.Agent, ctx Context) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("issueKillAction() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: "leaving issueKillAction()"}.Debug()
	}()
	// generate an `agentdestroy` action for this agent
	killAction := mig.Action{
		ID:            mig.GenID(),
		Name:          fmt.Sprintf("Kill agent %s", agent.Name),
		Target:        fmt.Sprintf("queueloc='%s'", agent.QueueLoc),
		ValidFrom:     time.Now().Add(-60 * time.Second).UTC(),
		ExpireAfter:   time.Now().Add(30 * time.Minute).UTC(),
		SyntaxVersion: 2,
	}
	var opparams struct {
		PID     int    `json:"pid"`
		Version string `json:"version"`
	}
	opparams.PID = agent.PID
	opparams.Version = agent.Version
	killOperation := mig.Operation{
		Module:     "agentdestroy",
		Parameters: opparams,
	}
	killAction.Operations = append(killAction.Operations, killOperation)

	// sign the action with the scheduler PGP key
	secring, err := getSecring(ctx)
	if err != nil {
		panic(err)
	}
	pgpsig, err := killAction.Sign(ctx.PGP.PrivKeyID, secring)
	if err != nil {
		panic(err)
	}
	killAction.PGPSignatures = append(killAction.PGPSignatures, pgpsig)
	var jsonAction []byte
	jsonAction, err = json.Marshal(killAction)
	if err != nil {
		panic(err)
	}

	// write the action to the spool for scheduling
	dest := fmt.Sprintf("%s/%.0f.json", ctx.Directories.Action.New, killAction.ID)
	err = safeWrite(ctx, dest, jsonAction)
	if err != nil {
		panic(err)
	}

	// mark the agent as `destroyed` in the database
	err = ctx.DB.MarkAgentDestroyed(agent)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("issued kill action for agent '%s' with PID '%d'", agent.Name, agent.PID)}.Warning()
	return
}

// getAction queries the database and retrieves the detail of an action
func getAction(respWriter http.ResponseWriter, request *http.Request) {
	var err error
	opid := mig.GenID()
	loc := fmt.Sprintf("http://%s:%d%s", ctx.Server.IP, ctx.Server.Port, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			emsg := fmt.Sprintf("%v", e)
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: emsg}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: emsg})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getAction()"}.Debug()
	}()
	actionID, err := strconv.ParseFloat(request.URL.Query()["actionid"][0], 64)
	if err != nil {
		err = fmt.Errorf("Wrong parameters 'actionid': '%v'", err)
		panic(err)
	}

	// retrieve the action
	var a mig.Action
	if actionID > 0 {
		a, err = ctx.DB.ActionByID(actionID)
		if err != nil {
			if a.ID == -1 {
				// not found, return 404
				resource.SetError(cljs.Error{
					Code:    fmt.Sprintf("%.0f", opid),
					Message: fmt.Sprintf("Action ID '%.0f' not found", actionID)})
				respond(404, resource, respWriter, request, opid)
				return
			} else {
				panic(err)
			}
		}
	} else {
		// bad request, return 400
		resource.SetError(cljs.Error{
			Code:    fmt.Sprintf("%.0f", opid),
			Message: fmt.Sprintf("Invalid Action ID '%.0f'", actionID)})
		respond(400, resource, respWriter, request, opid)
		return
	}

	// retrieve investigators
	a.Investigators, err = ctx.DB.InvestigatorByActionID(a.ID)
	if err != nil {
		panic(err)
	}
	// store the results in the resource
	actionItem, err := actionToItem(a, true, ctx)
	if err != nil {
		panic(err)
	}
	resource.AddItem(actionItem)
	respond(200, resource, respWriter, request, opid)
}

// processNewAction is called when a new action is available. It pulls
// the action from the directory, parse it, retrieve a list of targets from
// the backend database, and create individual command for each target.
func processNewAction(actionPath string, ctx Context) (err error) {
	var ea mig.ExtendedAction
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("processNewAction() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: ea.Action.ID, Desc: "leaving processNewAction()"}.Debug()
	}()

	// load the action file
	a, err := mig.ActionFromFile(actionPath)
	if err != nil {
		panic(err)
	}
	ea.Action = a

	ea.StartTime = time.Now()

	// generate an action id
	if ea.Action.ID < 1 {
		ea.Action.ID = mig.GenID()
	}

	desc := fmt.Sprintf("new action received: Name='%s' Target='%s' ValidFrom='%s' ExpireAfter='%s'",
		ea.Action.Name, ea.Action.Target, ea.Action.ValidFrom, ea.Action.ExpireAfter)
	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: ea.Action.ID, Desc: desc}

	// TODO: replace with action.Validate(), to include signature verification
	if time.Now().Before(ea.Action.ValidFrom) {
		// queue new action
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: ea.Action.ID, Desc: fmt.Sprintf("action '%s' not ready for scheduling", ea.Action.Name)}
		return
	}
	if time.Now().After(ea.Action.ExpireAfter) {
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: ea.Action.ID, Desc: fmt.Sprintf("removing expired action '%s'", ea.Action.Name)}
		// delete expired action
		os.Remove(actionPath)
		return
	}

	// expand the action in one command per agent
	ea.CommandIDs, err = prepareCommands(ea.Action, ctx)
	if err != nil {
		panic(err)
	}

	// move action to Fly-ing state
	ea.Status = "inflight"
	ea.Counters.Sent = len(ea.CommandIDs)
	err = flyAction(ctx, ea, actionPath)
	if err != nil {
		panic(err)
	}

	return
}

// startAgentsListener will create an AMQP consumer for this agent if none exist
func startAgentListener(reg mig.KeepAlive, ctx Context) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("startAgentListener() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: "leaving startAgentListener()"}.Debug()
	}()

	// If a listener already exists for this agent, exit
	for _, q := range activeAgentsList {
		if q == reg.QueueLoc {
			desc := fmt.Sprintf("startAgentListener() already has a listener for '%s'", reg.QueueLoc)
			ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: desc}.Debug()
			return
		}
	}

	//create a queue for agent
	queue := fmt.Sprintf("mig.sched.%s", reg.QueueLoc)
	_, err = ctx.MQ.Chan.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		panic(err)
	}

	err = ctx.MQ.Chan.QueueBind(queue, queue, "mig", false, nil)
	if err != nil {
		panic(err)
	}

	agentChan, err := ctx.MQ.Chan.Consume(queue, "", true, false, false, false, nil)
	if err != nil {
		panic(err)
	}

	// start a goroutine for this queue
	go func() {
		for msg := range agentChan {
			ctx.OpID = mig.GenID()
			err = recvAgentResults(msg, ctx)
			if err != nil {
				ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: fmt.Sprintf("%v", err)}.Err()
				// TODO: agent is sending bogus results, do something about it
			}
		}
	}()

	desc := fmt.Sprintf("startAgentactiveAgentsListener: started recvAgentResults goroutine for agent '%s'", reg.Name)
	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: desc}.Debug()

	// add the new active queue to the activeAgentsList
	activeAgentsList = append(activeAgentsList, reg.QueueLoc)

	return
}

func getDashboard(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	loc := fmt.Sprintf("http://%s:%d%s", ctx.Server.IP, ctx.Server.Port, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getDashboard()"}.Debug()
	}()

	// get summary of agents active in the last 5 minutes
	sum, err := ctx.DB.SumAgentsByVersion(time.Now().Add(-5 * time.Minute))
	if err != nil {
		panic(err)
	}
	count, err := ctx.DB.CountNewAgents(time.Now().Add(-24 * time.Hour))
	if err != nil {
		panic(err)
	}
	double, err := ctx.DB.CountDoubleAgents(time.Now().Add(-5 * time.Minute))
	if err != nil {
		panic(err)
	}
	disappeared, err := ctx.DB.CountDisappearedAgents(
		time.Now().Add(-7*24*time.Hour), time.Now().Add(-5*time.Minute))
	if err != nil {
		panic(err)
	}
	sumItem, err := agentsSummaryToItem(sum, count, double, disappeared, ctx)
	resource.AddItem(sumItem)

	// add the last 10 actions
	actions, err := ctx.DB.Last10Actions()
	if err != nil {
		panic(err)
	}
	for _, action := range actions {
		// retrieve investigators
		action.Investigators, err = ctx.DB.InvestigatorByActionID(action.ID)
		if err != nil {
			panic(err)
		}
		// store the results in the resource
		actionItem, err := actionToItem(action, false, ctx)
		if err != nil {
			panic(err)
		}
		resource.AddItem(actionItem)
	}
	respond(200, resource, respWriter, request, opid)
}

// getIP returns a the public IP of the caller as read from X-Forwarded-For
func getIP(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	defer func() {
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getIP()"}.Debug()
	}()
	if request.Header.Get("X-FORWARDED-FOR") != "" {
		respond(200, []byte(request.Header.Get("X-FORWARDED-FOR")), respWriter, request)
	} else {
		// request.RemoteAddr contains IP:Port, so strip the port and return just the IP
		respond(200, []byte(request.RemoteAddr[:strings.LastIndex(request.RemoteAddr, ":")]), respWriter, request)
	}
}

// getCommand takes an actionid and a commandid and returns a command
func getCommand(respWriter http.ResponseWriter, request *http.Request) {
	var err error
	opid := mig.GenID()
	loc := fmt.Sprintf("http://%s:%d%s", ctx.Server.IP, ctx.Server.Port, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			emsg := fmt.Sprintf("%v", e)
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: emsg}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: emsg})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getCommand()"}.Debug()
	}()
	commandID, err := strconv.ParseFloat(request.URL.Query()["commandid"][0], 64)
	if err != nil {
		err = fmt.Errorf("Wrong parameters 'commandid': '%v'", err)
		panic(err)
	}

	// retrieve the command
	var cmd mig.Command
	if commandID > 0 {
		cmd, err = ctx.DB.CommandByID(commandID)
		if err != nil {
			if fmt.Sprintf("%v", err) == "Error while retrieving command: 'sql: no rows in result set'" {
				// not found, return 404
				resource.SetError(cljs.Error{
					Code:    fmt.Sprintf("%.0f", opid),
					Message: fmt.Sprintf("Command ID '%.0f' not found", commandID)})
				respond(404, resource, respWriter, request, opid)
				return
			} else {
				panic(err)
			}
		}
	} else {
		// bad request, return 400
		resource.SetError(cljs.Error{
			Code:    fmt.Sprintf("%.0f", opid),
			Message: fmt.Sprintf("Invalid Command ID '%.0f'", commandID)})
		respond(400, resource, respWriter, request, opid)
		return
	}
	// store the results in the resource
	commandItem, err := commandToItem(cmd)
	if err != nil {
		panic(err)
	}
	resource.AddItem(commandItem)
	respond(200, resource, respWriter, request, opid)
}

func searchAgents(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	resource := cljs.New(request.URL.String())
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%d", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving searchAgents()"}.Debug()
	}()
	respond(501, resource, respWriter, request, opid)
}

// cancelCommand receives an action ID and a command ID and issues a cancellation order
func cancelCommand(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	loc := fmt.Sprintf("http://%s:%d%s", ctx.Server.IP, ctx.Server.Port, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving cancelCommand()"}.Debug()
	}()
	respond(501, resource, respWriter, request, opid)
}

func getAgent(respWriter http.ResponseWriter, request *http.Request) {
	opid := mig.GenID()
	loc := fmt.Sprintf("http://%s:%d%s", ctx.Server.IP, ctx.Server.Port, request.URL.String())
	resource := cljs.New(loc)
	defer func() {
		if e := recover(); e != nil {
			ctx.Channels.Log <- mig.Log{OpID: opid, Desc: fmt.Sprintf("%v", e)}.Err()
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: fmt.Sprintf("%v", e)})
			respond(500, resource, respWriter, request, opid)
		}
		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getAgentsDashboard()"}.Debug()
	}()
	agentID, err := strconv.ParseFloat(request.URL.Query()["agentid"][0], 64)
	if err != nil {
		err = fmt.Errorf("Wrong parameters 'agentid': '%v'", err)
		panic(err)
	}

	// retrieve the command
	var agt mig.Agent
	if agentID > 0 {
		agt, err = ctx.DB.AgentByID(agentID)
		if err != nil {
			if fmt.Sprintf("%v", err) == "Error while retrieving agent: 'sql: no rows in result set'" {
				// not found, return 404
				resource.SetError(cljs.Error{
					Code:    fmt.Sprintf("%.0f", opid),
					Message: fmt.Sprintf("Agent ID '%.0f' not found", agentID)})
				respond(404, resource, respWriter, request, opid)
				return
			} else {
				panic(err)
			}
		}
	} else {
		// bad request, return 400
		resource.SetError(cljs.Error{
			Code:    fmt.Sprintf("%.0f", opid),
			Message: fmt.Sprintf("Invalid Agent ID '%.0f'", agentID)})
		respond(400, resource, respWriter, request, opid)
		return
	}
	// store the results in the resource
	agentItem, err := agentToItem(agt)
	if err != nil {
		panic(err)
	}
	resource.AddItem(agentItem)
	respond(200, resource, respWriter, request, opid)
}

// safeWrite performs a two steps write:
// 1) a temp file is written
// 2) the temp file is moved into the target folder
// this prevents the dir watcher from waking up before the file is fully written
func safeWrite(ctx Context, destination string, data []byte) (err error) {
	if len(data) == 0 {
		return fmt.Errorf("data slice is empty. file not written")
	}
	// write the file temp dir
	tmp := fmt.Sprintf("%s/%.0f", ctx.Directories.Tmp, mig.GenID())
	err = ioutil.WriteFile(tmp, data, 0640)
	if err != nil {
		return fmt.Errorf("safeWrite: %v", err)
	}
	// move to destination
	err = os.Rename(tmp, destination)
	if err != nil {
		return fmt.Errorf("safeWrite: %v", err)
	}
	return
}

// recvAgentResults listens on the AMQP channel for command results from agents
// each iteration processes one command received from one agent. The json.Body
// in the message is extracted and written into ctx.Directories.Command.Done
func recvAgentResults(msg amqp.Delivery, ctx Context) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("recvAgentResults() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: "leaving recvAgentResults()"}.Debug()
	}()

	// write to disk Returned directory
	dest := fmt.Sprintf("%s/%d", ctx.Directories.Command.Returned, mig.GenID())
	err = safeWrite(ctx, dest, msg.Body)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: fmt.Sprintf("Received result from '%s'", msg.RoutingKey)}.Debug()

	return
}