Skip to content

Commit

Permalink
Merge pull request #11 from madhanrm/monitor
Browse files Browse the repository at this point in the history
Adding Monitor and Events
  • Loading branch information
madhanrm authored Apr 22, 2020
2 parents e7056ff + c8b581a commit 8589a0d
Show file tree
Hide file tree
Showing 147 changed files with 2,744 additions and 21,729 deletions.
3 changes: 1 addition & 2 deletions Makefile
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
GOCMD=go
GOBUILD=$(GOCMD) build -v -mod=vendor
GOBUILD=$(GOCMD) build -v
GOHOSTOS=$(strip $(shell $(GOCMD) env get GOHOSTOS))

TAG ?= $(shell git describe --tags)
Expand All @@ -14,7 +14,6 @@ all: format library

.PHONY: vendor
vendor:
GO111MODULE=on go mod vendor -v
GO111MODULE=on go mod tidy

library:
Expand Down
87 changes: 87 additions & 0 deletions pkg/base/event/event.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
// +build windows
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

package event

import (
"github.com/microsoft/wmi/pkg/base/session"
wmi "github.com/microsoft/wmi/pkg/wmiinstance"
)

type CallbackContext struct {
callbackData string
callback func(*wmi.WmiInstance, string)
}

// NewCallbackContext
func NewCallbackContext(cb func(*wmi.WmiInstance, string), data string) *CallbackContext {
return &CallbackContext{
callback: cb,
callbackData: data,
}
}

// Execute the callback
func (cb *CallbackContext) Execute(instance *wmi.WmiInstance) {
if cb.callback == nil {
return
}
cb.callback(instance, cb.callbackData)
}

func onInstanceReady(ctx interface{}, wmiInstances []*wmi.WmiInstance) {
context := ctx.(*CallbackContext)
if context == nil {
return
}

if len(wmiInstances) < 1 {
return
}

context.Execute(wmiInstances[0])
}

func onCompleted(ctx interface{}, wmiInstances []*wmi.WmiInstance) {
// Not used
}

func onProgress(ctx interface{}, wmiInstances []*wmi.WmiInstance) {
// Not used
}

func onInstancePut(ctx interface{}, wmiInstances []*wmi.WmiInstance) {
// Not used
}

// RegisterWmiCallback
func RegisterWmiCallback(context *CallbackContext, wmiNamespace, hostName, queryString string) (eventSink *wmi.WmiEventSink, err error) {
wmiSession, err := session.GetSession(wmiNamespace, hostName, "", "", "")
if err != nil {
return nil, err
}

eventSink, err = wmi.CreateWmiEventSink(wmiSession, context, onInstanceReady, onCompleted, onProgress, onInstancePut)
if err != nil {
return
}
defer func() {
if err != nil {
eventSink.Close()
eventSink = nil
}
}()

_, err = eventSink.Connect()
if err != nil {
return
}

_, err = wmiSession.ExecNotificationQueryAsync(eventSink, queryString)
if err != nil {
return
}

return
}
6 changes: 5 additions & 1 deletion pkg/base/host/host.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,13 @@ type WmiHost struct {
credential *credential.WmiCredential
}

func NewWmiLocalHost() *WmiHost {
return NewWmiHost("localhost")
}

// NewWmiHost
func NewWmiHost(hostname string) *WmiHost {
return &WmiHost{HostName: hostname}
return NewWmiHostWithCredential(hostname, "", "", "")
}

// NewWmiHostWithCredential
Expand Down
27 changes: 25 additions & 2 deletions pkg/base/instance/instancemanager.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,15 @@
package instance

import (
"fmt"
//"log"
"strings"
"sync"

"github.com/microsoft/wmi/pkg/base/credential"
"github.com/microsoft/wmi/pkg/base/host"
"github.com/microsoft/wmi/pkg/base/query"
wmisession "github.com/microsoft/wmi/pkg/base/session"
"github.com/microsoft/wmi/pkg/errors"
wmi "github.com/microsoft/wmi/pkg/wmiinstance"
)

