BeatSaber-IPA-Reloaded/IPA.Loader/AntiMalware/WindowsWin32AntiMalware.cs

#nullable enable
using IPA.AntiMalware.ComAPI;
using IPA.Logging;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace IPA.AntiMalware
{
    internal class WindowsWin32AntiMalware : IAntiMalware, IDisposable
    {
        internal static WindowsWin32AntiMalware? TryInitialize()
        {
            try
            {
                return new();
            }
            catch (Exception e)
            {
                Logger.AntiMalware.Warn("Could not initialize antimalware engine:");
                Logger.AntiMalware.Warn(e);
                return null;
            }
        }

        private readonly IntPtr handle;
        private bool disposedValue;

        private WindowsWin32AntiMalware()
        {
            AmsiInitialize(AmsiConstants.AppName, out handle);
        }

        private static ScanResult ScanResultFromAmsiResult(AmsiResult result)
            => result switch
            {
                AmsiResult.Clean => ScanResult.KnownSafe,
                AmsiResult.NotDetected => ScanResult.NotDetected,
                AmsiResult.Detected => ScanResult.Detected,
                _ => ScanResult.MaybeMalware
            };

        public ScanResult ScanFile(FileInfo file)
        {
            var data = File.ReadAllBytes(file.FullName);
            return ScanData(data, file.FullName);
        }

        public ScanResult ScanData(byte[] data, string? contentName = null)
        {
            contentName ??= $"unknown_data_{Guid.NewGuid()}";

            AmsiScanBuffer(handle, data, (uint)data.Length, contentName, IntPtr.Zero, out var result);

            Logger.AntiMalware.Debug($"Scanned data named '{contentName}' and got '{result}'");
            return ScanResultFromAmsiResult(result);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // we have no disposable managed state
                }

                AmsiUninitialize(handle);
                disposedValue = true;
            }
        }

        ~WindowsWin32AntiMalware()
        {
            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
            Dispose(disposing: false);
        }

        public void Dispose()
        {
            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
            Dispose(disposing: true);
            GC.SuppressFinalize(this);
        }

        [DllImport("amsi", CallingConvention = CallingConvention.Winapi, CharSet = CharSet.Unicode, ExactSpelling = true)]
        [DefaultDllImportSearchPaths(DllImportSearchPath.System32)]
        private static extern void AmsiInitialize([MarshalAs(UnmanagedType.LPWStr)] string appName, [Out] out IntPtr handle);

        [DllImport("amsi", CallingConvention = CallingConvention.Winapi, CharSet = CharSet.Unicode, ExactSpelling = true)]
        [DefaultDllImportSearchPaths(DllImportSearchPath.System32)]
        private static extern void AmsiUninitialize(IntPtr handle);

        [DllImport("amsi", CallingConvention = CallingConvention.Winapi, CharSet = CharSet.Unicode, ExactSpelling = true)]
        [DefaultDllImportSearchPaths(DllImportSearchPath.System32)]
        private static extern void AmsiScanBuffer(IntPtr context,
            [MarshalAs(UnmanagedType.LPArray, SizeParamIndex = 2)] byte[] buffer, uint length,
            [MarshalAs(UnmanagedType.LPWStr)] string contentName, 
            IntPtr session, 
            [Out] out AmsiResult result);
    }
}