BeatSaber-IPA-Reloaded/IPA/Patcher/Patcher.cs

using Mono.Cecil;
using Mono.Cecil.Cil;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace IPA.Patcher
{
    class PatchedModule
    {
        private static readonly string[] ENTRY_TYPES = { "Input", "Display" };

        private FileInfo _File;
        private ModuleDefinition _Module;

        internal struct PatchData {
            public bool IsPatched;
            public Version Version;
        }

        public static PatchedModule Load(string engineFile)
        {
            return new PatchedModule(engineFile);
        }

        private PatchedModule(string engineFile)
        {
            _File = new FileInfo(engineFile);

            LoadModules();
        }

        private void LoadModules()
        {
            var resolver = new DefaultAssemblyResolver();
            resolver.AddSearchDirectory(_File.DirectoryName);

            var parameters = new ReaderParameters
            {
                AssemblyResolver = resolver,
            };
            
            _Module = ModuleDefinition.ReadModule(_File.FullName, parameters);
        }
        
        public PatchData Data
        {
            get
            {
                var IIdata = new PatchData { IsPatched = false, Version = null };
                foreach (var @ref in _Module.AssemblyReferences) {
                    if (@ref.Name == "IllusionInjector") IIdata = new PatchData { IsPatched = true, Version = new Version(0, 0, 0, 0) }; 
                    if (@ref.Name == "IllusionPlugin") IIdata = new PatchData { IsPatched = true, Version = new Version(0, 0, 0, 0) };
                    if (@ref.Name == "IPA.Injector") return new PatchData { IsPatched = true, Version = @ref.Version };
                }
                return IIdata;
            }
        }

        public void Patch(Version v)
        {
            // First, let's add the reference
            var nameReference = new AssemblyNameReference("IPA.Injector", v);
            var injectorPath = Path.Combine(_File.DirectoryName, "IPA.Injector.dll");
            var injector = ModuleDefinition.ReadModule(injectorPath);

            bool hasIPAInjector = false;
            for (int i = 0; i < _Module.AssemblyReferences.Count; i++)
            {
                if (_Module.AssemblyReferences[i].Name == "IllusionInjector")
                    _Module.AssemblyReferences.RemoveAt(i--);
                if (_Module.AssemblyReferences[i].Name == "IllusionPlugin")
                    _Module.AssemblyReferences.RemoveAt(i--);
                if (_Module.AssemblyReferences[i].Name == "IPA.Injector")
                {
                    hasIPAInjector = true;
                    _Module.AssemblyReferences[i].Version = v;
                }
            }

            if (!hasIPAInjector)
            {
                _Module.AssemblyReferences.Add(nameReference);

                int patched = 0;
                foreach (var type in FindEntryTypes())
                {
                    if (PatchType(type, injector))
                    {
                        patched++;
                    }
                }

                if (patched > 0)
                {
                    _Module.Write(_File.FullName);
                }
                else
                {
                    throw new Exception("Could not find any entry type!");
                }
            }
            else
            {
                _Module.Write(_File.FullName);
            }
        }

        private bool PatchType(TypeDefinition targetType, ModuleDefinition injector)
        {
            var targetMethod = targetType.Methods.FirstOrDefault(m => m.IsConstructor && m.IsStatic);
            if (targetMethod != null)
            {
                var methodReference = _Module.Import(injector.GetType("IPA.Injector.Injector").Methods.First(m => m.Name == "Inject"));
                targetMethod.Body.Instructions.Insert(0, Instruction.Create(OpCodes.Call, methodReference));
                return true;
            }
            return false;
        }


        private IEnumerable<TypeDefinition> FindEntryTypes()
        {
            return _Module.GetTypes().Where(m => ENTRY_TYPES.Contains(m.Name));
        }
    }
}