Merge pull request #226 from Icinga:fix/add_type_classnames_and_added…

…_check Fix: Class names for Add-Type and check if class was already added Fixes class names of `Add-Type` for `Get-IcingaDiskAttributes` and `Get-IcingaUNCPathSize`. Also adds checks if the class was already added inside the context
Icinga · Aug 24, 2021 · 73fba17 · 73fba17
2 parents e4b1451 + 2a951ed
commit 73fba17
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 55 deletions.
diff --git a/doc/31-Changelog.md b/doc/31-Changelog.md
@@ -14,6 +14,7 @@ Released closed milestones can be found on [GitHub](https://github.com/Icinga/ic
 ## Bugfixes
 
 * [#224](https://github.com/Icinga/icinga-powershell-plugins/pull/224) Fixes `Invoke-IcingaCheckUsedPartitionSpace` label names, which might have caused conflicts with Graphite/InfluxDB, because of changed metric unit
+* [#226](https://github.com/Icinga/icinga-powershell-plugins/pull/226) Fixes class names of `Add-Type` for `Get-IcingaDiskAttributes` and `Get-IcingaUNCPathSize` and adds checks if the class was already added inside the session
 
 ### Enhancements
 

diff --git a/provider/disks/Get-IcingaDiskAttributes.psm1 b/provider/disks/Get-IcingaDiskAttributes.psm1
@@ -34,55 +34,57 @@ function Get-IcingaDiskAttributes()
     $DISK_ATTRIBUTE_OFFLINE         = 0x0000000000000001;
     $DISK_ATTRIBUTE_READ_ONLY       = 0x0000000000000002;
 
-    Add-Type -TypeDefinition @"
-        using System;
-        using System.IO;
-        using System.Diagnostics;
-        using System.Runtime.InteropServices;
+    if ((Test-IcingaAddTypeExist -Type 'IcingaDiskAttributes') -eq $FALSE) {
+        Add-Type -TypeDefinition @"
+            using System;
+            using System.IO;
+            using System.Diagnostics;
+            using System.Runtime.InteropServices;
 
-        public static class kernel32 {
-            [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
-            public static extern IntPtr CreateFile(
-                [MarshalAs(UnmanagedType.LPTStr)] string filename,
-                [MarshalAs(UnmanagedType.U4)] FileAccess access,
-                [MarshalAs(UnmanagedType.U4)] FileShare share,
-                IntPtr securityAttributes,
-                [MarshalAs(UnmanagedType.U4)] FileMode creationDisposition,
-                [MarshalAs(UnmanagedType.U4)] FileAttributes flagsAndAttributes,
-                IntPtr templateFile
-            );
+            public static class IcingaDiskAttributes {
+                [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
+                public static extern IntPtr CreateFile(
+                    [MarshalAs(UnmanagedType.LPTStr)] string filename,
+                    [MarshalAs(UnmanagedType.U4)] FileAccess access,
+                    [MarshalAs(UnmanagedType.U4)] FileShare share,
+                    IntPtr securityAttributes,
+                    [MarshalAs(UnmanagedType.U4)] FileMode creationDisposition,
+                    [MarshalAs(UnmanagedType.U4)] FileAttributes flagsAndAttributes,
+                    IntPtr templateFile
+                );
 
-            public struct Icinga_Disk_Data {
-                [MarshalAs(UnmanagedType.U4)]public UInt32 Version;
-                [MarshalAs(UnmanagedType.U4)]public UInt32 Reserved1;
-                [MarshalAs(UnmanagedType.U8)]public UInt64 Attributes;
-            }
+                public struct Icinga_Disk_Data {
+                    [MarshalAs(UnmanagedType.U4)]public UInt32 Version;
+                    [MarshalAs(UnmanagedType.U4)]public UInt32 Reserved1;
+                    [MarshalAs(UnmanagedType.U8)]public UInt64 Attributes;
+                }
 
-            [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
-            public static extern bool DeviceIoControl(
-                IntPtr hDevice,
-                uint dwIoControlCode,
-                IntPtr lpInBuffer,
-                uint nInBufferSize,
-                out Icinga_Disk_Data lpOutBuffer,
-                uint nOutBufferSize,
-                out uint lpBytesReturned,
-                IntPtr lpOverlapped
-            );
+                [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
+                public static extern bool DeviceIoControl(
+                    IntPtr hDevice,
+                    uint dwIoControlCode,
+                    IntPtr lpInBuffer,
+                    uint nInBufferSize,
+                    out Icinga_Disk_Data lpOutBuffer,
+                    uint nOutBufferSize,
+                    out uint lpBytesReturned,
+                    IntPtr lpOverlapped
+                );
 
-            [DllImport("kernel32.dll", SetLastError=true)]
-            public static extern bool CloseHandle(IntPtr hObject);
-        }
+                [DllImport("kernel32.dll", SetLastError=true)]
+                public static extern bool CloseHandle(IntPtr hObject);
+            }
 "@
+    }
 
     [bool]$DiskOffline  = $FALSE;
     [bool]$DiskReadOnly = $FALSE;
-    $KernelHandle       = [kernel32]::CreateFile($PhysicalDisk, 0, [System.IO.FileShare]::ReadWrite, [System.IntPtr]::Zero, [System.IO.FileMode]::Open, 0, [System.IntPtr]::Zero);
+    $KernelHandle       = [IcingaDiskAttributes]::CreateFile($PhysicalDisk, 0, [System.IO.FileShare]::ReadWrite, [System.IntPtr]::Zero, [System.IO.FileMode]::Open, 0, [System.IntPtr]::Zero);
 
     if ($KernelHandle) {
-        $DiskData = New-Object -TypeName Kernel32+Icinga_Disk_Data;
+        $DiskData = New-Object -TypeName IcingaDiskAttributes+Icinga_Disk_Data;
         $Value    = New-Object -TypeName UInt32;
-        $Result   = [kernel32]::DeviceIoControl($KernelHandle, $IOCTL_DISK_GET_DISK_ATTRIBUTES, [System.IntPtr]::Zero, 0, [ref]$DiskData, [System.Runtime.InteropServices.Marshal]::SizeOf($DiskData), [ref]$Value, [System.IntPtr]::Zero);
+        $Result   = [IcingaDiskAttributes]::DeviceIoControl($KernelHandle, $IOCTL_DISK_GET_DISK_ATTRIBUTES, [System.IntPtr]::Zero, 0, [ref]$DiskData, [System.Runtime.InteropServices.Marshal]::SizeOf($DiskData), [ref]$Value, [System.IntPtr]::Zero);
         if ($Result) {
             if (($DiskData.attributes -band $DISK_ATTRIBUTE_OFFLINE) -eq $DISK_ATTRIBUTE_OFFLINE) {
                 $DiskOffline = $TRUE;
@@ -91,7 +93,7 @@ function Get-IcingaDiskAttributes()
                 $DiskReadOnly = $TRUE;
             }
         }
-        $Result = [kernel32]::CloseHandle($KernelHandle);
+        $Result = [IcingaDiskAttributes]::CloseHandle($KernelHandle);
     }
 
     return @{

diff --git a/provider/disks/Get-IcingaUNCPathSize.psm1 b/provider/disks/Get-IcingaUNCPathSize.psm1
@@ -12,22 +12,24 @@ function Get-IcingaUNCPathSize()
             -Force;
     }
 
-    # Register our kernel32.dll Windows API function call
-Add-Type -TypeDefinition @"
-    using System;
-    using System.Runtime.InteropServices;
-
-    public static class kernel32 {
-        [DllImport("kernel32.dll", PreserveSig = true, CharSet = CharSet.Auto)]
-
-        public static extern int GetDiskFreeSpaceEx(
-            IntPtr lpDirectoryName,           // UNC Path for share
-            out long lpFreeBytesAvailable,    // Free Bytes available on path
-            out long lpTotalNumberOfBytes,    // Bytes available on target disk / path
-            out long lpTotalNumberOfFreeBytes // Total available space on target disk / path
-        );
-    }
+    if ((Test-IcingaAddTypeExist -Type 'IcingaUNCPath') -eq $FALSE) {
+        # Register our kernel32.dll Windows API function call
+        Add-Type -TypeDefinition @"
+            using System;
+            using System.Runtime.InteropServices;
+
+            public static class IcingaUNCPath {
+                [DllImport("kernel32.dll", PreserveSig = true, CharSet = CharSet.Auto)]
+
+                public static extern int GetDiskFreeSpaceEx(
+                    IntPtr lpDirectoryName,           // UNC Path for share
+                    out long lpFreeBytesAvailable,    // Free Bytes available on path
+                    out long lpTotalNumberOfBytes,    // Bytes available on target disk / path
+                    out long lpTotalNumberOfFreeBytes // Total available space on target disk / path
+                );
+            }
 "@
+    }
 
     # Setup variables as object which we can use to reference data into
     $ShareFree = New-Object -TypeName long;
@@ -38,7 +40,7 @@ Add-Type -TypeDefinition @"
     [System.IntPtr]$ptrPath = [System.Runtime.InteropServices.Marshal]::StringToHGlobalAuto($Path);
 
     # Call our function we registered within the Add-Type definition
-    [kernel32]::GetDiskFreeSpaceEx($ptrPath, [ref]$ShareFree, [ref]$ShareSize, [ref]$TotalFree) | Out-Null;
+    [IcingaUNCPath]::GetDiskFreeSpaceEx($ptrPath, [ref]$ShareFree, [ref]$ShareSize, [ref]$TotalFree) | Out-Null;
     $ShareFreePercent = 0;
 
     if ($ShareSize -ne 0) {