Merge pull request #176 from ameihm0912/action-compress

[medium] introduce optional action compression between client and agent
This commit is contained in:
Aaron Meihm 2016-01-22 14:51:47 -06:00
Родитель adf87c11f4 1d73fab48e
Коммит 3b34f1f853
13 изменённых файлов: 298 добавлений и 44 удалений

Просмотреть файл

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