From 440eb6eca8a4f588107b9bfaede2c01a6260cf06 Mon Sep 17 00:00:00 2001
From: he3als <he3als@ameliorated.io>
Date: Mon, 25 Mar 2024 19:02:17 +0000
Subject: [PATCH] feat: enable required services on startup

---
 src/Misc/CheckService.cs | 84 ++++++++++++++++++++++++++++++++++++++++
 src/Program.cs           | 24 ++++++++++++
 src/amecs.csproj         |  1 +
 3 files changed, 109 insertions(+)
 create mode 100644 src/Misc/CheckService.cs

diff --git a/src/Misc/CheckService.cs b/src/Misc/CheckService.cs
new file mode 100644
index 0000000..73eacd6
--- /dev/null
+++ b/src/Misc/CheckService.cs
@@ -0,0 +1,84 @@
+using System;
+using System.ComponentModel;
+using System.Runtime.InteropServices;
+using System.ServiceProcess;
+
+namespace amecs.Misc
+{
+    public static class ServiceManager
+    {
+        [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+        private static extern bool ChangeServiceConfig(
+            IntPtr hService,
+            uint nServiceType,
+            uint nStartType,
+            uint nErrorControl,
+            string lpBinaryPathName,
+            string lpLoadOrderGroup,
+            IntPtr lpdwTagId,
+            [In] char[] lpDependencies,
+            string lpServiceStartName,
+            string lpPassword,
+            string lpDisplayName);
+
+        [DllImport("advapi32.dll", SetLastError = true, CharSet = CharSet.Auto)]
+        static extern IntPtr OpenService(
+            IntPtr hScManager, string lpServiceName, uint dwDesiredAccess);
+
+        [DllImport("advapi32.dll", EntryPoint = "OpenSCManagerW", ExactSpelling = true, CharSet = CharSet.Unicode,
+            SetLastError = true)]
+        private static extern IntPtr OpenSCManager(
+            string machineName, string databaseName, uint dwAccess);
+
+        [DllImport("advapi32.dll", EntryPoint = "CloseServiceHandle")]
+        private static extern int CloseServiceHandle(IntPtr hScObject);
+
+        private const uint ServiceNoChange = 0xFFFFFFFF;
+        private const uint ServiceQueryConfig = 0x00000001;
+        private const uint ServiceChangeConfig = 0x00000002;
+        private const uint ScManagerAllAccess = 0x000F003F;
+
+        public static void ChangeStartMode(ServiceController svc, ServiceStartMode mode)
+        {
+            var scManagerHandle = OpenSCManager(null, null, ScManagerAllAccess);
+            if (scManagerHandle == IntPtr.Zero)
+            {
+                throw new ExternalException("Open Service Manager Error");
+            }
+
+            var serviceHandle = OpenService(
+                scManagerHandle,
+                svc.ServiceName,
+                ServiceQueryConfig | ServiceChangeConfig);
+
+            if (serviceHandle == IntPtr.Zero)
+            {
+                throw new ExternalException("Open Service Error");
+            }
+
+            var result = ChangeServiceConfig(
+                serviceHandle,
+                ServiceNoChange,
+                (uint)mode,
+                ServiceNoChange,
+                null,
+                null,
+                IntPtr.Zero,
+                null,
+                null,
+                null,
+                null);
+
+            if (result == false)
+            {
+                var nError = Marshal.GetLastWin32Error();
+                var win32Exception = new Win32Exception(nError);
+                throw new ExternalException("Could not change service start type: "
+                    + win32Exception.Message);
+            }
+
+            CloseServiceHandle(serviceHandle);
+            CloseServiceHandle(scManagerHandle);
+        }
+    }
+}
diff --git a/src/Program.cs b/src/Program.cs
index 2bede50..08d0d19 100644
--- a/src/Program.cs
+++ b/src/Program.cs
@@ -2,9 +2,11 @@
 using System.DirectoryServices.AccountManagement;
 using System.Linq;
 using System.Security.Principal;
+using System.ServiceProcess;
 using Ameliorated.ConsoleUtils;
 using Microsoft.Win32;
 using amecs.Actions;
+using amecs.Misc;
 using Menu = Ameliorated.ConsoleUtils.Menu;
 
 namespace amecs
@@ -19,6 +21,28 @@ namespace amecs
         {
             ConsoleTUI.Initialize("Central AME Script");
 
+            // LanmanServer and LanmanWorkstation must both be enabled for amecs to function.
+            // This is due to the use of DirectoryServices.
+            try
+            {
+                var server = new ServiceController("LanmanServer");
+                var workstation = new ServiceController("LanmanWorkstation");
+
+                foreach (var service in new[] {server, workstation})
+                {
+                    if (service.Status == ServiceControllerStatus.Running) continue;
+                    ServiceManager.ChangeStartMode(service, ServiceStartMode.Automatic);
+                    service.Start();
+                    service.WaitForStatus(ServiceControllerStatus.Running, TimeSpan.FromMilliseconds(10000));
+                }
+            }
+            catch (Exception e)
+            {
+                ConsoleTUI.ShowErrorBox("Could not enable required services: " + e, "Central AME Script");
+                Environment.Exit(1);
+            }
+            
+            // Elevation
             try
             {
                 NSudo.GetSystemPrivilege();
diff --git a/src/amecs.csproj b/src/amecs.csproj
index 28a4d5f..d9174a9 100644
--- a/src/amecs.csproj
+++ b/src/amecs.csproj
@@ -211,6 +211,7 @@
         <Compile Include="Extra\NVCP.cs" />
         <Compile Include="Extra\WSH.cs" />
         <Compile Include="Globals.cs" />
+        <Compile Include="Misc\CheckService.cs" />
         <Compile Include="Misc\FolderPicker.cs" />
         <Compile Include="Misc\SelectWindowsImage.cs" />
         <Compile Include="NSudo.cs" />