Expand Down Expand Up @@ -46,6 +47,12 @@ func newWmiInstanceManager(hostname, namespaceName, userName, password, domainNa

}

func GetWmiInstanceManagerFromWHost(whost *host.WmiHost, namespaceName string) (*WmiInstanceManager, error) {
return GetWmiInstanceManagerFromCred(whost.HostName, namespaceName, whost.GetCredential())
}
func GetWmiInstanceManagerFromCred(hostname, namespaceName string, cred *credential.WmiCredential) (*WmiInstanceManager, error) {
return GetWmiInstanceManager(hostname, namespaceName, cred.UserName, cred.Password, cred.Domain)
}
func GetWmiInstanceManager(hostname, namespaceName, userName, password, domainName string) (*WmiInstanceManager, error) {
mapId := strings.Join([]string{hostname, namespaceName, domainName}, "_")
if val, ok := instanceManagerMap[mapId]; ok {
Expand Down Expand Up @@ -74,9 +81,11 @@ func (im *WmiInstanceManager) QueryInstanceEx(queryString string) (*wmi.WmiInsta
}

if len(instances) == 0 {
return nil, fmt.Errorf("No Instance Found")
return nil, errors.Wrapf(errors.NotFound, "Query [%s] failed with no instance", queryString)
}

//log.Printf("QueryInstanceEx [%s]=>[%d]instances\n", queryString, len(instances))

return instances[0], nil
}

Expand Down Expand Up @@ -104,3 +113,17 @@ func GetWmiInstance(hostname, namespaceName, userName, password, domainName stri
}
return im.QueryInstance(inquery)
}

func GetWmiInstancesFromHost(host *host.WmiHost, namespaceName string, inquery *query.WmiQuery) (*wmi.WmiInstanceCollection, error) {
im, err := GetWmiInstanceManagerFromWHost(host, namespaceName)
if err != nil {
return nil, err
}
instances, err := im.QueryInstances(inquery.String())
if err != nil {
return nil, err
}
winstances := wmi.WmiInstanceCollection{}
winstances = append(winstances, instances...)
return &winstances, nil
}
101 changes: 101 additions & 0 deletions pkg/base/monitor/monitor.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
// +build windows
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

package monitor

import (
"fmt"
"sync"

"github.com/microsoft/wmi/pkg/base/event"
"github.com/microsoft/wmi/pkg/base/query"
"github.com/microsoft/wmi/pkg/constant"
wmi "github.com/microsoft/wmi/pkg/wmiinstance"
)

// Monitor is a generic monitor to subscribe to Wmi Events based on a Query String
type Monitor struct {
mux sync.Mutex
// eventSinks for each of the entities that are being monitored here
// multiple event sinks can be setup for a single entity
eventSinks map[string][]*wmi.WmiEventSink
propertyNameToQuery string
wmiNamespaceName string
callbackContext interface{}
callbackFunction func(interface{}, string)
}

func (c *Monitor) onModified(instance *wmi.WmiInstance, cbData string) {
/*
prop, err := instance.GetProperty(c.propertyNameToQuery)
if err != nil {
fmt.Printf("Err: %v - [%s]\n", err, c.propertyNameToQuery)
return
}
propVal, ok := prop.(string)
if !ok {
return
}
*/
c.callbackFunction(c.callbackContext, cbData)
return
}

// CreateMonitor createa a new Monitor
func CreateMonitor(wmiNamespace string, callbackContext interface{},
callback func(interface{}, string)) *Monitor {
return &Monitor{
wmiNamespaceName: wmiNamespace,
eventSinks: map[string][]*wmi.WmiEventSink{},
callbackFunction: callback,
callbackContext: callbackContext,
}
}

// AddEntityWithFilter would add the entity to be monitored for changes
func (c *Monitor) AddEntityWithFilter(entityName, wqlQueryString string, filters query.WmiQueryFilterCollection) (err error) {
if _, ok := c.eventSinks[entityName]; ok {
return nil // error
}
qString := fmt.Sprintf("%s %s", wqlQueryString, filters.String())
fmt.Printf("Event Query [%s]\n", qString)
esink, err := c.getEventSink(entityName, qString)
if err != nil {
return
}
c.mux.Lock()
defer c.mux.Unlock()
c.eventSinks[entityName] = append(c.eventSinks[entityName], esink)
return
}

