зеркало из https://github.com/mozilla/mig.git
Merge pull request #176 from ameihm0912/action-compress
[medium] introduce optional action compression between client and agent
This commit is contained in:
Коммит
3b34f1f853
74
action.go
74
action.go
|
@ -7,6 +7,9 @@
|
|||
package mig /* import "mig.ninja/mig" */
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
@ -74,12 +77,77 @@ type Threat struct {
|
|||
Type string `json:"type,omitempty"`
|
||||
}
|
||||
|
||||
// an operation is an object that map to an agent module.
|
||||
// the parameters of the operation are passed to the module as argument,
|
||||
// and thus their format depend on the module itself.
|
||||
// an operation is an object that maps to an agent module.
|
||||
// the parameters of the operation are passed to the module as an argument,
|
||||
// and thus their format depends on the module itself.
|
||||
type Operation struct {
|
||||
Module string `json:"module"`
|
||||
Parameters interface{} `json:"parameters"`
|
||||
|
||||
// If WantCompressed is set in the operation, the parameters
|
||||
// will be compressed in PostAction() when the client sends the
|
||||
// action to the API. This will also result in IsCompressed being
|
||||
// marked as true, so the receiving agent knows it must decompress
|
||||
// the parameter data.
|
||||
IsCompressed bool `json:"is_compressed,omitempty"`
|
||||
WantCompressed bool `json:"want_compressed,omitempty"`
|
||||
}
|
||||
|
||||
// Compress the parameters stored within an operation
|
||||
func (op *Operation) CompressOperationParam() (err error) {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
err = fmt.Errorf("CompressOperationParam() -> %v", e)
|
||||
}
|
||||
}()
|
||||
jb, err := json.Marshal(op.Parameters)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
var b bytes.Buffer
|
||||
wb64 := base64.NewEncoder(base64.StdEncoding, &b)
|
||||
w := gzip.NewWriter(wb64)
|
||||
_, err = w.Write(jb)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
w.Close()
|
||||
wb64.Close()
|
||||
op.Parameters = string(b.Bytes())
|
||||
op.IsCompressed = true
|
||||
return
|
||||
}
|
||||
|
||||
// Decompress the parameters stored within an operation
|
||||
func (op *Operation) DecompressOperationParam() (err error) {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
err = fmt.Errorf("DecompressOperationParam() -> %v", e)
|
||||
}
|
||||
}()
|
||||
if !op.IsCompressed {
|
||||
return nil
|
||||
}
|
||||
pstr, ok := op.Parameters.(string)
|
||||
if !ok {
|
||||
panic("Compressed parameter was not a string")
|
||||
}
|
||||
b := bytes.NewBuffer([]byte(pstr))
|
||||
rb64 := base64.NewDecoder(base64.StdEncoding, b)
|
||||
r, err := gzip.NewReader(rb64)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
rb, err := ioutil.ReadAll(r)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
err = json.Unmarshal(rb, &op.Parameters)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
op.IsCompressed = false
|
||||
return
|
||||
}
|
||||
|
||||
// ActionFromFile() reads an action from a local file on the file system
|
||||
|
|
|
@ -672,6 +672,33 @@ func (cli Client) MakeSignedToken() (token string, err error) {
|
|||
return
|
||||
}
|
||||
|
||||
// CompressAction takens a MIG action, and applies compression to any operations
|
||||
// within the action for which compression is requested.
|
||||
//
|
||||
// This function should be called on the action prior to signing it for submission
|
||||
// to the API.
|
||||
func (cli Client) CompressAction(a mig.Action) (comp_action mig.Action, err error) {
|
||||
comp_action = a
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
err = fmt.Errorf("CompressAction() -> %v", e)
|
||||
}
|
||||
}()
|
||||
for i := range comp_action.Operations {
|
||||
if !comp_action.Operations[i].WantCompressed {
|
||||
continue
|
||||
}
|
||||
if comp_action.Operations[i].IsCompressed {
|
||||
continue
|
||||
}
|
||||
err = comp_action.Operations[i].CompressOperationParam()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// SignAction takes a MIG Action, signs it with the key identified in the configuration
|
||||
// and returns the signed action
|
||||
func (cli Client) SignAction(a mig.Action) (signed_action mig.Action, err error) {
|
||||
|
|
|
@ -99,6 +99,10 @@ func main() {
|
|||
a.Target = *target
|
||||
}
|
||||
|
||||
a, err = cli.CompressAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
asig, err := cli.SignAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
|
|
@ -31,8 +31,9 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
|
|||
}
|
||||
}()
|
||||
var (
|
||||
a mig.Action
|
||||
tcount int
|
||||
a mig.Action
|
||||
paramCompression bool
|
||||
tcount int
|
||||
)
|
||||
if tpl.ID == 0 {
|
||||
fmt.Println("Entering action launcher with empty template")
|
||||
|
@ -52,7 +53,7 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
|
|||
prompt := "\x1b[33;1mlauncher>\x1b[0m "
|
||||
for {
|
||||
// completion
|
||||
var symbols = []string{"addoperation", "deloperation", "exit", "help", "init",
|
||||
var symbols = []string{"addoperation", "compress", "deloperation", "exit", "help", "init",
|
||||
"json", "launch", "listagents", "load", "details", "filechecker", "netstat",
|
||||
"setname", "settarget", "settimes", "sign", "times"}
|
||||
readline.Completer = func(query, ctx string) []string {
|
||||
|
@ -98,6 +99,9 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
|
|||
fmt.Printf("Parameters creation failed with error: %v\n", err)
|
||||
break
|
||||
}
|
||||
if paramCompression {
|
||||
operation.WantCompressed = true
|
||||
}
|
||||
a.Operations = append(a.Operations, operation)
|
||||
opjson, err := json.MarshalIndent(operation, "", " ")
|
||||
if err != nil {
|
||||
|
@ -108,6 +112,35 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
|
|||
fmt.Println("Module", operation.Module, "is not available in this console...")
|
||||
fmt.Println("You can write your action by hand and import it using 'load <file>'")
|
||||
}
|
||||
case "compress":
|
||||
if len(orders) != 2 {
|
||||
fmt.Println("Wrong arguments: Expects 'compress <true|false>'")
|
||||
fmt.Println("example: compress true")
|
||||
break
|
||||
}
|
||||
switch strings.ToLower(orders[1]) {
|
||||
case "false":
|
||||
paramCompression = false
|
||||
// Disable compression on all existing operations
|
||||
for i := range a.Operations {
|
||||
a.Operations[i].WantCompressed = false
|
||||
err = a.Operations[i].DecompressOperationParam()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
// Invalidate any signatures applied to the action at this point
|
||||
hasSignatures = false
|
||||
a.PGPSignatures = nil
|
||||
case "true":
|
||||
paramCompression = true
|
||||
// Enable compression on all existing operations
|
||||
for i := range a.Operations {
|
||||
a.Operations[i].WantCompressed = true
|
||||
}
|
||||
default:
|
||||
fmt.Println("Argument to compress must be true or false")
|
||||
}
|
||||
case "deloperation":
|
||||
if len(orders) != 2 {
|
||||
fmt.Println("Wrong arguments. Expects 'deloperation <opnum>'")
|
||||
|
@ -137,6 +170,7 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
|
|||
case "help":
|
||||
fmt.Printf(`The following orders are available:
|
||||
addoperation <module> append a new operation of type <module> to the action operations
|
||||
compress <false|true> request parameter compression in operations stored in action
|
||||
listagents list agents targetted by an action
|
||||
deloperation <opnum> remove operation numbered <opnum> from operations array, count starts at zero
|
||||
details display the action details
|
||||
|
@ -160,11 +194,15 @@ times show the various timestamps of the action
|
|||
fmt.Printf("Unknown option '%s'\n", orders[1])
|
||||
}
|
||||
}
|
||||
tmpAction, err := getActionView(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
var ajson []byte
|
||||
if pack {
|
||||
ajson, err = json.Marshal(a)
|
||||
ajson, err = json.Marshal(tmpAction)
|
||||
} else {
|
||||
ajson, err = json.MarshalIndent(a, "", " ")
|
||||
ajson, err = json.MarshalIndent(tmpAction, "", " ")
|
||||
}
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
@ -220,6 +258,10 @@ times show the various timestamps of the action
|
|||
hasTimes = true
|
||||
}
|
||||
if !hasSignatures {
|
||||
a, err = cli.CompressAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
asig, err := cli.SignAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
@ -267,6 +309,10 @@ times show the various timestamps of the action
|
|||
fmt.Println("Times must be set prior to signing")
|
||||
break
|
||||
}
|
||||
a, err = cli.CompressAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
asig, err := cli.SignAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
@ -414,3 +460,33 @@ finish:
|
|||
a.Counters.Failed, a.Counters.TimeOut, a.Counters.InFlight)
|
||||
return
|
||||
}
|
||||
|
||||
// Return a view of an action suitable for JSON display in the console; this
|
||||
// function essentially strips compression from the parameters, but leaves
|
||||
// other fields intact -- the output should be used for display purposes only.
|
||||
func getActionView(a mig.Action) (ret mig.Action, err error) {
|
||||
ret = a
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
err = fmt.Errorf("getActionView() -> %v", e)
|
||||
}
|
||||
}()
|
||||
// Create a copy of the original operations to modify, so we don't
|
||||
// change any of the original action parameters
|
||||
ret.Operations = make([]mig.Operation, len(a.Operations))
|
||||
copy(ret.Operations, a.Operations)
|
||||
for i := range ret.Operations {
|
||||
if !ret.Operations[i].IsCompressed {
|
||||
continue
|
||||
}
|
||||
err = ret.Operations[i].DecompressOperationParam()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
// Reset the IsCompressed flag, purely for visual purposes to
|
||||
// indicate the parameters are compressed when the JSON is
|
||||
// viewed
|
||||
ret.Operations[i].IsCompressed = true
|
||||
}
|
||||
return
|
||||
}
|
||||
|
|
|
@ -132,8 +132,12 @@ times show the various timestamps of the action
|
|||
fmt.Println(i.Name, "- Key ID:", i.PGPFingerprint)
|
||||
}
|
||||
case "json":
|
||||
tmpAction, err := getActionView(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
var ajson []byte
|
||||
ajson, err = json.MarshalIndent(a, "", " ")
|
||||
ajson, err = json.MarshalIndent(tmpAction, "", " ")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
|
|
@ -47,6 +47,7 @@ usage: %s <module> <global options> <module parameters>
|
|||
* run on local system: -t local
|
||||
-v verbose output, includes debug information and raw queries
|
||||
-V print version
|
||||
-z <bool> compress action before sending it to agents
|
||||
|
||||
Progress information is sent to stderr, silence it with "2>/dev/null".
|
||||
Results are sent to stdout, redirect them with "1>/path/to/file".
|
||||
|
@ -76,6 +77,7 @@ func main() {
|
|||
migrc, show, render, target, expiration, afile string
|
||||
printAndExit bool
|
||||
verbose, showversion bool
|
||||
compressAction bool
|
||||
modargs []string
|
||||
run interface{}
|
||||
)
|
||||
|
@ -96,6 +98,7 @@ func main() {
|
|||
fs.StringVar(&afile, "i", "/path/to/file", "Load action from file")
|
||||
fs.BoolVar(&verbose, "v", false, "Enable verbose output")
|
||||
fs.BoolVar(&showversion, "V", false, "Show version")
|
||||
fs.BoolVar(&compressAction, "z", false, "Request compression of action parameters")
|
||||
|
||||
// if first argument is missing, or is help, print help
|
||||
// otherwise, pass the remainder of the arguments to the module for parsing
|
||||
|
@ -181,10 +184,14 @@ func main() {
|
|||
if err != nil || op.Parameters == nil {
|
||||
panic(err)
|
||||
}
|
||||
// If compression has been enabled, flag it in the operation.
|
||||
if compressAction {
|
||||
op.WantCompressed = true
|
||||
}
|
||||
// If running against the local target, don't post the action to the MIG API
|
||||
// but run it locally instead.
|
||||
if target == "local" {
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, op.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, op.Parameters, false)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
@ -252,6 +259,10 @@ readytolaunch:
|
|||
// add extra 60 seconds taken for clock skew
|
||||
a.ExpireAfter = a.ExpireAfter.Add(60 * time.Second).UTC()
|
||||
|
||||
a, err = cli.CompressAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
asig, err := cli.SignAction(a)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
|
|
@ -38,13 +38,14 @@ type moduleResult struct {
|
|||
}
|
||||
|
||||
type moduleOp struct {
|
||||
err error
|
||||
id float64
|
||||
mode string
|
||||
params interface{}
|
||||
resultChan chan moduleResult
|
||||
position int
|
||||
expireafter time.Time
|
||||
err error
|
||||
id float64
|
||||
mode string
|
||||
isCompressed bool
|
||||
params interface{}
|
||||
resultChan chan moduleResult
|
||||
position int
|
||||
expireafter time.Time
|
||||
}
|
||||
|
||||
var runningOps = make(map[float64]moduleOp)
|
||||
|
@ -212,7 +213,7 @@ func runModuleDirectly(mode string, paramargs interface{}, pretty bool) (out str
|
|||
// If parameters are being supplied as an argument, use these vs.
|
||||
// expecting parameters to be supplied on Stdin.
|
||||
if paramargs != nil {
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, paramargs)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, paramargs, false)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
@ -486,12 +487,13 @@ func parseCommands(ctx Context, msg []byte) (err error) {
|
|||
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,
|
||||
position: counter,
|
||||
expireafter: cmd.Action.ExpireAfter,
|
||||
id: mig.GenID(),
|
||||
mode: operation.Module,
|
||||
isCompressed: operation.IsCompressed,
|
||||
params: operation.Parameters,
|
||||
resultChan: resultChan,
|
||||
position: counter,
|
||||
expireafter: cmd.Action.ExpireAfter,
|
||||
}
|
||||
|
||||
desc := fmt.Sprintf("sending operation %d to module %s", counter, operation.Module)
|
||||
|
@ -554,7 +556,7 @@ func runModule(ctx Context, op moduleOp) (err error) {
|
|||
}
|
||||
|
||||
// Build parameters message
|
||||
modParams, err := modules.MakeMessage(modules.MsgClassParameters, op.params)
|
||||
modParams, err := modules.MakeMessage(modules.MsgClassParameters, op.params, op.isCompressed)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
|
|
@ -90,6 +90,10 @@ func (e *entity) launchAction() (err error) {
|
|||
// time begins in the past.
|
||||
period += -window
|
||||
act.ExpireAfter = act.ValidFrom.Add(period)
|
||||
act, err = cli.CompressAction(act)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
asig, err := cli.SignAction(act)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
|
|
@ -52,7 +52,7 @@ func TestNameSearch(t *testing.T) {
|
|||
s.Names = append(s.Names, "!^"+tp.name+"FOOBAR$")
|
||||
s.Options.MatchAll = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -85,7 +85,7 @@ func TestContentSearch(t *testing.T) {
|
|||
s.Contents = append(s.Contents, "!^FOOBAR$")
|
||||
s.Options.MatchAll = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -127,7 +127,7 @@ func TestDecompressedContentSearch(t *testing.T) {
|
|||
s.Options.MatchAll = true
|
||||
s.Options.Decompress = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -158,7 +158,7 @@ func TestSize(t *testing.T) {
|
|||
s.Paths = append(s.Paths, basedir)
|
||||
s.Sizes = append(s.Sizes, tp.size)
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -191,7 +191,7 @@ func TestMTime(t *testing.T) {
|
|||
s.Mtimes = append(s.Mtimes, tp.mtime)
|
||||
s.Options.MatchAll = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -224,7 +224,7 @@ func TestMode(t *testing.T) {
|
|||
s.Modes = append(s.Modes, tp.mode)
|
||||
s.Options.MatchAll = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -265,7 +265,7 @@ func TestHashes(t *testing.T) {
|
|||
s.SHA3 = append(s.SHA3, tp.sha3)
|
||||
}
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -306,7 +306,7 @@ func TestDecompressedHash(t *testing.T) {
|
|||
}
|
||||
s.Options.Decompress = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -341,7 +341,7 @@ func TestAllHashes(t *testing.T) {
|
|||
s.SHA3 = append(s.SHA3, tp.sha3)
|
||||
s.Options.MatchAll = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -373,7 +373,7 @@ func TestMaxDepth(t *testing.T) {
|
|||
s.Options.MatchAll = true
|
||||
s.Options.MaxDepth = 5
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -443,7 +443,7 @@ func TestMacroal(t *testing.T) {
|
|||
s.Options.MatchAll = true
|
||||
s.Options.Macroal = true
|
||||
r.Parameters.Searches["s1"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -591,7 +591,7 @@ func TestMismatch(t *testing.T) {
|
|||
var r run
|
||||
r.Parameters = *newParameters()
|
||||
r.Parameters.Searches["s1"] = mt.search
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
|
@ -70,7 +70,7 @@ func TestFindGoTestProcess(t *testing.T) {
|
|||
s.Options.MaxLength = 10000000
|
||||
s.Options.LogFailures = true
|
||||
r.Parameters.Searches["testsearch"] = s
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
@ -119,7 +119,7 @@ func TestSearches(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters)
|
||||
msg, err := modules.MakeMessage(modules.MsgClassParameters, r.Parameters, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
|
@ -12,9 +12,13 @@ package modules /* import "mig.ninja/mig/modules" */
|
|||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
)
|
||||
|
||||
// Message defines the input messages received by modules.
|
||||
|
@ -63,13 +67,41 @@ type Runner interface {
|
|||
|
||||
// MakeMessage creates a new modules.Message with a given class and parameters and
|
||||
// return the byte slice of the json marshalled message
|
||||
func MakeMessage(class MessageClass, params interface{}) (rawMsg []byte, err error) {
|
||||
func MakeMessage(class MessageClass, params interface{}, comp bool) (rawMsg []byte, err error) {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
err = fmt.Errorf("Failed to make modules.Message: %v", e)
|
||||
}
|
||||
}()
|
||||
|
||||
var msg Message
|
||||
msg.Class = class
|
||||
msg.Parameters = params
|
||||
// If the compression flag is set, treat Parameters as a compressed
|
||||
// byte string.
|
||||
if comp {
|
||||
pstr, ok := msg.Parameters.(string)
|
||||
if !ok {
|
||||
panic("Compressed parameter was not a string")
|
||||
}
|
||||
b := bytes.NewBuffer([]byte(pstr))
|
||||
rb64 := base64.NewDecoder(base64.StdEncoding, b)
|
||||
r, err := gzip.NewReader(rb64)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
rb, err := ioutil.ReadAll(r)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
err = json.Unmarshal(rb, &msg.Parameters)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
rawMsg, err = json.Marshal(msg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to make modules.Message: %v", err)
|
||||
panic(err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
|
|
@ -7,6 +7,10 @@
|
|||
package modules /* import "mig.ninja/mig/modules" */
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
|
@ -60,7 +64,7 @@ func TestRegister(t *testing.T) {
|
|||
func TestMakeMessage(t *testing.T) {
|
||||
var p params
|
||||
p.SomeParam = "foo"
|
||||
raw, err := MakeMessage(MsgClassParameters, p)
|
||||
raw, err := MakeMessage(MsgClassParameters, p, false)
|
||||
if err != nil {
|
||||
t.Fatalf(err.Error())
|
||||
}
|
||||
|
@ -68,7 +72,29 @@ func TestMakeMessage(t *testing.T) {
|
|||
t.Fatalf("Invalid module message class `parameters`")
|
||||
}
|
||||
|
||||
raw, err = MakeMessage(MsgClassStop, nil)
|
||||
// Test parameter decompression
|
||||
jb, err := json.Marshal(p)
|
||||
if err != nil {
|
||||
t.Fatalf(err.Error())
|
||||
}
|
||||
var b bytes.Buffer
|
||||
wb64 := base64.NewEncoder(base64.StdEncoding, &b)
|
||||
w := gzip.NewWriter(wb64)
|
||||
_, err = w.Write(jb)
|
||||
if err != nil {
|
||||
t.Fatalf(err.Error())
|
||||
}
|
||||
w.Close()
|
||||
wb64.Close()
|
||||
raw, err = MakeMessage(MsgClassParameters, string(b.Bytes()), true)
|
||||
if err != nil {
|
||||
t.Fatalf(err.Error())
|
||||
}
|
||||
if string(raw) != `{"class":"parameters","parameters":{"someparam":"foo"}}` {
|
||||
t.Fatalf("Invalid module message class `parameters`")
|
||||
}
|
||||
|
||||
raw, err = MakeMessage(MsgClassStop, nil, false)
|
||||
if err != nil {
|
||||
t.Fatalf(err.Error())
|
||||
}
|
||||
|
|
|
@ -88,7 +88,7 @@ func fileModuleLocator(pattern string, regex bool, root string, depth int) ([]st
|
|||
args = append(args, "-maxdepth", strconv.Itoa(depth))
|
||||
param, err := run.(modules.HasParamsParser).ParamsParser(args)
|
||||
|
||||
buf, err := modules.MakeMessage(modules.MsgClassParameters, param)
|
||||
buf, err := modules.MakeMessage(modules.MsgClassParameters, param, false)
|
||||
if err != nil {
|
||||
return ret, nil
|
||||
}
|
||||
|
|
Загрузка…
Ссылка в новой задаче