From cf37f0055b776af19647dddb8f0a0d7100bcd6bc Mon Sep 17 00:00:00 2001
From: Konstantin Tyukalov <v-ktyukalov@microsoft.com>
Date: Thu, 19 Sep 2024 19:44:29 +0400
Subject: [PATCH] Add wmi util

---
 src/Agent.Sdk/Util/WmiUtil.cs | 48 +++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 src/Agent.Sdk/Util/WmiUtil.cs

diff --git a/src/Agent.Sdk/Util/WmiUtil.cs b/src/Agent.Sdk/Util/WmiUtil.cs
new file mode 100644
index 0000000000..9879500cad
--- /dev/null
+++ b/src/Agent.Sdk/Util/WmiUtil.cs
@@ -0,0 +1,48 @@
+
+using System.Management;
+using System.Runtime.Versioning;
+
+namespace Agent.Sdk.Util;
+
+[SupportedOSPlatform("windows")]
+public class WmiUtil
+{
+    public static Task<List<ManagementBaseObject>> QueryGet(string query, CancellationToken cancellationToken)
+    {
+        var output = new List<ManagementBaseObject>();
+        var completionSource = new TaskCompletionSource<List<ManagementBaseObject>>();
+
+        var observer = new ManagementOperationObserver();
+        observer.ObjectReady += (sender, obj) =>
+        {
+            output.Add(obj.NewObject);
+        };
+        observer.Completed += (sender, e) =>
+        {
+            switch (e.Status)
+            {
+                case ManagementStatus.CallCanceled:
+                    completionSource.SetCanceled(cancellationToken);
+                    break;
+
+                case ManagementStatus.NoError:
+                    completionSource.SetResult(output);
+                    break;
+
+                default:
+                    completionSource.SetException(new Exception($"WMI Get Query failed with status {e.Status}"));
+                    break;
+            }
+        };
+
+        cancellationToken.Register(() =>
+        {
+            observer.Cancel();
+        });
+
+        using var searcher = new ManagementObjectSearcher(query);
+        searcher.Get(observer);
+
+        return completionSource.Task;
+    }
+}