diff --git a/examples/samples_with_config/wmic.go b/examples/samples_with_config/wmic.go index a0df38e9..a594770b 100644 --- a/examples/samples_with_config/wmic.go +++ b/examples/samples_with_config/wmic.go @@ -8,6 +8,7 @@ import ( "flag" "fmt" "os" + "runtime" "time" "github.com/oiweiwei/go-msrpc/dcerpc" @@ -180,6 +181,8 @@ func main() { flags |= wmi.QueryFlagType(wmi.GenericFlagTypeForwardOnly) } + now := time.Now() + enum, err := svcs.ExecQuery(ctx, &iwbemservices.ExecQueryRequest{ This: &dcom.ORPCThis{Version: srv.COMVersion}, QueryLanguage: &oaut.String{Data: "WQL"}, @@ -244,14 +247,15 @@ func main() { return } - var cls wmio.Class - - now := time.Now() - if limit > 0 && limit < page { page = limit } + // classes should store the class definitions across the calls. + var classes = make(map[string]*wmio.Class) + + count, buffer := 0, 0 + for i := 0; limit == 0 || i < limit; i += page { ret, err := wco.Next(ctx, &iwbemwcosmartenum.NextRequest{ @@ -270,7 +274,7 @@ func main() { break } - oa, err := wmi.UnmarshalObjectArrayWithClass(ret.Buffer, cls) + oa, err := wmi.UnmarshalObjectArrayWithClasses(ret.Buffer, classes) if err != nil { fmt.Fprintln(os.Stderr, "unmarshal_object_array_with_class", err) return @@ -280,13 +284,33 @@ func main() { if po.Object.Class != nil { fmt.Println(J(po.Object.Properties())) } else { - cls = po.Object.Instance.CurrentClass fmt.Println(J(po.Object.Values())) } } + + count += len(oa.Objects) + buffer += len(ret.Buffer) } + fmt.Fprintln(os.Stderr, "buffer fetched:", buffer/1024, "KiB") + fmt.Fprintln(os.Stderr, "records fetched:", count) fmt.Fprintln(os.Stderr, "query execution time:", time.Now().Sub(now)) fmt.Fprintln(os.Stderr, "script execution time:", time.Now().Sub(startTime)) + PrintMemUsage() + +} + +func PrintMemUsage() { + var m runtime.MemStats + runtime.ReadMemStats(&m) + // For info on each, see: https://golang.org/pkg/runtime/#MemStats + fmt.Fprintf(os.Stderr, "alloc / total / sys / numgc: %v MiB / %v MiB / %v MiB / %v\n", + bToMb(m.Alloc), + bToMb(m.TotalAlloc), + bToMb(m.Sys), + m.NumGC) +} +func bToMb(b uint64) uint64 { + return b / 1024 / 1024 } diff --git a/examples/wmic.go b/examples/wmic.go index 90fe0c4a..d77f8518 100644 --- a/examples/wmic.go +++ b/examples/wmic.go @@ -192,6 +192,8 @@ func main() { flags |= wmi.QueryFlagType(wmi.GenericFlagTypeForwardOnly) } + now := time.Now() + enum, err := svcs.ExecQuery(ctx, &iwbemservices.ExecQueryRequest{ This: &dcom.ORPCThis{Version: srv.COMVersion}, QueryLanguage: &oaut.String{Data: "WQL"}, @@ -256,9 +258,7 @@ func main() { return } - var cls wmio.Class - - now := time.Now() + var classes = make(map[string]*wmio.Class) if limit > 0 && limit < page { page = limit @@ -282,7 +282,7 @@ func main() { break } - oa, err := wmi.UnmarshalObjectArrayWithClass(ret.Buffer, cls) + oa, err := wmi.UnmarshalObjectArrayWithClasses(ret.Buffer, classes) if err != nil { fmt.Fprintln(os.Stderr, "unmarshal_object_array_with_class", err) return @@ -292,7 +292,6 @@ func main() { if po.Object.Class != nil { fmt.Println(j(po.Object.Properties())) } else { - cls = po.Object.Instance.CurrentClass fmt.Println(j(po.Object.Values())) } } diff --git a/msrpc/dcom/wmi/object_array.go b/msrpc/dcom/wmi/object_array.go index 43d6f109..fd7a191b 100644 --- a/msrpc/dcom/wmi/object_array.go +++ b/msrpc/dcom/wmi/object_array.go @@ -13,6 +13,14 @@ func UnmarshalObjectArrayWithClass(b []byte, cls wmio.Class) (*ObjectArray, erro return o, o.DecodeWithClass(wmio.NewCodec(b), cls) } +// UnmarshalObjectArrayWithClasses decodes the ObjectArray structure +// using the classes lookup map. Note that classes lookup map is updated +// during the call and must be reused for the next object decoding. +func UnmarshalObjectArrayWithClasses(b []byte, classes map[string]*wmio.Class) (*ObjectArray, error) { + o := &ObjectArray{} + return o, o.DecodeWithClasses(wmio.NewCodec(b), classes) +} + // The ObjectArray structure MUST be used to encode multiple CIM // objects that are returned in response to the IWbemWCOSmartEnum::Next // method. This structure is also used to encode parameters of the @@ -102,7 +110,25 @@ type DataPacketObject struct { Object *wmio.Object } +var ( + // Object is type WBEMOBJECT_CLASS: + // Structure contains the complete CIM Class definition. + ObjectTypeClass uint8 = 1 + // Object is type WBEMOBJECT_INSTANCE: + // Structure contains the complete CIM Instance definition. + ObjectTypeInstance uint8 = 2 + // Object is type WBEMOBJECT_INSTANCE_NOCLASS. + // Structure contains CIM Instance without the CIM Class definition. + ObjectTypeInstanceNoClass uint8 = 3 +) + +var DefaultClassID = "" + func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error { + return o.DecodeWithClasses(r, map[string]*wmio.Class{DefaultClassID: &cls}) +} + +func (o *ObjectArray) DecodeWithClasses(r *wmio.Codec, classes map[string]*wmio.Class) error { r.ReadData(&o.ByteOrdering) o.Signature = make([]byte, 8) @@ -120,6 +146,12 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error { for i := 0; i < int(o.NumObjects); i++ { + var cls wmio.Class + + if _, ok := classes[DefaultClassID]; ok { + cls = *classes[DefaultClassID] + } + var po DataPacketObject r.ReadData(&po.SizeOfHeader) @@ -136,12 +168,15 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error { var err error switch po.ObjectType { - case 1: + case ObjectTypeClass: + if po.Object, err = wmio.UnmarshalWithClass(r.Bytes()[:po.SizeOfData2], wmio.Class{}); err != nil { return err } + cls = po.Object.Class.CurrentClass - case 2: + + case ObjectTypeInstance: if po.ClassID, err = ReadClassID(r); err != nil { return err @@ -152,18 +187,28 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error { } cls = po.Object.Instance.CurrentClass - case 3: - if cls.Name == "" { - return fmt.Errorf("class name is empty") - } + + // store classes into map. + classes[po.ClassID.String()] = &cls + + case ObjectTypeInstanceNoClass: if po.ClassID, err = ReadClassID(r); err != nil { return err } + if _, ok := classes[po.ClassID.String()]; ok { + cls = *(classes[po.ClassID.String()]) + } + + if cls.Name == "" { + return fmt.Errorf("class name is empty") + } + if po.Object, err = wmio.UnmarshalWithClass(r.Bytes()[:po.SizeOfData2], cls); err != nil { return err } + default: return fmt.Errorf("unknown class type %d", po.ObjectType) } diff --git a/msrpc/dcom/wmio/codec.go b/msrpc/dcom/wmio/codec.go index a9e676b2..dc59a669 100644 --- a/msrpc/dcom/wmio/codec.go +++ b/msrpc/dcom/wmio/codec.go @@ -157,6 +157,10 @@ func (c *Codec) DecodeWithSize32(sz uint32, data any) error { return nil } + if sz > uint32(c.buf.Len()) { + return c.Errf("decode_with_size32: expected size is greater than actual: %d > %d", sz, c.buf.Len()) + } + b := make([]byte, sz) c.ReadData(b) diff --git a/msrpc/dcom/wmio/wmio.go b/msrpc/dcom/wmio/wmio.go index b092f312..216c111f 100644 --- a/msrpc/dcom/wmio/wmio.go +++ b/msrpc/dcom/wmio/wmio.go @@ -123,7 +123,18 @@ func (o *Object) Values() Values { var values = make(Values) for _, prop := range o.Instance.Properties { - values[prop.Name] = prop.Value.Value + switch value := prop.Value.Value.(type) { + case *Object: + values[prop.Name] = value.Values() + case []*Object: + vls := make([]any, len(value)) + for i := range value { + vls[i] = value[i].Values() + } + values[prop.Name] = vls + default: + values[prop.Name] = value + } } return values