// RemoveEntity to remove the entity being monitored for changes
func (c *Monitor) RemoveEntity(entityName string) (err error) {
if _, ok := c.eventSinks[entityName]; !ok {
return nil // error NOT Found
}
c.mux.Lock()
defer c.mux.Unlock()

delete(c.eventSinks, entityName)
return
}

// Close the monitor
func (c *Monitor) Close() error {
for k := range c.eventSinks {
for _, s := range c.eventSinks[k] {
s.Close()
}
delete(c.eventSinks, k)
}
return nil
}

// getEventSink
func (c *Monitor) getEventSink(entityName, wqlQueryString string) (*wmi.WmiEventSink, error) {
//qString := fmt.Sprintf("%s %s", queryString, additionalFilter)
context := event.NewCallbackContext(c.onModified, entityName)
return event.RegisterWmiCallback(context, c.wmiNamespaceName, constant.HostName, wqlQueryString)
}
19 changes: 19 additions & 0 deletions pkg/base/query/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ import (
"fmt"
)

// https://docs.microsoft.com/en-us/windows/win32/wmisdk/wql-sql-for-wmi

type CompareOperator string

const (
Expand All @@ -16,6 +18,7 @@ const (
GreaterThanEquals CompareOperator = ">="
NotEquals CompareOperator = "<>"
Like CompareOperator = "LIKE"
Isa CompareOperator = "ISA"
)

type WmiQueryFilter struct {
Expand Down Expand Up @@ -56,6 +59,12 @@ func (q *WmiQueryFilter) String() string {
}
}

// HasFilter
func (q *WmiQuery) HasFilter() bool {
return len(q.Filters) > 0
}

// String
func (q *WmiQuery) String() (queryString string) {
queryString = fmt.Sprintf("SELECT * FROM %s", q.ClassName)

Expand All @@ -72,3 +81,13 @@ func (q *WmiQuery) String() (queryString string) {
queryString = queryString + fmt.Sprintf(" %s ", q.Filters[len(q.Filters)-1].String())
return
}

type WmiQueryFilterCollection []*WmiQueryFilter

func (c *WmiQueryFilterCollection) String() string {
queryString := ""
for _, query := range *c {
queryString = fmt.Sprintf("%s AND %s", queryString, query.String())
}
return queryString
}
13 changes: 3 additions & 10 deletions pkg/base/session/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@ import (
var (
sessionManager *wmi.WmiSessionManager
sessionsMap map[string]*wmi.WmiSession
localHostName string
)

// StartWMI
func StartWMI() {
func init() {
localHostName, _ = os.Hostname()
sessionsMap = make(map[string]*wmi.WmiSession)
sessionManager = wmi.NewWmiSessionManager()
}
Expand Down Expand Up @@ -64,14 +65,6 @@ func GetSession(namespaceName string, serverName string, domain string, userName
return sessionsMap[sessionsMapId], nil
}

var (
localHostName string
)

func init() {
localHostName, _ = os.Hostname()
}

////////////// Private functions ////////////////////////////
func createSession(sessionName string, serverName string, domain string, username string, password string) (*wmi.WmiSession, error) {
// TODO: ideally, we should also compare the domain here.
Expand Down
2 changes: 0 additions & 2 deletions pkg/base/session/session_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@ package session
import "testing"

func TestGetSession(t *testing.T) {

StartWMI()
defer StopWMI()

_, err := GetSession("temp", "localhost", "", "", "")
Expand Down
17 changes: 17 additions & 0 deletions pkg/constant/constant.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

package constant

type WMINamespace string

const (
Virtualization WMINamespace = "root/virtualization/v2"
CimV2 WMINamespace = "root/cimv2"
StadardCimV2 WMINamespace = "root/standardcimv2"
FailoverCluster WMINamespace = "root/mscluster"
)

const (
HostName string = "localhost"
)
Loading

0 comments on commit 8589a0d

Please sign in to comment.