Serge Bazanski | cc25bdf | 2018-10-25 14:02:58 +0200 | [diff] [blame] | 1 | // +build windows |
| 2 | |
| 3 | package wmi |
| 4 | |
| 5 | import ( |
| 6 | "fmt" |
| 7 | "reflect" |
| 8 | "runtime" |
| 9 | "sync" |
| 10 | |
| 11 | "github.com/go-ole/go-ole" |
| 12 | "github.com/go-ole/go-ole/oleutil" |
| 13 | ) |
| 14 | |
| 15 | // SWbemServices is used to access wmi. See https://msdn.microsoft.com/en-us/library/aa393719(v=vs.85).aspx |
| 16 | type SWbemServices struct { |
| 17 | //TODO: track namespace. Not sure if we can re connect to a different namespace using the same instance |
| 18 | cWMIClient *Client //This could also be an embedded struct, but then we would need to branch on Client vs SWbemServices in the Query method |
| 19 | sWbemLocatorIUnknown *ole.IUnknown |
| 20 | sWbemLocatorIDispatch *ole.IDispatch |
| 21 | queries chan *queryRequest |
| 22 | closeError chan error |
| 23 | lQueryorClose sync.Mutex |
| 24 | } |
| 25 | |
| 26 | type queryRequest struct { |
| 27 | query string |
| 28 | dst interface{} |
| 29 | args []interface{} |
| 30 | finished chan error |
| 31 | } |
| 32 | |
| 33 | // InitializeSWbemServices will return a new SWbemServices object that can be used to query WMI |
| 34 | func InitializeSWbemServices(c *Client, connectServerArgs ...interface{}) (*SWbemServices, error) { |
| 35 | //fmt.Println("InitializeSWbemServices: Starting") |
| 36 | //TODO: implement connectServerArgs as optional argument for init with connectServer call |
| 37 | s := new(SWbemServices) |
| 38 | s.cWMIClient = c |
| 39 | s.queries = make(chan *queryRequest) |
| 40 | initError := make(chan error) |
| 41 | go s.process(initError) |
| 42 | |
| 43 | err, ok := <-initError |
| 44 | if ok { |
| 45 | return nil, err //Send error to caller |
| 46 | } |
| 47 | //fmt.Println("InitializeSWbemServices: Finished") |
| 48 | return s, nil |
| 49 | } |
| 50 | |
| 51 | // Close will clear and release all of the SWbemServices resources |
| 52 | func (s *SWbemServices) Close() error { |
| 53 | s.lQueryorClose.Lock() |
| 54 | if s == nil || s.sWbemLocatorIDispatch == nil { |
| 55 | s.lQueryorClose.Unlock() |
| 56 | return fmt.Errorf("SWbemServices is not Initialized") |
| 57 | } |
| 58 | if s.queries == nil { |
| 59 | s.lQueryorClose.Unlock() |
| 60 | return fmt.Errorf("SWbemServices has been closed") |
| 61 | } |
| 62 | //fmt.Println("Close: sending close request") |
| 63 | var result error |
| 64 | ce := make(chan error) |
| 65 | s.closeError = ce //Race condition if multiple callers to close. May need to lock here |
| 66 | close(s.queries) //Tell background to shut things down |
| 67 | s.lQueryorClose.Unlock() |
| 68 | err, ok := <-ce |
| 69 | if ok { |
| 70 | result = err |
| 71 | } |
| 72 | //fmt.Println("Close: finished") |
| 73 | return result |
| 74 | } |
| 75 | |
| 76 | func (s *SWbemServices) process(initError chan error) { |
| 77 | //fmt.Println("process: starting background thread initialization") |
| 78 | //All OLE/WMI calls must happen on the same initialized thead, so lock this goroutine |
| 79 | runtime.LockOSThread() |
| 80 | defer runtime.LockOSThread() |
| 81 | |
| 82 | err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED) |
| 83 | if err != nil { |
| 84 | oleCode := err.(*ole.OleError).Code() |
| 85 | if oleCode != ole.S_OK && oleCode != S_FALSE { |
| 86 | initError <- fmt.Errorf("ole.CoInitializeEx error: %v", err) |
| 87 | return |
| 88 | } |
| 89 | } |
| 90 | defer ole.CoUninitialize() |
| 91 | |
| 92 | unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator") |
| 93 | if err != nil { |
| 94 | initError <- fmt.Errorf("CreateObject SWbemLocator error: %v", err) |
| 95 | return |
| 96 | } else if unknown == nil { |
| 97 | initError <- ErrNilCreateObject |
| 98 | return |
| 99 | } |
| 100 | defer unknown.Release() |
| 101 | s.sWbemLocatorIUnknown = unknown |
| 102 | |
| 103 | dispatch, err := s.sWbemLocatorIUnknown.QueryInterface(ole.IID_IDispatch) |
| 104 | if err != nil { |
| 105 | initError <- fmt.Errorf("SWbemLocator QueryInterface error: %v", err) |
| 106 | return |
| 107 | } |
| 108 | defer dispatch.Release() |
| 109 | s.sWbemLocatorIDispatch = dispatch |
| 110 | |
| 111 | // we can't do the ConnectServer call outside the loop unless we find a way to track and re-init the connectServerArgs |
| 112 | //fmt.Println("process: initialized. closing initError") |
| 113 | close(initError) |
| 114 | //fmt.Println("process: waiting for queries") |
| 115 | for q := range s.queries { |
| 116 | //fmt.Printf("process: new query: len(query)=%d\n", len(q.query)) |
| 117 | errQuery := s.queryBackground(q) |
| 118 | //fmt.Println("process: s.queryBackground finished") |
| 119 | if errQuery != nil { |
| 120 | q.finished <- errQuery |
| 121 | } |
| 122 | close(q.finished) |
| 123 | } |
| 124 | //fmt.Println("process: queries channel closed") |
| 125 | s.queries = nil //set channel to nil so we know it is closed |
| 126 | //TODO: I think the Release/Clear calls can panic if things are in a bad state. |
| 127 | //TODO: May need to recover from panics and send error to method caller instead. |
| 128 | close(s.closeError) |
| 129 | } |
| 130 | |
| 131 | // Query runs the WQL query using a SWbemServices instance and appends the values to dst. |
| 132 | // |
| 133 | // dst must have type *[]S or *[]*S, for some struct type S. Fields selected in |
| 134 | // the query must have the same name in dst. Supported types are all signed and |
| 135 | // unsigned integers, time.Time, string, bool, or a pointer to one of those. |
| 136 | // Array types are not supported. |
| 137 | // |
| 138 | // By default, the local machine and default namespace are used. These can be |
| 139 | // changed using connectServerArgs. See |
| 140 | // http://msdn.microsoft.com/en-us/library/aa393720.aspx for details. |
| 141 | func (s *SWbemServices) Query(query string, dst interface{}, connectServerArgs ...interface{}) error { |
| 142 | s.lQueryorClose.Lock() |
| 143 | if s == nil || s.sWbemLocatorIDispatch == nil { |
| 144 | s.lQueryorClose.Unlock() |
| 145 | return fmt.Errorf("SWbemServices is not Initialized") |
| 146 | } |
| 147 | if s.queries == nil { |
| 148 | s.lQueryorClose.Unlock() |
| 149 | return fmt.Errorf("SWbemServices has been closed") |
| 150 | } |
| 151 | |
| 152 | //fmt.Println("Query: Sending query request") |
| 153 | qr := queryRequest{ |
| 154 | query: query, |
| 155 | dst: dst, |
| 156 | args: connectServerArgs, |
| 157 | finished: make(chan error), |
| 158 | } |
| 159 | s.queries <- &qr |
| 160 | s.lQueryorClose.Unlock() |
| 161 | err, ok := <-qr.finished |
| 162 | if ok { |
| 163 | //fmt.Println("Query: Finished with error") |
| 164 | return err //Send error to caller |
| 165 | } |
| 166 | //fmt.Println("Query: Finished") |
| 167 | return nil |
| 168 | } |
| 169 | |
| 170 | func (s *SWbemServices) queryBackground(q *queryRequest) error { |
| 171 | if s == nil || s.sWbemLocatorIDispatch == nil { |
| 172 | return fmt.Errorf("SWbemServices is not Initialized") |
| 173 | } |
| 174 | wmi := s.sWbemLocatorIDispatch //Should just rename in the code, but this will help as we break things apart |
| 175 | //fmt.Println("queryBackground: Starting") |
| 176 | |
| 177 | dv := reflect.ValueOf(q.dst) |
| 178 | if dv.Kind() != reflect.Ptr || dv.IsNil() { |
| 179 | return ErrInvalidEntityType |
| 180 | } |
| 181 | dv = dv.Elem() |
| 182 | mat, elemType := checkMultiArg(dv) |
| 183 | if mat == multiArgTypeInvalid { |
| 184 | return ErrInvalidEntityType |
| 185 | } |
| 186 | |
| 187 | // service is a SWbemServices |
| 188 | serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", q.args...) |
| 189 | if err != nil { |
| 190 | return err |
| 191 | } |
| 192 | service := serviceRaw.ToIDispatch() |
| 193 | defer serviceRaw.Clear() |
| 194 | |
| 195 | // result is a SWBemObjectSet |
| 196 | resultRaw, err := oleutil.CallMethod(service, "ExecQuery", q.query) |
| 197 | if err != nil { |
| 198 | return err |
| 199 | } |
| 200 | result := resultRaw.ToIDispatch() |
| 201 | defer resultRaw.Clear() |
| 202 | |
| 203 | count, err := oleInt64(result, "Count") |
| 204 | if err != nil { |
| 205 | return err |
| 206 | } |
| 207 | |
| 208 | enumProperty, err := result.GetProperty("_NewEnum") |
| 209 | if err != nil { |
| 210 | return err |
| 211 | } |
| 212 | defer enumProperty.Clear() |
| 213 | |
| 214 | enum, err := enumProperty.ToIUnknown().IEnumVARIANT(ole.IID_IEnumVariant) |
| 215 | if err != nil { |
| 216 | return err |
| 217 | } |
| 218 | if enum == nil { |
| 219 | return fmt.Errorf("can't get IEnumVARIANT, enum is nil") |
| 220 | } |
| 221 | defer enum.Release() |
| 222 | |
| 223 | // Initialize a slice with Count capacity |
| 224 | dv.Set(reflect.MakeSlice(dv.Type(), 0, int(count))) |
| 225 | |
| 226 | var errFieldMismatch error |
| 227 | for itemRaw, length, err := enum.Next(1); length > 0; itemRaw, length, err = enum.Next(1) { |
| 228 | if err != nil { |
| 229 | return err |
| 230 | } |
| 231 | |
| 232 | err := func() error { |
| 233 | // item is a SWbemObject, but really a Win32_Process |
| 234 | item := itemRaw.ToIDispatch() |
| 235 | defer item.Release() |
| 236 | |
| 237 | ev := reflect.New(elemType) |
| 238 | if err = s.cWMIClient.loadEntity(ev.Interface(), item); err != nil { |
| 239 | if _, ok := err.(*ErrFieldMismatch); ok { |
| 240 | // We continue loading entities even in the face of field mismatch errors. |
| 241 | // If we encounter any other error, that other error is returned. Otherwise, |
| 242 | // an ErrFieldMismatch is returned. |
| 243 | errFieldMismatch = err |
| 244 | } else { |
| 245 | return err |
| 246 | } |
| 247 | } |
| 248 | if mat != multiArgTypeStructPtr { |
| 249 | ev = ev.Elem() |
| 250 | } |
| 251 | dv.Set(reflect.Append(dv, ev)) |
| 252 | return nil |
| 253 | }() |
| 254 | if err != nil { |
| 255 | return err |
| 256 | } |
| 257 | } |
| 258 | //fmt.Println("queryBackground: Finished") |
| 259 | return errFieldMismatch |
| 260 | } |