dPaste

import (
    "crypto/rand"
    "errors"
    "os"
    "plugin"
    "reflect"
    "sync"
    "unsafe"
)
​
// This is the actual module we have loaded
type ModQuery struct {
    Module       *plugin.Plugin
    MethodMap    map[string]reflect.Value
    ParamMap     map[string][]reflect.Value
    Result       []interface{}
    Error        error
    ReplyTracker chan interface{}
    ModuleName   string
}
​
// This returns replies from function calls
type FunctionReply struct {
    ModuleName    string
    FunctionHash  [16]byte
    ReturnData    []interface{}
    OutputChannel chan []interface{}
    Error         error
}
​
// This structure contains all the loaded modules
type LoadedModules struct {
    Loaded        map[string]*ModQuery
    InUseMap      map[string]bool
    ActiveQueries map[string]map[[16]byte]bool
    Signal        chan string
    ReplyTracker  chan interface{}
    Mutex         *sync.Mutex
}
​
type CallTrack struct {
    ModuleName string
    Hash       [16]byte
}
​
// Create the overarching module handler
func NewModuleHandler() *LoadedModules {
    return &LoadedModules{
        Loaded:        make(map[string]*ModQuery),
        Signal:        make(chan string),
        InUseMap:      make(map[string]bool),
        ReplyTracker:  make(chan interface{}),
        ActiveQueries: make(map[string]map[[16]byte]bool),
        Mutex:         &sync.Mutex{},
    }
}
​
// Create a new query for a loaded module
func NewQuery(ModuleName string, ReplyTracker chan interface{}) *ModQuery {
    return &ModQuery{
        MethodMap:    make(map[string]reflect.Value),
        ParamMap:     make(map[string][]reflect.Value),
        ReplyTracker: ReplyTracker,
        ModuleName:   ModuleName,
    }
}
​
// Handle replies we receive
func (lm *LoadedModules) ReplyHandler(InputChannel chan interface{}) {
    var GotReply = false
    for {
        select {
        case input := <-InputChannel:
            switch input.(type) {
            case *FunctionReply:
                fr := input.(*FunctionReply)
                lm.Mutex.Lock()
                if mod, ok := lm.ActiveQueries[fr.ModuleName]; ok {
                    if _, ok := mod[fr.FunctionHash]; ok {
                        GotReply = true
                    }
                }
                lm.Mutex.Unlock()
            case *CallTrack:
                lm.Mutex.Lock()
                ct := input.(*CallTrack)
                if _, ok := lm.ActiveQueries[ct.ModuleName]; !ok {
                    lm.ActiveQueries[ct.ModuleName] = make(map[[16]byte]bool)
                }
                lm.ActiveQueries[ct.ModuleName][ct.Hash] = true
                lm.Mutex.Unlock()
            }
            if GotReply {
                fr := input.(*FunctionReply)
                fr.OutputChannel <- fr.ReturnData
                lm.Mutex.Lock()
                delete(lm.ActiveQueries[fr.ModuleName], fr.FunctionHash)
                if len(lm.ActiveQueries[fr.ModuleName]) == 0 {
                    lm.InUseMap[fr.ModuleName] = false
                }
                lm.Mutex.Unlock()
            }
        }
    }
}
​
// Load a module from a file - this should be build with --buildmode=plugin
func (lm *LoadedModules) AddModule(fName, mName string) error {
    var err error
    if _, modLoaded := lm.Loaded[mName]; modLoaded {
        return errors.New("module was already loaded")
    }
    if _, err := os.Stat(fName); err != nil {
        if os.IsNotExist(err) {
            return errors.New("specified module doesnt exist")
        }
        return errors.New("error getting module file data")
    }
    NewMod := NewQuery(mName, lm.ReplyTracker)
    if NewMod.Module, err = plugin.Open(fName); err != nil {
        return err
    }
    PlugMap := reflect.ValueOf(NewMod).Elem().Field(3)
    MapValues := reflect.NewAt(PlugMap.Type(), unsafe.Pointer(PlugMap.UnsafeAddr())).Elem()
    MethodMap := MapValues.Interface().(map[string]interface{})
    for MethodName, Method := range MethodMap {
        if reflect.ValueOf(Method).Kind() == reflect.Func {
            err := NewMod.addMethod(MethodName, reflect.ValueOf(Method))
            if err != nil {
                d.Log(LogHiveSlave, "Failed to load method from module: [s]: %v", MethodName, err)
            }
        }
    }
​
    lm.Mutex.Lock()
    lm.Loaded[mName] = NewMod
    return nil
}
​
// Remove a module that's been added
func (lm *LoadedModules) RemoveModule(mName string) {
    var RemovalOk = false
    lm.Mutex.Lock()
    if _, modLoaded := lm.Loaded[mName]; modLoaded {
        if Used, Exists := lm.InUseMap[mName]; Exists && !Used {
            RemovalOk = true
        }
    }
    if RemovalOk {
        lm.Loaded[mName] = nil
        delete(lm.Loaded, mName)
        delete(lm.InUseMap, mName)
        delete(lm.ActiveQueries, mName)
    }
    lm.Mutex.Unlock()
}
​
// Simplistic way to generate a hash for function call tracking
func NewHash() [16]byte {
    var result [16]byte
    var HashBytes = make([]byte, 8)
    _, _ = rand.Read(HashBytes)
    copy((*[8]byte)(unsafe.Pointer(&result[0]))[:], HashBytes)
    return result
}
​
// Execute a function from the specific module
func (lm *LoadedModules) ExecuteFunc(
    Method,
    Module string,
    retChan chan []interface{},
    parameters ...interface{}) error {
​
    var mod *ModQuery
    var verified bool
​
    if mod, verified = lm.Loaded[Module]; !verified {
        return errors.New("module not loaded")
    } else {
        QueryHash := NewHash()
        lm.ReplyTracker <- &CallTrack{
            ModuleName: Module,
            Hash:       QueryHash,
        }
        go mod.ExecuteMethod(Method, retChan, QueryHash, parameters)
        return nil
    }
}
​
// Add a method from a module to a point where it is executable
func (q *ModQuery) addMethod(mName string, mFunc reflect.Value) error {
    if mFunc.Kind() != reflect.Func {
        return errors.New("need a function as an input")
    } else {
        q.MethodMap[mName] = mFunc
    }
    return nil
}
​
func (q *ModQuery) ExecuteMethod(
    Method string, ReturnChan chan []interface{}, Hash [16]byte, params ...interface{}) {
​
    var staticParamCount int
    if f, haveFunc := q.MethodMap[Method]; haveFunc {
        if len(params) > 0 {
            // Generate an array of parameters as reflection values
            ParameterValues := generateParameters(params)
            // Get the type of function we will be executing so we can pass its input parameters
            funcType := reflect.TypeOf(f.Interface())
            // Now check if this function is varadic
            if funcType.IsVariadic() {
                // Get the type we need to input for varadic parameters
                vType := funcType.In(funcType.NumIn() - 1)
                staticParamCount = funcType.NumIn() - 1
                if len(ParameterValues) < staticParamCount {
                    newErr := errors.New("incorrect parameter count")
                    q.ReplyTracker <- FuncReplyError(q.ModuleName, Hash, newErr, ReturnChan)
                    return
                }
                if len(ParameterValues) > staticParamCount {
                    vParams := ParameterValues[staticParamCount:]
                    for _, vp := range vParams {
                        if reflect.SliceOf(vp.Type()) != vType {
                            newErr := errors.New("incorrectly type varadic argument")
                            q.ReplyTracker <- FuncReplyError(q.ModuleName, Hash, newErr, ReturnChan)
                            return
                        }
                    }
                }
            } else {
                // Function was not varadic, verify function input arguments
                staticParamCount = funcType.NumIn()
                if len(ParameterValues) != staticParamCount {
                    newErr := errors.New("invalid parameter count to non-varadic function")
                    q.ReplyTracker <- FuncReplyError(q.ModuleName, Hash, newErr, ReturnChan)
                    return
                }
            }
            // Verify the typing of static parameters
            for i := 0; i < staticParamCount; i++ {
                if funcType.In(i).Kind() != ParameterValues[i].Kind() {
                    newErr := errors.New("invalid static parameter type")
                    q.ReplyTracker <- FuncReplyError(q.ModuleName, Hash, newErr, ReturnChan)
                    return
                }
            }
            go q.SendResults(Hash, ReturnChan, f.Call(ParameterValues))
        } else {
            go q.SendResults(Hash, ReturnChan, f.Call([]reflect.Value{}))
        }
    } else {
        newErr := errors.New("method not found")
        q.ReplyTracker <- FuncReplyError(q.ModuleName, Hash, newErr, ReturnChan)
    }
}
​
// Convert a slice of interfaces to a slice of reflect.Value's
func generateParameters(p []interface{}) []reflect.Value {
    var result = make([]reflect.Value, len(p))
    for i, param := range p {
        result[i] = reflect.ValueOf(param)
    }
    return result
}
​
// Generate an error function reply
func FuncReplyError(mn string, h [16]byte, err error, Out chan []interface{}) *FunctionReply {
    return &FunctionReply{
        ModuleName:    mn,
        FunctionHash:  h,
        ReturnData:    nil,
        OutputChannel: Out,
        Error:         err,
    }
}
​
// SendResults sends the results of a called function to the reply tracker
func (q *ModQuery) SendResults(Hash [16]byte, Out chan []interface{}, FuncResult []reflect.Value) {
    result := &FunctionReply{
        ModuleName:    q.ModuleName,
        FunctionHash:  Hash,
        OutputChannel: Out,
        Error:         nil,
    }
    if len(FuncResult) > 0 {
        result.ReturnData = make([]interface{}, len(FuncResult))
        for i, r := range FuncResult {
            result.ReturnData[i] = r.Interface()
        }
    }
    q.ReplyTracker <- result
}